Merge pull request #245 from SuperCowPowers/cli-work

Cli work

tox.ini

@@ -9,7 +9,7 @@ deps =
 install =
     pip install -e .
 commands =
-    py.test test/test_server_spinup.py workbench/workers workbench/clients
+    py.test test/test_server_spinup.py workbench/server workbench/workers workbench/clients
     coverage combine
 
 [pytest]

workbench/clients/__init__.py

@@ -1 +0,0 @@
-''' Workbench Clients. '''

workbench/clients/help_client.py

@@ -1,7 +1,5 @@
 ''' This client calls a bunch of help commands from workbench '''
 import zerorpc
-import pprint
-import os
 import workbench_client
 
 def run():
@@ -16,13 +14,11 @@ def run():
 
     # Call help methods
     print workbench.help()
-    print workbench.help_basic()
-    print workbench.help_commands()
-    print workbench.help_command('store_sample')
-    print workbench.help_workers()
-    print workbench.help_worker('meta')
-    print workbench.help_advanced()
-    print workbench.help_everything()
+    print workbench.help('basic')
+    print workbench.help('commands')
+    print workbench.help('store_sample')
+    print workbench.help('workers')
+    print workbench.help('meta')
 
     # Call a test worker
     print workbench.test_worker('meta')

workbench/clients/short_md5s.py

@@ -0,0 +1,40 @@
+"""This client tests workbench support for short md5s """
+
+import zerorpc
+import os
+import pprint
+import workbench_client
+
+def run():
+    """This client tests workbench support for short md5s """
+
+    # Grab server args
+    args = workbench_client.grab_server_args()
+
+    # Start up workbench connection
+    workbench = zerorpc.Client(timeout=300, heartbeat=60)
+    workbench.connect('tcp://'+args['server']+':'+args['port'])
+
+    # Pull in a bunch of files
+    data_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),'../data/pe/bad')
+    file_list = [os.path.join(data_path, child) for child in os.listdir(data_path)]
+    data_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),'../data/pe/good')
+    file_list += [os.path.join(data_path, child) for child in os.listdir(data_path)]
+    for filename in file_list:
+
+        # Skip OS generated files
+        if '.DS_Store' in filename: continue
+
+        with open(filename,'rb') as f:
+            base_name = os.path.basename(filename)
+            md5 = workbench.store_sample(base_name, f.read(), 'exe')
+            results = workbench.work_request('meta', md5[:6])
+            pprint.pprint(results)
+
+
+def test():
+    """Executes short md5 test."""
+    run()
+
+if __name__ == '__main__':
+    run()

workbench/clients/workbench

@@ -1,144 +1,12 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
-"""Workbench CLI Client"""
-import os, sys
-import pprint
-import hashlib
-import zerorpc
-import tabulate
-import workbench.clients.workbench_client as workbench_client
-
-def try_get_tables(data):
-    ''' Try to determine if this data has table keys '''
-
-    table_list = []
-    if 'tables' in data.keys():
-        for table in data['tables']:
-            if isinstance(data[table], list):
-                table_list.append(table)
-
-    # At this point the list has data in it or is an empty list
-    return table_list
-
-def table_xform(table, force_to_str=False):
-    ''' Transform a table from a list of dicts to dict of lists'''
-    import collections
-
-    new_table = collections.defaultdict(list)
-    for row in table:
-        for key, value in row.items():
-            if force_to_str:
-                new_table[key].append(str(value)[:40])
-            elif isinstance(value, str):
-                new_table[key].append(value[:40])
-            else:
-                new_table[key].append(value)
-    return new_table
-
-def print_table(table):
-    ''' Print the table using tabulate '''
-    try:
-        print tabulate.tabulate(table_xform(table), headers="keys")
-    except TypeError:
-        print tabulate.tabulate(table_xform(table, force_to_str=True), headers="keys")
+"""Workbench Interactive Shell using IPython"""
+from workbench.clients import workbench_shell
 
 def run():
-    ''' Running the workbench CLI '''
-    # Announce Version
-    print '<<< Workbench Version %s >>>' % sys.modules['workbench'].__version__
-
-    # Grab arguments
-    args = workbench_client.grab_server_args()
-
-    # Start up workbench connection
-    workbench = zerorpc.Client(timeout=300, heartbeat=60)
-    workbench.connect('tcp://'+args['server']+':'+args['port'])
-
-    # If no command than assume they want help
-    if not args['commands']:
-        args['commands'] = ['help']
-
-    # Check to see if the command is valid if it's not assume they want a 'work_request'
-    command = args['commands'][0]
-    orig_command = command
-    if command in workbench.list_all_commands():
-        parameters = args['commands'][1:]
-    else:
-        command = 'work_request'
-        parameters = args['commands']
-
-    # Do they want a help command ?
-    if 'help' in command:
-        parameters += {'cli':True}
-
-    # Do they want 'store_sample'?
-    if command == 'store_sample':
-        file_path = parameters[0]
-
-        # Do they want everything under a directory?
-        if os.path.isdir(file_path):
-            file_list = [os.path.join(file_path, child) for child in os.listdir(file_path)]
-        else:
-            file_list = [file_path]
-
-        # Upload the files into workbench
-        for path in file_list:
-            with open(path, 'rb') as my_file:
-                raw_bytes = my_file.read()
-                md5 = hashlib.md5(raw_bytes).hexdigest()
-                if not workbench.has_sample(md5):
-                    print 'Storing Sample...'
-                    md5 = workbench.store_sample(os.path.basename(path), raw_bytes, 'unknown')
-                else:
-                    print 'Sample already in Workbench...'
-                print md5
-
-    # Do they want a batch_work_request?
-    elif command == 'batch_work_request':
-        output = workbench(command, parameters[0])
-        for row in output:
-            pprint.pprint(row)
-
-    # Do they want to list_samples?
-    elif command == 'list_samples':
-
-        # They might have a predicate
-        predicate = json.loads(parameters[0]) if len(parameters) else None
-        output = workbench(command, predicate)
-        print_table(output)
-
-    # Okay must be a 'normal' command so send the command and any optional parameters to workbench
-    else:
-        try:
-            output = workbench(command, *parameters)
-        except zerorpc.exceptions.RemoteError:
-            print 'Failed to run command: %s' % orig_command
-            exit(1)
-
-        # Try to do different stuff based on the output
-        if isinstance(output, str):
-            print output
-
-        # Okay for non-string output it will always be keyed by the worker name
-        # which in this case in parameter[0]
-        elif isinstance(output, dict):
-
-            worker_name = parameters[0]
-            output = output[worker_name]
-
-            # If the data contains table keys/information then print as tables
-            # if not we just kinda punt and do a pprint on the data.
-            table_list = try_get_tables(output)
-            if table_list:
-                for table in table_list:
-                    print '\n\n<<<<<< %s >>>>>\n' % table
-                    print_table(output[table])
-            else:
-                pprint.pprint(output)
-        else:
-            print 'Critical: I have no idea what to do with output of type: %s' % type(output)
-
+    work_shell = workbench_shell.WorkbenchShell()
+    work_shell.run()
 
 if __name__ == '__main__':
     run()