b5d6c9922b
Interactive shell for sharkd with some auto-completion, history, and JSON beautification support. Quick and dirty, tested with Python 3.7.4 on Arch Linux. Change-Id: I5dec1ac07230bfcbc3165a36b5c7f088ee65a8f7 Reviewed-on: https://code.wireshark.org/review/34583 Reviewed-by: Peter Wu <peter@lekensteyn.nl>
312 lines
10 KiB
Python
Executable file
312 lines
10 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
# Convenience shell for using sharkd, including history and tab completion.
|
|
#
|
|
# Copyright (c) 2019 Peter Wu <peter@lekensteyn.nl>
|
|
#
|
|
# SPDX-License-Identifier: GPL-2.0-or-later
|
|
#
|
|
import argparse
|
|
import contextlib
|
|
import glob
|
|
import json
|
|
import logging
|
|
import os
|
|
import readline
|
|
import selectors
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
# grep -Po 'tok_req, "\K\w+' sharkd_session.c
|
|
all_commands = """
|
|
load
|
|
status
|
|
analyse
|
|
info
|
|
check
|
|
complete
|
|
frames
|
|
tap
|
|
follow
|
|
iograph
|
|
intervals
|
|
frame
|
|
setcomment
|
|
setconf
|
|
dumpconf
|
|
download
|
|
bye
|
|
""".split()
|
|
all_commands += """
|
|
!pretty
|
|
!histfile
|
|
!debug
|
|
""".split()
|
|
|
|
|
|
class SharkdShell:
|
|
def __init__(self, pretty, history_file):
|
|
self.pretty = pretty
|
|
self.history_file = history_file
|
|
|
|
def ignore_sigint(self):
|
|
# Avoid terminating the sharkd child when ^C in the shell.
|
|
signal.signal(signal.SIGINT, signal.SIG_IGN)
|
|
|
|
def sharkd_process(self):
|
|
sharkd = 'sharkd'
|
|
env = os.environ.copy()
|
|
# Avoid loading user preferences which may trigger deprecation warnings.
|
|
env['WIRESHARK_CONFIG_DIR'] = '/nonexistent'
|
|
proc = subprocess.Popen([sharkd, '-'],
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
env=env,
|
|
preexec_fn=self.ignore_sigint)
|
|
banner = proc.stderr.read1().decode('utf8')
|
|
if banner.strip() != 'Hello in child.':
|
|
_logger.warning('Unexpected banner: %r', banner)
|
|
return proc
|
|
|
|
def completer(self, text, state):
|
|
if state == 0:
|
|
origline = readline.get_line_buffer()
|
|
line = origline.lstrip()
|
|
skipped = len(origline) - len(line)
|
|
startpos = readline.get_begidx() - skipped
|
|
curpos = readline.get_endidx() - skipped
|
|
# _logger.debug('Completing: head=%r cur=%r tail=%r',
|
|
# line[:startpos], line[startpos:curpos], line[curpos:])
|
|
completions = []
|
|
if startpos == 0:
|
|
completions = all_commands
|
|
elif line[:1] == '!':
|
|
cmd = line[1:startpos].strip()
|
|
if cmd == 'pretty':
|
|
completions = ['jq', 'indent', 'off']
|
|
elif cmd == 'histfile':
|
|
# spaces in paths are not supported for now.
|
|
completions = glob.glob(glob.escape(text) + '*')
|
|
elif cmd == 'debug':
|
|
completions = ['on', 'off']
|
|
completions = [x for x in completions if x.startswith(text)]
|
|
if len(completions) == 1:
|
|
completions = [completions[0] + ' ']
|
|
self.completions = completions
|
|
try:
|
|
return self.completions[state]
|
|
except IndexError:
|
|
return None
|
|
|
|
def wrap_exceptions(self, fn):
|
|
# For debugging, any exception in the completion function is usually
|
|
# silently ignored by readline.
|
|
def wrapper(*args):
|
|
try:
|
|
return fn(*args)
|
|
except Exception as e:
|
|
_logger.exception(e)
|
|
raise
|
|
return wrapper
|
|
|
|
def add_history(self, line):
|
|
# Emulate HISTCONTROL=ignorespace to avoid adding to history.
|
|
if line.startswith(' '):
|
|
return
|
|
# Emulate HISTCONTROL=ignoredups to avoid duplicate history entries.
|
|
nitems = readline.get_current_history_length()
|
|
lastline = readline.get_history_item(nitems)
|
|
if lastline != line:
|
|
readline.add_history(line)
|
|
|
|
def parse_command(self, cmd):
|
|
'''Converts a user-supplied command to a sharkd one.'''
|
|
# Support 'foo {...}' as alias for '{"req": "foo", ...}'
|
|
if cmd[0].isalpha():
|
|
if ' ' in cmd:
|
|
req, cmd = cmd.split(' ', 1)
|
|
else:
|
|
req, cmd = cmd, '{}'
|
|
elif cmd[0] == '!':
|
|
return self.parse_special_command(cmd[1:])
|
|
else:
|
|
req = None
|
|
try:
|
|
c = json.loads(cmd)
|
|
if req is not None:
|
|
c['req'] = req
|
|
except json.JSONDecodeError as e:
|
|
_logger.error('Invalid command: %s', e)
|
|
return
|
|
if type(c) != dict or not 'req' in c:
|
|
_logger.error('Missing req key in request')
|
|
return
|
|
return c
|
|
|
|
def parse_special_command(self, cmd):
|
|
args = cmd.split()
|
|
if not args:
|
|
_logger.warning('Missing command')
|
|
return
|
|
if args[0] == 'pretty':
|
|
choices = ['jq', 'indent']
|
|
if len(args) >= 2:
|
|
self.pretty = args[1] if args[1] in choices else None
|
|
print('Pretty printing is now', self.pretty or 'disabled')
|
|
elif args[0] == 'histfile':
|
|
if len(args) >= 2:
|
|
self.history_file = args[1] if args[1] != 'off' else None
|
|
print('History is now', self.history_file or 'disabled')
|
|
elif args[0] == 'debug':
|
|
if len(args) >= 2 and args[1] in ('on', 'off'):
|
|
_logger.setLevel(
|
|
logging.DEBUG if args[1] == 'on' else logging.INFO)
|
|
print('Debug logging is now',
|
|
['off', 'on'][_logger.level == logging.DEBUG])
|
|
else:
|
|
_logger.warning('Unsupported command %r', args[0])
|
|
|
|
@contextlib.contextmanager
|
|
def wrap_history(self):
|
|
'''Loads history at startup and saves history on exit.'''
|
|
readline.set_auto_history(False)
|
|
try:
|
|
if self.history_file:
|
|
readline.read_history_file(self.history_file)
|
|
h_len = readline.get_current_history_length()
|
|
except FileNotFoundError:
|
|
h_len = 0
|
|
try:
|
|
yield
|
|
finally:
|
|
new_items = readline.get_current_history_length() - h_len
|
|
if new_items > 0 and self.history_file:
|
|
open(self.history_file, 'a').close()
|
|
readline.append_history_file(new_items, self.history_file)
|
|
|
|
def shell_prompt(self):
|
|
'''Sets up the interactive prompt.'''
|
|
readline.parse_and_bind("tab: complete")
|
|
readline.set_completer(self.wrap_exceptions(self.completer))
|
|
readline.set_completer_delims(' ')
|
|
return self.wrap_history()
|
|
|
|
def read_command(self):
|
|
while True:
|
|
try:
|
|
origline = input('# ')
|
|
except EOFError:
|
|
raise
|
|
except KeyboardInterrupt:
|
|
print('^C', file=sys.stderr)
|
|
continue
|
|
cmd = origline.strip()
|
|
if not cmd:
|
|
return
|
|
self.add_history(origline)
|
|
c = self.parse_command(cmd)
|
|
if c:
|
|
return json.dumps(c)
|
|
|
|
def want_input(self):
|
|
'''Request the prompt to be displayed.'''
|
|
os.write(self.user_input_wr, b'x')
|
|
|
|
def main_loop(self):
|
|
sel = selectors.DefaultSelector()
|
|
user_input_rd, self.user_input_wr = os.pipe()
|
|
self.want_input()
|
|
with self.sharkd_process() as proc, self.shell_prompt():
|
|
self.process = proc
|
|
sel.register(proc.stdout, selectors.EVENT_READ, self.handle_stdout)
|
|
sel.register(proc.stderr, selectors.EVENT_READ, self.handle_stderr)
|
|
sel.register(user_input_rd, selectors.EVENT_READ, self.handle_user)
|
|
interrupts = 0
|
|
while True:
|
|
try:
|
|
events = sel.select()
|
|
_logger.debug('got events: %r', events)
|
|
if not events:
|
|
break
|
|
for key, mask in events:
|
|
key.data(key)
|
|
interrupts = 0
|
|
except KeyboardInterrupt:
|
|
print('Interrupt again to abort immediately.', file=sys.stderr)
|
|
interrupts += 1
|
|
if interrupts >= 2:
|
|
break
|
|
if self.want_command:
|
|
self.ask_for_command_and_run_it()
|
|
# Process died? Stop the shell.
|
|
if proc.poll() is not None:
|
|
break
|
|
|
|
def handle_user(self, key):
|
|
'''Received a notification that another prompt can be displayed.'''
|
|
os.read(key.fileobj, 4096)
|
|
self.want_command = True
|
|
|
|
def ask_for_command_and_run_it(self):
|
|
cmd = self.read_command()
|
|
if not cmd:
|
|
# Give a chance for the event loop to run again.
|
|
self.want_input()
|
|
return
|
|
self.want_command = False
|
|
_logger.debug('Running: %r', cmd)
|
|
self.process.stdin.write((cmd + '\n').encode('utf8'))
|
|
self.process.stdin.flush()
|
|
|
|
def handle_stdout(self, key):
|
|
resp = key.fileobj.readline().decode('utf8')
|
|
_logger.debug('Response: %r', resp)
|
|
if not resp:
|
|
raise EOFError
|
|
self.want_input()
|
|
resp = resp.strip()
|
|
if resp:
|
|
try:
|
|
if self.pretty == 'jq':
|
|
subprocess.run(['jq', '.'], input=resp,
|
|
universal_newlines=True)
|
|
elif self.pretty == 'indent':
|
|
r = json.loads(resp)
|
|
json.dump(r, sys.stdout, indent=' ')
|
|
print('')
|
|
else:
|
|
print(resp)
|
|
except Exception as e:
|
|
_logger.warning('Dumping output as-is due to: %s', e)
|
|
print(resp)
|
|
|
|
def handle_stderr(self, key):
|
|
data = key.fileobj.read1().decode('utf8')
|
|
print(data, end="", file=sys.stderr)
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--debug', action='store_true',
|
|
help='Enable verbose logging')
|
|
parser.add_argument('--pretty', choices=['jq', 'indent'],
|
|
help='Pretty print responses (one of: %(choices)s)')
|
|
parser.add_argument('--histfile',
|
|
help='Log shell history to this file')
|
|
|
|
|
|
def main(args):
|
|
logging.basicConfig()
|
|
_logger.setLevel(logging.DEBUG if args.debug else logging.INFO)
|
|
shell = SharkdShell(args.pretty, args.histfile)
|
|
try:
|
|
shell.main_loop()
|
|
except EOFError:
|
|
print('')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main(parser.parse_args())
|