#!/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]()