From 89b608ae766705950efc5f4914b01b9a32b6a0e7 Mon Sep 17 00:00:00 2001 From: Yaroslav Brustinov Date: Fri, 13 May 2016 20:11:51 +0300 Subject: add master daemon --- scripts/master_daemon.py | 281 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 281 insertions(+) create mode 100755 scripts/master_daemon.py (limited to 'scripts/master_daemon.py') 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]() + + -- cgit 1.2.3-korg