MPM: add ability to run scripts to MPM shell

MPM shell now supports script execution. It utilizes the cmdqueue of
Pythons cmd.Cmd class for this. The script to execute is a text file
containing the commands one per line.

The output decoration of MPM shell changed. Commands are decorated with
">" whereas responses use "<" at line start. Multiline responses are
decorated in each line.

Cleanup overwritten methods of cmd.Cmd to allow proper shutdow in
interactive as well as in scripted mode.

Improved pylint score.

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())
-