1#!/usr/bin/env python3
2# Convenience shell for using sharkd, including history and tab completion.
3#
4# Copyright (c) 2019 Peter Wu <peter@lekensteyn.nl>
5#
6# SPDX-License-Identifier: GPL-2.0-or-later
7#
8import argparse
9import contextlib
10import glob
11import json
12import logging
13import os
14import readline
15import selectors
16import signal
17import subprocess
18import sys
19
20_logger = logging.getLogger(__name__)
21
22# grep -Po 'tok_req, "\K\w+' sharkd_session.c
23all_commands = """
24load
25status
26analyse
27info
28check
29complete
30frames
31tap
32follow
33iograph
34intervals
35frame
36setcomment
37setconf
38dumpconf
39download
40bye
41""".split()
42all_commands += """
43!pretty
44!histfile
45!debug
46""".split()
47
48
49class SharkdShell:
50    def __init__(self, pretty, history_file):
51        self.pretty = pretty
52        self.history_file = history_file
53
54    def ignore_sigint(self):
55        # Avoid terminating the sharkd child when ^C in the shell.
56        signal.signal(signal.SIGINT, signal.SIG_IGN)
57
58    def sharkd_process(self):
59        sharkd = 'sharkd'
60        env = os.environ.copy()
61        # Avoid loading user preferences which may trigger deprecation warnings.
62        env['WIRESHARK_CONFIG_DIR'] = '/nonexistent'
63        proc = subprocess.Popen([sharkd, '-'],
64                                stdin=subprocess.PIPE,
65                                stdout=subprocess.PIPE,
66                                stderr=subprocess.PIPE,
67                                env=env,
68                                preexec_fn=self.ignore_sigint)
69        banner = proc.stderr.read1().decode('utf8')
70        if banner.strip() != 'Hello in child.':
71            _logger.warning('Unexpected banner: %r', banner)
72        return proc
73
74    def completer(self, text, state):
75        if state == 0:
76            origline = readline.get_line_buffer()
77            line = origline.lstrip()
78            skipped = len(origline) - len(line)
79            startpos = readline.get_begidx() - skipped
80            curpos = readline.get_endidx() - skipped
81            # _logger.debug('Completing: head=%r cur=%r tail=%r',
82            #              line[:startpos], line[startpos:curpos], line[curpos:])
83            completions = []
84            if startpos == 0:
85                completions = all_commands
86            elif line[:1] == '!':
87                cmd = line[1:startpos].strip()
88                if cmd == 'pretty':
89                    completions = ['jq', 'indent', 'off']
90                elif cmd == 'histfile':
91                    # spaces in paths are not supported for now.
92                    completions = glob.glob(glob.escape(text) + '*')
93                elif cmd == 'debug':
94                    completions = ['on', 'off']
95            completions = [x for x in completions if x.startswith(text)]
96            if len(completions) == 1:
97                completions = [completions[0] + ' ']
98            self.completions = completions
99        try:
100            return self.completions[state]
101        except IndexError:
102            return None
103
104    def wrap_exceptions(self, fn):
105        # For debugging, any exception in the completion function is usually
106        # silently ignored by readline.
107        def wrapper(*args):
108            try:
109                return fn(*args)
110            except Exception as e:
111                _logger.exception(e)
112                raise
113        return wrapper
114
115    def add_history(self, line):
116        # Emulate HISTCONTROL=ignorespace to avoid adding to history.
117        if line.startswith(' '):
118            return
119        # Emulate HISTCONTROL=ignoredups to avoid duplicate history entries.
120        nitems = readline.get_current_history_length()
121        lastline = readline.get_history_item(nitems)
122        if lastline != line:
123            readline.add_history(line)
124
125    def parse_command(self, cmd):
126        '''Converts a user-supplied command to a sharkd one.'''
127        # Support 'foo {...}' as alias for '{"req": "foo", ...}'
128        if cmd[0].isalpha():
129            if ' ' in cmd:
130                req, cmd = cmd.split(' ', 1)
131            else:
132                req, cmd = cmd, '{}'
133        elif cmd[0] == '!':
134            return self.parse_special_command(cmd[1:])
135        else:
136            req = None
137        try:
138            c = json.loads(cmd)
139            if req is not None:
140                c['req'] = req
141        except json.JSONDecodeError as e:
142            _logger.error('Invalid command: %s', e)
143            return
144        if type(c) != dict or not 'req' in c:
145            _logger.error('Missing req key in request')
146            return
147        return c
148
149    def parse_special_command(self, cmd):
150        args = cmd.split()
151        if not args:
152            _logger.warning('Missing command')
153            return
154        if args[0] == 'pretty':
155            choices = ['jq', 'indent']
156            if len(args) >= 2:
157                self.pretty = args[1] if args[1] in choices else None
158            print('Pretty printing is now', self.pretty or 'disabled')
159        elif args[0] == 'histfile':
160            if len(args) >= 2:
161                self.history_file = args[1] if args[1] != 'off' else None
162            print('History is now', self.history_file or 'disabled')
163        elif args[0] == 'debug':
164            if len(args) >= 2 and args[1] in ('on', 'off'):
165                _logger.setLevel(
166                    logging.DEBUG if args[1] == 'on' else logging.INFO)
167            print('Debug logging is now',
168                  ['off', 'on'][_logger.level == logging.DEBUG])
169        else:
170            _logger.warning('Unsupported command %r', args[0])
171
172    @contextlib.contextmanager
173    def wrap_history(self):
174        '''Loads history at startup and saves history on exit.'''
175        readline.set_auto_history(False)
176        try:
177            if self.history_file:
178                readline.read_history_file(self.history_file)
179            h_len = readline.get_current_history_length()
180        except FileNotFoundError:
181            h_len = 0
182        try:
183            yield
184        finally:
185            new_items = readline.get_current_history_length() - h_len
186            if new_items > 0 and self.history_file:
187                open(self.history_file, 'a').close()
188                readline.append_history_file(new_items, self.history_file)
189
190    def shell_prompt(self):
191        '''Sets up the interactive prompt.'''
192        readline.parse_and_bind("tab: complete")
193        readline.set_completer(self.wrap_exceptions(self.completer))
194        readline.set_completer_delims(' ')
195        return self.wrap_history()
196
197    def read_command(self):
198        while True:
199            try:
200                origline = input('# ')
201            except EOFError:
202                raise
203            except KeyboardInterrupt:
204                print('^C', file=sys.stderr)
205                continue
206            cmd = origline.strip()
207            if not cmd:
208                return
209            self.add_history(origline)
210            c = self.parse_command(cmd)
211            if c:
212                return json.dumps(c)
213
214    def want_input(self):
215        '''Request the prompt to be displayed.'''
216        os.write(self.user_input_wr, b'x')
217
218    def main_loop(self):
219        sel = selectors.DefaultSelector()
220        user_input_rd, self.user_input_wr = os.pipe()
221        self.want_input()
222        with self.sharkd_process() as proc, self.shell_prompt():
223            self.process = proc
224            sel.register(proc.stdout, selectors.EVENT_READ, self.handle_stdout)
225            sel.register(proc.stderr, selectors.EVENT_READ, self.handle_stderr)
226            sel.register(user_input_rd, selectors.EVENT_READ, self.handle_user)
227            interrupts = 0
228            while True:
229                try:
230                    events = sel.select()
231                    _logger.debug('got events: %r', events)
232                    if not events:
233                        break
234                    for key, mask in events:
235                        key.data(key)
236                    interrupts = 0
237                except KeyboardInterrupt:
238                    print('Interrupt again to abort immediately.', file=sys.stderr)
239                    interrupts += 1
240                    if interrupts >= 2:
241                        break
242                if self.want_command:
243                    self.ask_for_command_and_run_it()
244                # Process died? Stop the shell.
245                if proc.poll() is not None:
246                    break
247
248    def handle_user(self, key):
249        '''Received a notification that another prompt can be displayed.'''
250        os.read(key.fileobj, 4096)
251        self.want_command = True
252
253    def ask_for_command_and_run_it(self):
254        cmd = self.read_command()
255        if not cmd:
256            # Give a chance for the event loop to run again.
257            self.want_input()
258            return
259        self.want_command = False
260        _logger.debug('Running: %r', cmd)
261        self.process.stdin.write((cmd + '\n').encode('utf8'))
262        self.process.stdin.flush()
263
264    def handle_stdout(self, key):
265        resp = key.fileobj.readline().decode('utf8')
266        _logger.debug('Response: %r', resp)
267        if not resp:
268            raise EOFError
269        self.want_input()
270        resp = resp.strip()
271        if resp:
272            try:
273                if self.pretty == 'jq':
274                    subprocess.run(['jq', '.'], input=resp,
275                                   universal_newlines=True)
276                elif self.pretty == 'indent':
277                    r = json.loads(resp)
278                    json.dump(r, sys.stdout, indent='  ')
279                    print('')
280                else:
281                    print(resp)
282            except Exception as e:
283                _logger.warning('Dumping output as-is due to: %s', e)
284                print(resp)
285
286    def handle_stderr(self, key):
287        data = key.fileobj.read1().decode('utf8')
288        print(data, end="", file=sys.stderr)
289
290
291parser = argparse.ArgumentParser()
292parser.add_argument('--debug', action='store_true',
293                    help='Enable verbose logging')
294parser.add_argument('--pretty', choices=['jq', 'indent'],
295                    help='Pretty print responses (one of: %(choices)s)')
296parser.add_argument('--histfile',
297                    help='Log shell history to this file')
298
299
300def main(args):
301    logging.basicConfig()
302    _logger.setLevel(logging.DEBUG if args.debug else logging.INFO)
303    shell = SharkdShell(args.pretty, args.histfile)
304    try:
305        shell.main_loop()
306    except EOFError:
307        print('')
308
309
310if __name__ == '__main__':
311    main(parser.parse_args())
312