add master daemon

1 files changed, 281 insertions, 0 deletions

diff --git a/scripts/master_daemon.py b/scripts/master_daemon.py
new file mode 100755
index 00000000..caaaf5fc
--- /dev/null
+++ b/scripts/master_daemon.py
@@ -0,0 +1,281 @@
+#!/usr/bin/python
+import os, sys, getpass
+import tempfile
+import argparse
+import socket
+from time import time, sleep
+import subprocess, shlex, shutil, multiprocessing
+from distutils.dir_util import mkpath
+from glob import glob
+import logging
+logging.basicConfig(level = logging.FATAL) # keep quiet
+
+sys.path.append(os.path.join('external_libs', 'jsonrpclib-pelix-0.2.5'))
+from jsonrpclib.SimpleJSONRPCServer import SimpleJSONRPCServer
+
+sys.path.append(os.path.join('external_libs', 'termstyle'))
+import termstyle
+
+
+### Server functions ###
+
+def check_connectivity():
+    return True
+
+def add(a, b): # for sanity checks
+    return a + b
+
+def is_trex_daemon_running():
+    ret_code, stdout, stderr = run_command('ps -u root --format comm')
+    if ret_code:
+        raise Exception('Failed to list running processes, stderr: %s' % stderr)
+    if 'trex_daemon_ser' in stdout: # name is cut
+        return True
+    return False
+
+def restart_trex_daemon():
+    if is_trex_daemon_running:
+        stop_trex_daemon()
+    start_trex_daemon()
+    return True
+
+def stop_trex_daemon():
+    if not is_trex_daemon_running():
+        return False
+    return_code, stdout, stderr = run_command('%s stop' % trex_daemon_path)
+    if return_code:
+        raise Exception('Could not stop trex_daemon_server, err: %s' % stderr)
+    for i in range(20):
+        if not is_trex_daemon_running():
+            return True
+        sleep(0.1)
+    raise Exception('Could not stop trex_daemon_server')
+
+def start_trex_daemon():
+    if is_trex_daemon_running():
+        return False
+    return_code, stdout, stderr = run_command('%s start' % trex_daemon_path)
+    if return_code:
+        raise Exception('Could not run trex_daemon_server, err: %s' % stderr)
+    for i in range(20):
+        if is_trex_daemon_running():
+            return True
+        sleep(0.1)
+    raise Exception('Could not run trex_daemon_server')
+
+def update_trex(package_path = 'http://trex-tgn.cisco.com/trex/release/latest'):
+    if not args.allow_update:
+        raise Exception('Updading server not allowed')
+    try:
+        cur_dir = args.trex_dir
+        bu_dir = '%s_BU%i' % (cur_dir, int(time()))
+        # rename current dir for backup
+        try:
+            if os.path.exists(bu_dir):
+                shutil.rmtree(bu_dir)
+            shutil.move(cur_dir, bu_dir)
+            mkpath(cur_dir)
+        except Exception as e:
+            raise Exception('Could not make backup of previous directory. Err: %s' % e)
+        # getting new package
+        if package_path.startswith('http://'):
+            file_name = package_path.split('/')[-1]
+            ret_code, stdout, stderr = run_command('wget %s -O %s' % (package_path, cur_dir), timeout = 600)
+        elif os.path.normpath(package_path).startswith('/auto/') and package_path.endswith('.tar.gz'):
+            file_name = os.path.basename(package_path)
+            ret_code, stdout, stderr = run_command('rsync -cL %s %s' % (package_path, cur_dir), timeout = 300)
+        else:
+            raise Exception('package_path should be http address or path starting with /auto')
+        if ret_code:
+            raise Exception('Could not get requested package: %s' % stderr)
+        # unpacking
+        ret_code, stdout, stderr = run_command('tar -mxzf %s -C %s' % (os.path.join(cur_dir, file_name), cur_dir))
+        if ret_code:
+            raise Exception('Could not untar the package: %s' % stderr)
+        # version is the name of dir
+        dir_cont = glob(os.path.join(cur_dir, 'v[0-9].[0-9][0-9]'))
+        if not len(dir_cont):
+            raise Exception('Did not found directory with version name after unpacking.')
+        if len(dir_cont) > 1:
+            raise Exception('Found more than one directory with version name after unpacking.')
+        version = os.path.basename(dir_cont[0])
+        # moving files from dir with version one level up
+        for root, dirs, files in os.walk(os.path.join(cur_dir, version)):
+            for f in files:
+                shutil.move(os.path.join(root, f), os.path.join(cur_dir, f))
+            for d in dirs:
+                shutil.move(os.path.join(root, d), os.path.join(cur_dir, d))
+        # no errors, remove BU dir
+        if os.path.exists(bu_dir):
+            shutil.rmtree(bu_dir)
+        return True
+    except: # something went wrong, return backup dir
+        if os.path.exists(cur_dir):
+            shutil.rmtree(cur_dir)
+        shutil.move(bu_dir, cur_dir)
+        raise
+
+
+def unpack_client():
+    if not args.allow_update:
+        raise Exception('Unpacking of client not allowed')
+    # determining client archive
+    client_pkg_files = glob(os.path.join(args.trex_dir, 'trex_client*.tar.gz'))
+    if not len(client_pkg_files):
+        raise Exception('Could not find client package')
+    if len(client_pkg_files) > 1:
+        raise Exception('Found more than one client packages')
+    client_pkg_name = os.path.basename(client_pkg_files[0])
+    client_new_path = os.path.join(args.trex_dir, 'trex_client')
+    if os.path.exists(client_new_path):
+        if os.path.islink(client_new_path) or os.path.isfile(client_new_path):
+            os.unlink(client_new_path)
+        else:
+            shutil.rmtree(client_new_path)
+    # unpacking
+    ret_code, stdout, stderr = run_command('tar -mxzf %s -C %s' % (client_pkg_files[0], args.trex_dir))
+    if ret_code:
+        raise Exception('Could not untar the client package: %s' % stderr)
+    return True
+
+### /Server functions ###
+
+def fail(msg):
+    print(msg)
+    sys.exit(-1)
+
+def run_command(command, timeout = 15):
+    command = 'timeout %s %s' % (timeout, command)
+    # pipes might stuck, even with timeout
+    with tempfile.TemporaryFile() as stdout_file, tempfile.TemporaryFile() as stderr_file:
+        proc = subprocess.Popen(shlex.split(command), stdout=stdout_file, stderr=stderr_file)
+        proc.wait()
+        stdout_file.seek(0)
+        stderr_file.seek(0)
+        return (proc.returncode, stdout_file.read().decode(errors = 'replace'), stderr_file.read().decode(errors = 'replace'))
+
+def show_master_daemon_status():
+    if get_master_daemon_pid():
+        print(termstyle.red('Master daemon is running'))
+    else:
+        print(termstyle.red('Master daemon is NOT running'))
+
+def start_master_daemon():
+    if get_master_daemon_pid():
+        print(termstyle.red('Master daemon is already running'))
+        return
+    server = multiprocessing.Process(target = start_master_daemon_func)
+    server.daemon = True
+    server.start()
+    for i in range(10):
+        if get_master_daemon_pid():
+            print(termstyle.green('Master daemon is started'))
+            os._exit(0)
+        sleep(0.1)
+    fail(termstyle.red('Master daemon failed to run'))
+
+def restart_master_daemon():
+    if get_master_daemon_pid():
+        kill_master_daemon()
+    start_master_daemon()
+
+def start_master_daemon_func():
+    server = SimpleJSONRPCServer(('0.0.0.0', args.daemon_port))
+    print('Started master daemon (port %s)' % args.daemon_port)
+    server.register_function(add)
+    server.register_function(check_connectivity)
+    server.register_function(is_trex_daemon_running)
+    server.register_function(restart_trex_daemon)
+    server.register_function(start_trex_daemon)
+    server.register_function(stop_trex_daemon)
+    server.register_function(unpack_client)
+    server.register_function(update_trex)
+    server.serve_forever()
+
+
+def get_master_daemon_pid():
+    return_code, stdout, stderr = run_command('netstat -tlnp')
+    if return_code:
+        fail('Failed to determine which program holds port %s, netstat error: %s' % (args.daemon_port, stderr))
+    for line in stdout.splitlines():
+        if '0.0.0.0:%s' % args.daemon_port in line:
+            line_arr = line.split()
+            if len(line_arr) != 7:
+                fail('Could not parse netstat line to determine which process holds port %s: %s'(args.daemon_port, line))
+            if '/' not in line_arr[6]:
+                fail('Expecting pid/program name in netstat line of using port %s, got: %s'(args.daemon_port, line_arr[6]))
+            pid, program = line_arr[6].split('/')
+            if 'python' not in program and 'master_daemon' not in program:
+                fail('Some other program holds port %s, not our daemon: %s. Please verify.' % (args.daemon_port, program))
+            return pid
+    return None
+
+def kill_master_daemon():
+    pid = get_master_daemon_pid()
+    if not pid:
+        print(termstyle.red('Master daemon is NOT running'))
+        return True
+    return_code, stdout, stderr = run_command('kill %s' % pid) # usual kill
+    if return_code:
+        fail('Failed to kill master daemon, error: %s' % stderr)
+    for i in range(10):
+        if not get_master_daemon_pid():
+            print(termstyle.green('Master daemon is killed'))
+            return True
+        sleep(0.1)
+    return_code, stdout, stderr = run_command('kill -9 %s' % pid) # unconditional kill
+    if return_code:
+        fail('Failed to kill trex_daemon, error: %s' % stderr)
+    for i in range(10):
+        if not get_master_daemon_pid():
+            print(termstyle.green('Master daemon is killed'))
+            return True
+        sleep(0.1)
+    fail('Failed to kill master daemon, even with -9. Please review manually.') # should not happen
+
+# returns True if given path is under current dir or /tmp
+def _check_path_under_current_or_temp(path):
+    if not os.path.relpath(path, '/tmp').startswith(os.pardir):
+        return True
+    if not os.path.relpath(path, os.getcwd()).startswith(os.pardir):
+        return True
+    return False
+
+### Main ###
+
+if getpass.getuser() != 'root':
+    fail('Please run this program as root/with sudo')
+
+actions_help = '''Specify action command to be applied on master daemon.
+    (*) start      : start the master daemon.
+    (*) show       : prompt the status of master daemon process (running / not running).
+    (*) stop       : exit the master daemon process.
+    (*) restart    : stop, then start again the master daemon process
+    '''
+action_funcs = {'start': start_master_daemon,
+                'show': show_master_daemon_status,
+                'stop': kill_master_daemon,
+                'restart': restart_master_daemon,
+                }
+
+parser = argparse.ArgumentParser(description = 'Runs master daemon that can start/stop TRex daemon or update TRex version.')
+parser.add_argument('-p', '--daemon-port', type=int, default = 8091, dest='daemon_port', 
+                    help = 'Select port on which the master_daemon runs.\nDefault is 8091.', action = 'store')
+parser.add_argument('-d', '--trex-dir', type=str, default = os.getcwd(), dest='trex_dir',
+                    help = 'Path of TRex, default is current dir', action = 'store')
+parser.add_argument('--allow-update', default = False, dest='allow_update', action = 'store_true',
+                    help = "Allow update of TRex via RPC command. WARNING: It's security hole! Use on your risk!")
+parser.add_argument('action', choices=action_funcs.keys(),
+                    action='store', help=actions_help)
+parser.usage = None
+args = parser.parse_args()
+
+if not os.path.exists(args.trex_dir):
+    raise Exception("Given path '%s' does not exist" % args.trex_dir)
+if not _check_path_under_current_or_temp(args.trex_dir):
+    raise Exception('Only allowed to use path under /tmp or current directory')
+
+trex_daemon_path = os.path.join(args.trex_dir, 'trex_daemon_server')
+action_funcs[args.action]()
+
+