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