#!/usr/bin/python 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) 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 = 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(): 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 check_path = cur_dir last_err_path = None ret = -1 while ret and check_path != '/': ret, out = run_command("ls %s" % check_path, demote = True) 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) server_path = os.path.join(cur_dir, 'automation', 'trex_control_plane', 'stl', 'services', 'scapy_server') cmd = 'taskset -c {core} {python} ./scapy_zmq_server.py -s {port}'.format(core = args.core, python = sys.executable, port = args.port) proc, stdout_file = run_command(cmd, demote = True, is_daemon = True, cwd = server_path) ret = progress(is_running, 'Starting Scapy server', 'Scapy server is started', 'Scapy server failed to run', 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('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', timeout = 15) 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', timeout = 15) 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]()