diff options
Diffstat (limited to 'mpm/tools')
| -rwxr-xr-x | mpm/tools/mpm_shell.py | 117 |
1 files changed, 83 insertions, 34 deletions
diff --git a/mpm/tools/mpm_shell.py b/mpm/tools/mpm_shell.py index 2997a5eb2..538d54586 100755 --- a/mpm/tools/mpm_shell.py +++ b/mpm/tools/mpm_shell.py @@ -9,6 +9,7 @@ RPC shell to debug USRP MPM capable devices """ from __future__ import print_function +import re import cmd import time import argparse @@ -51,6 +52,11 @@ def parse_args(): '-j', '--hijack', type=str, help="Hijack running session (excludes --claim)." ) + parser.add_argument( + '-s', '--script', type=str, + help="Run shell in scripting mode. Specified script contains " + "MPM shell commands, one per line." + ) return parser.parse_args() @@ -85,35 +91,35 @@ class MPMClaimer(object): """ from mprpc import RPCClient from mprpc.exceptions import RPCError - cmd = None + command = None token = None exit_loop = False client = RPCClient(host, port, pack_params={'use_bin_type': True}) try: while not exit_loop: - if token and not cmd: + if token and not command: client.call('reclaim', token) - elif cmd == 'claim': + elif command == 'claim': if not token: token = client.call('claim', 'MPM Shell') else: print("Already have claim") token_q.put(token) - elif cmd == 'unclaim': + elif command == 'unclaim': if token: client.call('unclaim', token) token = None token_q.put(None) - elif cmd == 'exit': + elif command == 'exit': if token: client.call('unclaim', token) token = None token_q.put(None) exit_loop = True time.sleep(1) - cmd = None + command = None if not cmd_q.empty(): - cmd = cmd_q.get(False) + command = cmd_q.get(False) except RPCError as ex: print("Unexpected RPC error in claimer loop!") print(str(ex)) @@ -152,6 +158,9 @@ class MPMClaimer(object): return self.token def hijack(self, token): + """ + Take over existing session by providing session token. + """ if self.token: print("Already have token") return @@ -163,7 +172,7 @@ class MPMShell(cmd.Cmd): """ RPC Shell class. See cmd module. """ - def __init__(self, host, port, claim, hijack): + def __init__(self, host, port, claim, hijack, script): cmd.Cmd.__init__(self) self.prompt = "> " self.client = None @@ -179,6 +188,9 @@ class MPMShell(cmd.Cmd): elif hijack: self.hijack(hijack) self.update_prompt() + self._script = script + if self._script: + self.parse_script() def _add_command(self, command, docs, requires_token=False): """ @@ -193,6 +205,9 @@ class MPMShell(cmd.Cmd): setattr(self, cmd_name, new_command) self.remote_methods.append(command) + def _print_response(self, response): + print(re.sub("^", "< ", response, flags=re.MULTILINE)) + def rpc_template(self, command, requires_token, args=None): """ Template function to create new RPC shell commands @@ -200,8 +215,9 @@ class MPMShell(cmd.Cmd): from mprpc.exceptions import RPCError if requires_token and \ (self._claimer is None or self._claimer.get_token() is None): - print("Cannot execute '{}' -- no claim available!".format(command)) - return + self._print_response("Cannot execute `{}' -- " + "no claim available!".format(command)) + return False try: if args or requires_token: expanded_args = self.expand_args(args) @@ -211,21 +227,20 @@ class MPMShell(cmd.Cmd): else: response = self.client.call(command) except RPCError as ex: - print("RPC Command failed!") - print("Error: {}".format(ex)) - return + self._print_response("RPC Command failed!\nError: {}".format(ex)) + return False except Exception as ex: - print("Unexpected exception!") - print("Error: {}".format(ex)) - return + self._print_response("Unexpected exception!\nError: {}".format(ex)) + return True if isinstance(response, bool): if response: - print("Command executed successfully!") + self._print_response("Command succeeded.") else: - print("Command failed!") + self._print_response("Command failed!") else: - print("==> " + str(response)) - return response + self._print_response(str(response)) + + return False def get_names(self): " We need this for tab completion. " @@ -234,13 +249,25 @@ class MPMShell(cmd.Cmd): ########################################################################### # Cmd module specific ########################################################################### - def run(self): - " Go, go, go! " - try: - self.cmdloop() - except KeyboardInterrupt: - self.do_disconnect(None) - exit(0) + def default(self, line): + self._print_response("*** Unknown syntax: %s" % line) + + def preloop(self): + """ + In script mode add Execution start marker to ease parsing script output + :return: None + """ + if self._script: + print("Execute %s" % self._script) + + def precmd(self, line): + """ + Add command prepended by "> " in scripting mode to ease parsing script + output. + """ + if self.cmdqueue: + print("> %s" % line) + return line def postcmd(self, stop, line): """ @@ -248,6 +275,7 @@ class MPMShell(cmd.Cmd): - Update prompt """ self.update_prompt() + return stop ########################################################################### # Internal methods @@ -257,7 +285,6 @@ class MPMShell(cmd.Cmd): Launch a connection. """ from mprpc import RPCClient - from mprpc.exceptions import RPCError print("Attempting to connect to {host}:{port}...".format( host=host, port=port )) @@ -342,6 +369,24 @@ class MPMShell(cmd.Cmd): claim_status=claim_status, ) + def parse_script(self): + """ + Adding script command from file pointed to by self._script. + + The commands are read from file one per line and added to cmdqueue of + parent class. This way they will be executed instead of input from + stdin. An EOF command is appended to the list to ensure the shell exits + after script execution. + :return: None + """ + try: + with open(self._script, "r") as script: + for command in script: + self.cmdqueue.append(command.strip()) + except OSError as ex: + print("Failed to read script. (%s)" % ex) + self.cmdqueue.append("EOF") # terminate shell after script execution + def expand_args(self, args): """ Takes a string and returns a list @@ -409,26 +454,30 @@ class MPMShell(cmd.Cmd): """import a python module into the global namespace""" globals()[args] = import_module(args) + # pylint: disable=invalid-name def do_EOF(self, _): - " When catching EOF, exit the program. " + """ + When catching EOF, exit the program. + """ print("Exiting...") self.disconnect() - exit(0) + return True # orderly shutdown def main(): " Go, go, go! " args = parse_args() - my_shell = MPMShell(args.host, args.port, args.claim, args.hijack) + my_shell = MPMShell(args.host, args.port, args.claim, + args.hijack, args.script) try: - return my_shell.run() + my_shell.cmdloop() except KeyboardInterrupt: my_shell.disconnect() - except Exception as ex: + except Exception as ex: # pylint: disable=broad-except print("Uncaught exception: " + str(ex)) my_shell.disconnect() + return False return True if __name__ == "__main__": exit(not main()) - |