mtda-cli

#!/usr/bin/env python3

# System imports
import daemon
import getopt
import lockfile
import os
import os.path
import requests
import signal
import sys
import zerorpc

# Local imports
from mtda.main import MentorTestDeviceAgent
from mtda.client import Client

PASTEBIN_API_KEY  = "1a265f6e04cf3c1df5e153018390eb29"
PASTEBIN_ENDPOINT = "http://pastebin.com/api/api_post.php"

class Application:

    def __init__(self):
        self.agent  = None
        self.remote = None
        self.logfile = "/var/log/mtda.log"
        self.pidfile = "/var/run/mtda.pid"
        self.exiting = False

    def daemonize(self):
        context = daemon.DaemonContext(
            working_directory=os.getcwd(),
            stdout=open(self.logfile, 'w+'),
            stderr=open(self.logfile, 'w+'),
            umask=0o002,
            pidfile=lockfile.FileLock(self.pidfile)
        )

        context.signal_map = {
            signal.SIGTERM: 'terminate',
            signal.SIGHUP:  'terminate',
        }

        with context:
            status = self.server()
            return status

    def server(self):
        # Start our agent
        status = self.agent.start()
        if status == False:
            return False

        # Start our RPC server
        uri = "tcp://*:%d" % (self.agent.ctrlport)
        s = zerorpc.Server(self.agent, heartbeat=20)
        s.bind(uri)
        s.run()
        return True

    def client(self):
        return self.agent

    def console_clear(self, args):
        self.client().console_clear()

    def console_flush(self, args):
        data = self.client().console_flush()
        if data is not None:
            sys.stdout.write(data)
            sys.stdout.flush()

    def console_head(self, args):
        line = self.client().console_head()
        if line is not None:
            sys.stdout.write(line)
            sys.stdout.flush()
    
    def console_lines(self, args):
        lines = self.client().console_lines()
        sys.stdout.write("%d\n" % (lines))
        sys.stdout.flush()

    def console_interactive(self, args=None):
        client = self.agent
        server = self.client()

        # Print target information
        self.target_info()

        # Connect to the console
        client.console_remote(self.remote)

        # Input loop
        while self.exiting == False:
            c = client.console_getkey()
            if c == '\x01':
                c = client.console_getkey()
                self.console_menukey(c)
            else:
                server.console_send(c, True)

    def console_menukey(self, c):
        server = self.client()
        if c == 'a':
            status = server.target_lock()
            if status == True:
                server.console_print("\n*** Target was acquired ***\n")
        elif c == 'b':
            self.console_pastebin()
        elif c == 'i':
            self.target_info()
        elif c == 'p':
            previous_status = server.target_status()
            server.target_toggle()
            new_status = server.target_status()
            if previous_status != new_status:
                server.console_print("\n*** Target is now %s ***\n" % (new_status))
        elif c == 'q':
            self.exiting = True
        elif c == 'r':
            status = server.target_unlock()
            if status == True:
                server.console_print("\n*** Target was released ***\n")
        elif c == 's':
            previous_status = server.sd_status()
            server.sd_toggle()
            new_status = server.sd_status()
            if new_status != previous_status:
                server.console_print("\n*** SD now connected to %s ***\n" % (new_status))
        elif c == 't':
            server.toggle_timestamps()
        elif c == 'u':
            server.usb_toggle(1)

    def console_pastebin(self):
        data = {
                'api_dev_key'      : PASTEBIN_API_KEY,
                'api_option'       : 'paste',
                'api_paste_code'   : self.client().console_flush(),
                'api_paste_format' : 'python'
               }
        r = requests.post(url = PASTEBIN_ENDPOINT, data = data)
        server = self.client()
        server.console_print("\n*** console buffer posted to %s ***\n" % (r.text))

    def console_prompt(self, args):
        prompt = None
        if len(args) > 0:
            prompt = args[0]
        data = self.client().console_prompt(prompt)
        if data is not None:
            sys.stdout.write(data)
            sys.stdout.flush()

    def console_run(self, args):
        data = self.client().console_run(args[0])
        if data is not None:
            sys.stdout.write(data)
            sys.stdout.flush()

    def console_send(self, args):
        self.client().console_send(args[0])

    def console_tail(self, args):
        line = self.client().console_tail()
        if line is not None:
            sys.stdout.write(line)
            sys.stdout.flush()

    def console_help(self, args=None):
       print("The 'console' command accepts the following sub-commands:")
       print("   clear         Clear any data present in the console buffer")
       print("   flush         Flush content of the console buffer")
       print("   head          Fetch and print the first line from the console buffer")
       print("   interactive   Open the device console for interactive use")
       print("   lines         Print number of lines present in the console buffer")
       print("   prompt        Configure or print the target shell prompt")
       print("   run           Run the specified command via the device console")
       print("   send          Send characters to the device console")
       print("   tail          Fetch and print the last line from the console buffer")

    def console_cmd(self, args):
        if len(args) > 0:
            cmd = args[0]
            args.pop(0)

            cmds = {
               'clear'       : self.console_clear,
               'flush'       : self.console_flush,
               'head'        : self.console_head,
               'interactive' : self.console_interactive,
               'lines'       : self.console_lines,
               'prompt'      : self.console_prompt,
               'run'         : self.console_run,
               'send'        : self.console_send,
               'tail'        : self.console_tail
            }

            if cmd in cmds:
                cmds[cmd](args)
            else:
                print("unknown console command '%s'!" %(cmd), file=sys.stderr)

    def help_cmd(self, args=None):
        if args is not None and len(args) > 0:
            cmd = args[0]
            args.pop(0)

            cmds = {
               'console' : self.console_help,
               'sd'      : self.sd_help,
               'target'  : self.target_help,
               'usb'     : self.usb_help
            }

            if cmd in cmds:
                cmds[cmd](args)
            else:
                print("no help found for command '%s'!" %(cmd), file=sys.stderr)
        else:
            print("usage: mtda [options] <command> [<args>]")
            print("")
            print("The most commonly used mtda commands are:")
            print("   console   Interact with the device console")
            print("   target    Power control the device")
            print("   sd        Interact with the device SD card")
            print("   usb       Control USB devices attached to the device")
            print("")

    def sd_cmd(self, args):
        if len(args) > 0:
            cmd = args[0]
            args.pop(0)

            cmds = {
               'host'   : self.sd_host,
               'mount'  : self.sd_mount,
               'target' : self.sd_target,
               'update' : self.sd_update,
               'write'  : self.sd_write
            }

            if cmd in cmds:
                return cmds[cmd](args)
            else:
                print("unknown sd command '%s'!" %(cmd), file=sys.stderr)
                return 1

    def sd_help(self, args=None):
       print("The 'sd' command accepts the following sub-commands:")
       print("   host      Attach the device SD card to the host")
       print("   mount     Mount the device SD card on the host")
       print("   target    Attach the device SD card to the target")
       print("   update    Update the specified file on the SD card")
       print("   write     Write an image to the device SD card")

    def sd_host(self, args=None):
        status = self.client().sd_to_host()
        if status == False:
            print("failed to connect the device SD card to the host!", file=sys.stderr)
            return 1
        return 0

    def sd_mount(self, args=None):
        status = self.sd_host()
        if status != 0:
            return 1
        part = args[0] if len(args) > 0 else None
        status = self.agent.sd_mount(part)
        if status == False:
            print("'sd mount' failed!", file=sys.stderr)
            return 1
        return 0

    def sd_target(self, args):
        status = self.client().sd_to_target()
        if status == False:
            print("failed to connect the device SD card to the target!", file=sys.stderr)
            return 1
        return 0

    def _sd_write_cb(self, imgname, totalread, imgsize):
        # Print progress
        progress = int((float(totalread) / float(imgsize)) * float(100))
        blocks = int(round((20 * progress) / 100))
        spaces = ' ' * (20 - blocks)
        blocks = '#' * blocks
        totalread = int(totalread / 1024 / 1024)
        totalwritten = int(self.agent.sd_bytes_written() / 1024 / 1024)
        sys.stdout.write("\r{0}: [{1}] {2}% ({3} MiB read, {4} MiB written) "
            .format(imgname, str(blocks + spaces), progress, totalread, totalwritten))
        sys.stdout.flush()

    def sd_update(self, args=None):
        if len(args) == 0:
            print("'sd update' expects a file argument!", file=sys.stderr)
            return 1

        dest = args[0]
        src  = args[1] if len(args) > 1 else None

        status = self.agent.sd_update(dest, src, self._sd_write_cb)
        sys.stdout.write("\n")
        sys.stdout.flush()

        if status == False:
            print("'sd update' failed!", file=sys.stderr)
            return 1
        return 0

    def sd_write(self, args=None):
        if len(args) == 0:
            print("'sd write' expects a file argument!", file=sys.stderr)
            return 1

        status = self.agent.sd_write_image(args[0], self._sd_write_cb)
        sys.stdout.write("\n")
        sys.stdout.flush()

        if status == False:
            print("'sd write' failed!", file=sys.stderr)
            return 1
        return 0

    def target_help(self, args=None):
       print("The 'target' command accepts the following sub-commands:")
       print("   on      Power on the device")
       print("   off     Power off the device")
       print("   toggle  Toggle target power")

    def target_info(self, args=None):
        sys.stdout.write("\rFetching target information...\r")
        sys.stdout.flush()

        # Get general information
        client     = self.client()
        locked     = " (locked)" if client.target_locked() else ""
        remote     = "Local" if self.remote is None else self.remote
        session    = client.session()
        sd_status  = client.sd_status()
        sd_written = client.sd_bytes_written() / 1024 / 1024
        tgt_status = client.target_status()

        # Print general information
        print("Agent     : %s%30s" % (remote, ""))
        print("Session   : %s" % (session))
        print("Target    : %-6s%s" %(tgt_status, locked))
        print("SD on     : %-6s%s" %(sd_status, locked))
        print("SD writes : %u MiB" %(sd_written))

        # Print status of the USB ports
        ports = client.usb_ports()
        for ndx in range(0, ports):
            status = client.usb_status(ndx+1)
            print("USB #%-2d   : %s" %(ndx+1, status))

    def target_off(self, args=None):
        status = self.client().target_off()
        return 0 if (status == True) else 1

    def target_on(self, args=None):
        status = self.client().target_on()
        return 0 if (status == True) else 1

    def target_toggle(self, args=None):
        previous_status = self.client().target_status()
        self.client().target_toggle()
        new_status = self.client().target_status()
        return 0 if (new_status != previous_status) else 1

    def target_cmd(self, args):
        if len(args) > 0:
            cmd = args[0]
            args.pop(0)

            cmds = {
               'help'   : self.target_help,
               'off'    : self.target_off,
               'on'     : self.target_on,
               'toggle' : self.target_toggle
            }

            if cmd in cmds:
                return cmds[cmd](args)
            else:
                print("unknown target command '%s'!" %(cmd), file=sys.stderr)
                return 1

    def usb_cmd(self, args):
        if len(args) > 0:
            cmd = args[0]
            args.pop(0)

            cmds = {
               'help'   : self.usb_help,
               'off'    : self.usb_off,
               'on'     : self.usb_on
            }

            if cmd in cmds:
                return cmds[cmd](args)
            else:
                print("unknown usb command '%s'!" %(cmd), file=sys.stderr)
                return 1

    def usb_help(self, args=None):
       print("The 'usb' command accepts the following sub-commands:")
       print("   on      Power on the specified USB device")
       print("   off     Power off the specified USB device")

    def usb_off(self, args):
        if len(args) > 0:
            klass = args[0]
            args.pop(0)
            client  = self.client()
            result = client.usb_off_by_class(klass)
            return 0 if result else 1
        else:
            print("missing class argument to 'usb off' command!", file=sys.stderr)
            return 1

    def usb_on(self, args):
        if len(args) > 0:
            klass = args[0]
            args.pop(0)
            client  = self.client()
            result = client.usb_on_by_class(klass)
            return 0 if result else 1
        else:
            print("missing class argument to 'usb on' command!", file=sys.stderr)
            return 1

    def main(self):
        daemonize = False
        detach = True

        options, stuff = getopt.getopt(sys.argv[1:], 
            'dnr:',
            ['daemon', 'no-detach', 'remote='])
        for opt, arg in options:
            if opt in ('-d', '--daemon'):
                daemonize = True
            if opt in ('-n', '--no-detach'):
                detach = False 
            if opt in ('-r', '--remote'):
                self.remote = arg

        # Start our server
        if daemonize == True:
            self.agent = MentorTestDeviceAgent() 
            self.agent.load_config(self.remote, daemonize)
            self.remote = self.agent.remote
            if detach == True:
                status = self.daemonize()
            else:
                status = self.server()
            if status == False:
                print('Failed to start the MTDA server!', file=sys.stderr)
                return False
        else:
            # Start our agent
            self.agent = Client(self.remote)
            self.remote = self.agent.remote()
            self.agent.start()

        # Check for non-option arguments
        if len(stuff) > 0:
           cmd = stuff[0]
           stuff.pop(0)

           cmds = {
              'console' : self.console_cmd,
              'help'    : self.help_cmd,
              'sd'      : self.sd_cmd,
              'target'  : self.target_cmd,
              'usb'     : self.usb_cmd
           } 

           if cmd in cmds:
               status = cmds[cmd](stuff)
               sys.exit(status)
           else:
               print("unknown command '%s'!" %(cmd), file=sys.stderr)
               self.help_cmd()
               sys.exit(1)
        else:
            # Assume we want an interactive console if called without a command
            self.console_interactive()
        return True

if __name__ == '__main__':
    app = Application()
    app.main()