#! /bin/bash
"source" "find_python.sh" "--local"
"exec" "$PYTHON" "$0" "$@"

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

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

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

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

from netstat import netstat
from termstyle import termstyle


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


def progress(success_check, start_msg, success_msg, fail_msg, timeout = 35, poll_rate = 0.5, fail_check = None):
    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
        if fail_check and fail_check():
            print(termstyle.red(' ' + fail_msg))
            return 1
        time.sleep(poll_rate)
        sys.stdout.write('.')
        sys.stdout.flush()
    print(termstyle.red(' Timeout'))
    return 1


def demote_func():
    pw_record = pwd.getpwnam('nobody')
    os.setgid(pw_record.pw_gid)
    os.setuid(pw_record.pw_uid)


def run_command(command, timeout = 30, poll_rate = 0.1, cwd = None, demote = False, is_daemon = False):
    if not is_daemon:
        assert timeout > 0, 'Timeout should be positive'
        assert poll_rate > 0, 'Poll rate should be positive'

    preexec_fn = demote_func if demote else None

    try: # P2
        stdout_file = tempfile.TemporaryFile(bufsize = 0)
    except: # P3
        stdout_file = tempfile.TemporaryFile(buffering = 0)

    try:
        proc = subprocess.Popen(shlex.split(command), stdout = stdout_file, stderr = subprocess.STDOUT, cwd = cwd,
                                close_fds = True, universal_newlines = True, preexec_fn = preexec_fn)
        if is_daemon:
            return proc, stdout_file
        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())
    finally:
        if not is_daemon:
            stdout_file.close()


def get_daemon_pid():
    pid = None
    for conn in netstat.netstat(with_pid = True, search_local_port = args.port):
        if conn[2] == '0.0.0.0' 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, search_local_port = args.port):
        if conn[2] == '0.0.0.0' and conn[6] == 'LISTEN':
            return True
    return False


def show_daemon_status():
    if is_running():
        print(termstyle.green('{name} server is running'.format(name=args.name)))
    else:
        print(termstyle.red('{name} server is NOT running'.format(name=args.name)))


def start_daemon():
    if is_running():
        print(termstyle.red('{name} server is already running'.format(args.name)))
        return
    if args.check_readable:
        check_path = cur_dir
        last_err_path = None
        ret = -1
        while ret and check_path != '/':
            ret, out = run_command("ls %s" % check_path, demote = args.demote)
            if ret:
                last_err_path = check_path
                check_path = os.path.abspath(os.path.join(check_path, '..'))
        if last_err_path:
            msg = '''
    Error: current path is not readable by user "nobody" (starting at {path}).
    Two possible solutions:

    1. (Recommended)
        Copy TRex to some public location (/tmp or /scratch, assuming it has proper permissions (chmod 777 etc.))

    2. (Not recommended)
        Change permissions of current path. (Starting from directory {path}).
        chmod 777 {path} -R
    '''.format(path = last_err_path)
            fail(msg)

    if args.interactive:
        server_path = os.path.join(cur_dir, 'automation', 'trex_control_plane', 'interactive')
    else:
        server_path = cur_dir
    assert args.exe is not None, 'Must provide -e with an executable string for start command'
    sudo   = 'sudo' if args.sudo else ''
    python = sys.executable if args.py else ''
    cmd = '{sudo} taskset -c {core} {python} {exe} -p {port}'.format(sudo = sudo, core = args.core, python = python, exe = args.exe,port = args.port)

    proc, stdout_file = run_command(cmd, demote = args.demote, is_daemon = True, cwd = server_path)
    ret = progress(is_running, 'Starting {name} server'.format(name=args.name), '{name} server is started'.format(name=args.name), '{name} server failed to run'.format(name=args.name), fail_check = proc.poll)
    if proc.poll():
        stdout_file.seek(0)
        print('Output: %s' % stdout_file.read())
        stdout_file.close()
        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('{name} server is NOT running'.format(name=args.name)))
        return True
    run_command('kill %s' % pid) # usual kill
    ret = progress(inv(is_running), 'Killing {name} server'.format(name=args.name), '{name} server is killed'.format(name=args.name), 'failed', timeout = 15)
    if not ret:
        return
    _, out = run_command('kill -9 %s' % pid) # unconditional kill
    ret = progress(inv(is_running), 'Killing {name} server with -9'.format(name=args.name), '{name} server is killed'.format(name=args.name), 'failed', timeout = 15)
    if ret:
        fail('Failed to kill {name} server, even with -9. Please review manually.\nOutput: {out}'.format(name=args.name, out=out))


def get_default_port(server_name):
    name_2_port = {'Scapy': 4507, 'PyBird': 4509, 'Emu': 4510}
    port = name_2_port.get(server_name)
    if port is None:
        raise Exception('Unknown server name: "{0}". Choose one of: {1}'.format(server_name, list(name_2_port.keys())))
    return port


### 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 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 General server application.',
        formatter_class = RawTextHelpFormatter,
    )
    
    parser.add_argument('-n', '--name', type = str, required = True,
        help='Server name i.e: Scapy, Emu or PyBird')
    parser.add_argument('-p', '--port', type = int,
        help='Select tcp port on which server will listen.\nDefault is using script default.')
    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('--sudo', action = 'store_true', default = False,
        help='Run taskset with sudo')
    parser.add_argument('action', choices=action_funcs.keys(), action='store', help=actions_help)
    parser.add_argument('--py', action = 'store_true', default = False,
        help='True will run execute command with python, false as a bash')
    parser.add_argument('-e', '--exe', type = str,
        help='Command to execute, starting the server')
    parser.add_argument('-d', '--demote', action = 'store_true', default = False,
        help='Demote with commands, used in Scapy server')
    parser.add_argument('-r', '--readable', action = 'store_true', default = False, dest = 'check_readable',
        help='Check readable path, used in Scapy server')
    parser.add_argument('-i', '--interactive', action = 'store_true', default = False,
        help='True will run command from interactive path, false from scripts')
    parser.usage = None
    args = parser.parse_args()
    
    if args.port is None:
        args.port = get_default_port(args.name)

    action_funcs[args.action]()
