tools/sharkd_shell.py

   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 #
   8 import argparse
   9 import contextlib
  10 import glob
  11 import json
  12 import logging
  13 import os
  14 import readline
  15 import selectors
  16 import signal
  17 import subprocess
  18 import sys
  19
  20 _logger = logging.getLogger(__name__)
  21
  22 # grep -Po 'tok_req, "\K\w+' sharkd_session.c
  23 all_commands = """
  24 load
  25 status
  26 analyse
  27 info
  28 check
  29 complete
  30 frames
  31 tap
  32 follow
  33 iograph
  34 intervals
  35 frame
  36 setcomment
  37 setconf
  38 dumpconf
  39 download
  40 bye
  41 """.split()
  42 all_commands += """
  43 !pretty
  44 !histfile
  45 !debug
  46 """.split()
  47
  48
  49 class 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
 291 parser = argparse.ArgumentParser()
 292 parser.add_argument('--debug', action='store_true',
 293                     help='Enable verbose logging')
 294 parser.add_argument('--pretty', choices=['jq', 'indent'],
 295                     help='Pretty print responses (one of: %(choices)s)')
 296 parser.add_argument('--histfile',
 297                     help='Log shell history to this file')
 298
 299
 300 def 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
 310 if __name__ == '__main__':
 311     main(parser.parse_args())