summaryrefslogtreecommitdiffstats
path: root/scripts/scapy_daemon_server
blob: e19c8b3c9a3b8e935f0c512441ef0791e4883446 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
#!/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]()