#!/usr/bin/python

import os, sys
import tempfile
import time
import subprocess, shlex
from argparse import ArgumentParser, RawTextHelpFormatter
import errno

def fail(msg):
    print(msg)
    sys.exit(-1)

if os.getuid() != 0:
    fail('Please run this program as root/with sudo')

pwd = os.path.abspath(os.path.dirname(__file__))
ext_libs_path = os.path.join(pwd, 'external_libs')
if ext_libs_path not in sys.path:
    sys.path.append(ext_libs_path)

import netstat
try:
    from termstyle import termstyle
except ImportError:
    import termstyle


def inv(f):
    return lambda *a, **k: not f(*a, **k)


def progress(success_check, start_msg, success_msg, fail_msg, timeout = 10, poll_rate = 0.5):
    sys.stdout.write('%s...' % start_msg)
    sys.stdout.flush()
    for i in range(int(timeout/poll_rate)):
        if success_check():
            print(termstyle.green(' ' + success_msg))
            return 0
        time.sleep(poll_rate)
        sys.stdout.write('.')
        sys.stdout.flush()
    if success_check():
        print(termstyle.green(' ' + success_msg))
        return 0
    print(termstyle.red(' ' + fail_msg))
    return 1


def run_command(command, timeout = 15, poll_rate = 0.1, cwd = None):
    assert timeout > 0, 'Timeout should be positive'
    assert poll_rate > 0, 'Poll rate should be positive'
    try:
        tempfile.TemporaryFile(bufsize=0)
        temp_params = {'bufsize': 0}
    except:
        tempfile.TemporaryFile(buffering=0)
        temp_params = {'buffering': 0}
    with tempfile.TemporaryFile(**temp_params) as stdout_file:
        proc = subprocess.Popen(shlex.split(command), stdout = stdout_file, stderr = subprocess.STDOUT, cwd = cwd, close_fds = True, universal_newlines = True)
        for i in range(int(timeout/poll_rate)):
            time.sleep(poll_rate)
            if proc.poll() is not None: # process stopped
                break
        if proc.poll() is None:
            proc.kill() # timeout
            stdout_file.seek(0)
            return (errno.ETIMEDOUT, '%s\n\n...Timeout of %s second(s) is reached!' % (stdout_file.read(), timeout))
        stdout_file.seek(0)
        return (proc.returncode, stdout_file.read())


def get_daemon_pid():
    pid = None
    for conn in netstat.netstat():
        if conn[2] == '0.0.0.0' and int(conn[3]) == args.port and conn[6] == 'LISTEN':
            pid = conn[7]
            if pid is None:
                raise Exception('Found the connection, but could not determine pid: %s' % conn)
            break
    return pid


# faster variant of get_daemon_pid
def is_running():
    for conn in netstat.netstat(with_pid = False):
        if conn[2] == '0.0.0.0' and int(conn[3]) == args.port and conn[6] == 'LISTEN':
            return True
    return False


def show_daemon_status():
    if is_running():
        print(termstyle.green('Scapy server is running'))
    else:
        print(termstyle.red('Scapy server is NOT running'))


def start_daemon():
    if is_running():
        print(termstyle.red('Scapy server is already running'))
        return
    server_path = os.path.join(pwd, 'automation', 'trex_control_plane', 'stl', 'services', 'scapy_server')
    with tempfile.TemporaryFile() as stdout_file:
        subprocess.Popen(shlex.split("taskset -c %s su -s /bin/bash -c '%s scapy_zmq_server.py -s %s' nobody" % (args.core, sys.executable, args.port)),
                    stdout = stdout_file, stderr = subprocess.STDOUT, cwd = server_path, close_fds = True, universal_newlines = True)
        ret = progress(is_running, 'Starting Scapy server', 'Scapy server is started', 'Scapy server failed to run')
        if ret:
            stdout_file.seek(0)
            print('Output: %s' % stdout_file.read())
            sys.exit(1)


def restart_daemon():
    if is_running():
        kill_daemon()
        time.sleep(0.5)
    start_daemon()


def kill_daemon():
    pid = get_daemon_pid()
    if not pid:
        print(termstyle.red('Scapy server is NOT running'))
        return True
    run_command('kill %s' % pid) # usual kill
    ret = progress(inv(is_running), 'Killing Scapy server', 'Scapy server is killed', 'failed')
    if not ret:
        return
    _, out = run_command('kill -9 %s' % pid) # unconditional kill
    ret = progress(inv(is_running), 'Killing Scapy server with -9', 'Scapy server is killed', 'failed')
    if ret:
        fail('Failed to kill Scapy server, even with -9. Please review manually.\nOutput: %s' % out)

### Main ###

if __name__ == '__main__':
    actions_help = '''Specify action command to be applied on server.
        (*) start      : start the application in as a daemon process.
        (*) show       : prompt an updated status of daemon process (running/ not running).
        (*) stop       : exit Scapy server daemon process.
        (*) restart    : stop, then start again the application as daemon process
        '''
    action_funcs = {'start': start_daemon,
                    'show': show_daemon_status,
                    'stop': kill_daemon,
                    'restart': restart_daemon,
                    }
    
    parser = ArgumentParser(description = 'Runs Scapy server application.',
        formatter_class = RawTextHelpFormatter,
    )
    
    parser.add_argument('-p', '--port', type = int, default = 4507,
            help='Select tcp port on which Scapy server will listen.\nDefault is 4507.')
    parser.add_argument('-c', '--core', type = int, default = 0,
            help='Core number to set affinity.\nAdvised to set on free core or TRex master thread core\nDefault is 0.')
    parser.add_argument('action', choices=action_funcs.keys(),
                            action='store', help=actions_help)
    parser.usage = None
    args = parser.parse_args()
    
    action_funcs[args.action]()