Merge pull request #240 from SuperCowPowers/cli-work

Cli work

workbench/clients/customer_report.py

@@ -15,7 +15,7 @@ def run():
     workbench = zerorpc.Client(timeout=300, heartbeat=60)
     workbench.connect('tcp://'+args['server']+':'+args['port'])
 
-    results = workbench.batch_work_request('view_customer', {})
+    results = workbench.batch_work_request('view_customer')
     pprint.pprint(list(results))
 
 def test():

workbench/clients/log_meta_stream.py

@@ -22,11 +22,12 @@ def run():
         with open(filename,'rb') as f:
 
             # Skip OS generated files
-            if '.DS_Store' in filename: continue
+            base_name = os.path.basename(filename)
+            if base_name == '.DS_Store': continue
 
-            md5 = workbench.store_sample(filename, f.read(), 'log')
+            md5 = workbench.store_sample(base_name, f.read(), 'log')
             results = workbench.work_request('view_log_meta', md5)
-            print 'Filename: %s\n' % (filename)
+            print 'Filename: %s\n' % (base_name)
             pprint.pprint(results)
             stream_log = workbench.stream_sample(md5, 20)
             for row in stream_log:

workbench/clients/pcap_bro_indexer.py

@@ -25,12 +25,13 @@ def run():
             continue
 
         with open(filename, 'rb') as pcap_file:
-            md5 = workbench.store_sample(filename, pcap_file.read(), 'pcap')
+            base_name = os.path.basename(filename)
+            md5 = workbench.store_sample(base_name, pcap_file.read(), 'pcap')
 
             # Index the view_pcap output (notice we can ask for any worker output)
             # Also (super important) it all happens on the server side.
             workbench.index_worker_output('view_pcap', md5, 'pcap_bro', None)
-            print '\n\n<<< PCAP Bro log Data: %s Indexed>>>' % (filename)
+            print '\n\n<<< PCAP Bro log Data: %s Indexed>>>' % (base_name)
 
 
 def test():

workbench/clients/pcap_bro_raw.py

@@ -26,7 +26,8 @@ def run():
         if '.DS_Store' in filename: continue
 
         with open(filename,'rb') as f:
-            md5 = workbench.store_sample(filename, f.read(), 'pcap')
+            base_name = os.path.basename(filename)
+            md5 = workbench.store_sample(base_name, f.read(), 'pcap')
             results = workbench.work_request('pcap_bro', md5)
 
             # Results is just a dictionary of Bro log file names and their MD5s in workbench

workbench/clients/pcap_bro_urls.py

@@ -25,7 +25,8 @@ def run():
         if '.DS_Store' in filename: continue
 
         with open(filename,'rb') as f:
-            pcap_md5 = workbench.store_sample(filename, f.read(), 'pcap')
+            base_name = os.path.basename(filename)
+            pcap_md5 = workbench.store_sample(base_name, f.read(), 'pcap')
             results = workbench.work_request('pcap_bro', pcap_md5)
 
             # Just grab the http log

workbench/clients/pcap_bro_view.py

@@ -27,9 +27,10 @@ def run():
 
         # Process the pcap file
         with open(filename,'rb') as f:
-            md5 = workbench.store_sample(filename, f.read(), 'pcap')
+            base_name = os.path.basename(filename)
+            md5 = workbench.store_sample(base_name, f.read(), 'pcap')
             results = workbench.work_request('view_pcap', md5)
-            print '\n<<< %s >>>' % filename
+            print '\n<<< %s >>>' % base_name
             pprint.pprint(results)
 
 def test():

workbench/clients/pcap_meta.py

@@ -24,9 +24,10 @@ def run():
         if '.DS_Store' in filename: continue
 
         with open(filename,'rb') as f:
-            md5 = workbench.store_sample(filename, f.read(), 'pcap')
+            base_name = os.path.basename(filename)
+            md5 = workbench.store_sample(base_name, f.read(), 'pcap')
             results = workbench.work_request('view_pcap', md5)
-            print 'Filename: %s results:' % (filename)
+            print 'Filename: %s results:' % (base_name)
             pprint.pprint(results)
 
 def test():

workbench/clients/pcap_meta_indexer.py

@@ -22,12 +22,13 @@ def run():
         if '.DS_Store' in filename: continue
 
         with open(filename,'rb') as pcap_file:
-            md5 = workbench.store_sample(filename, pcap_file.read(), 'pcap')
+            base_name = os.path.basename(filename)
+            md5 = workbench.store_sample(base_name, pcap_file.read(), 'pcap')
 
             # Index the view_pcap output (notice we can ask for any worker output)
             # Also (super important) it all happens on the server side.
             workbench.index_worker_output('view_pcap', md5, 'view_pcap', None)
-            print '\n\n<<< PCAP Data: %s Indexed>>>' % (filename)
+            print '\n\n<<< PCAP Data: %s Indexed>>>' % (base_name)
 
 def test():
     """Executes pcap_meta_indexer test."""

workbench/clients/pe_indexer.py

@@ -24,14 +24,15 @@ def run():
             continue
 
         with open(filename, 'rb') as f:
-            md5 = workbench.store_sample(filename, f.read(), 'pe')
+            base_name = os.path.basename(filename)
+            md5 = workbench.store_sample(base_name, f.read(), 'exe')
 
             # Index the strings and features output (notice we can ask for any worker output)
             # Also (super important) it all happens on the server side.
             workbench.index_worker_output('strings', md5, 'strings', None)
-            print '\n<<< Strings for PE: %s Indexed>>>' % (filename)
+            print '\n<<< Strings for PE: %s Indexed>>>' % (base_name)
             workbench.index_worker_output('pe_features', md5, 'pe_features', None)
-            print '<<< Features for PE: %s Indexed>>>' % (filename)
+            print '<<< Features for PE: %s Indexed>>>' % (base_name)
 
     # Well we should execute some queries against ElasticSearch at this point but as of
     # version 1.2+ the dynamic scripting disabled by default, see

workbench/clients/pe_peid.py

@@ -26,7 +26,8 @@ def run():
         if '.DS_Store' in filename: continue
 
         with open(filename,'rb') as f:
-            md5 = workbench.store_sample(filename, f.read(), 'pe')
+            base_name = os.path.basename(filename)
+            md5 = workbench.store_sample(base_name, f.read(), 'exe')
             results = workbench.work_request('pe_peid', md5)
             pprint.pprint(results)
 

workbench/clients/pe_sim_graph.py

@@ -20,7 +20,8 @@ def add_it(workbench, file_list, labels):
     for filename in file_list:
         if filename != '.DS_Store':
             with open(filename, 'rb') as pe_file:
-                md5 = workbench.store_sample(filename,  pe_file.read(), 'pe')
+                base_name = os.path.basename(filename)
+                md5 = workbench.store_sample(base_name,  pe_file.read(), 'exe')
                 workbench.add_node(md5, md5[:6], labels)
                 md5s.append(md5)
     return md5s
@@ -92,7 +93,7 @@ def run():
     workbench.clear_graph_db()
 
     # First throw them into workbench and add them as nodes into the graph
-    all_md5s = add_it(workbench, bad_files, ['pe', 'bad']) + add_it(workbench, good_files, ['pe', 'good'])
+    all_md5s = add_it(workbench, bad_files, ['exe', 'bad']) + add_it(workbench, good_files, ['exe', 'good'])
 
     # Compute pe_features on all files of type pe, just pull back the sparse features
     import_gen = workbench.batch_work_request('pe_features',
@@ -139,7 +140,7 @@ def run():
     '''
 
     # Compute pe_deep_sim on all files of type pe
-    results = workbench.batch_work_request('pe_deep_sim', {'type_tag': 'pe'})
+    results = workbench.batch_work_request('pe_deep_sim', {'type_tag': 'exe'})
 
     # Store the ssdeep sims as relationships
     for result in list(results):

workbench/clients/timeout_corner/stress_test.py

@@ -27,7 +27,7 @@ def process_files(path):
     if 'pdf' in path:
         type_tag = 'pdf'
     else:
-        type_tag = 'pe'
+        type_tag = 'exe'
 
     # Open a connection to workbench
     workbench = zerorpc.Client(timeout=300, heartbeat=60)
@@ -41,9 +41,10 @@ def process_files(path):
     for i in xrange(10):
         for filename in file_list:
             with open(filename, 'rb') as f:
-                md5 = workbench.store_sample(filename, f.read(), type_tag)
+                base_name = os.path.basename(filename)
+                md5 = workbench.store_sample(base_name, f.read(), type_tag)
                 workbench.work_request('view', md5)
-                print 'Filename: %s' % (filename)
+                print 'Filename: %s' % (base_name)
         total_files += num_files
     end = datetime.datetime.now()
     delta = end - start

workbench/clients/upload_file.py

@@ -33,7 +33,7 @@ def run():
         # Check to see if workbench already has the file
         filename = os.path.basename(my_file)
         raw_bytes = f.read()
-        md5 = workbench.store_sample(filename, raw_bytes, 'pe')
+        md5 = workbench.store_sample(filename, raw_bytes, 'exe')
         results = workbench.work_request('view', md5)
         print 'Filename: %s' % filename
         pprint.pprint(results)

workbench/clients/workbench

@@ -61,12 +61,17 @@ def run():
 
     # 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]
@@ -91,13 +96,25 @@ def run():
 
     # Do they want a batch_work_request?
     elif command == 'batch_work_request':
-        output = workbench(command, parameters[0], {})
+        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:
-        output = workbench(command, *parameters)
+        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):

workbench/clients/zip_file_extraction.py

@@ -20,15 +20,16 @@ def run():
     file_list = [os.path.join(data_path, child) for child in os.listdir(data_path)]
     for filename in file_list:
         with open(filename,'rb') as f:
-            md5 = workbench.store_sample(filename, f.read(), 'zip')
+            base_name = os.path.basename(filename)
+            md5 = workbench.store_sample(base_name, f.read(), 'zip')
             results = workbench.work_request('view', md5)
-            print 'Filename: %s ' % (filename)
+            print 'Filename: %s ' % (base_name)
             pprint.pprint(results)
 
             # The unzip worker gives you a list of md5s back
             # Run meta on all the unzipped files.
             results = workbench.work_request('unzip', md5)
-            print '\n*** Filename: %s ***' % (filename)
+            print '\n*** Filename: %s ***' % (base_name)
             for child_md5 in results['unzip']['payload_md5s']:
                 pprint.pprint(workbench.work_request('meta', child_md5))
 

workbench/notebooks/Adding_Worker.ipynb

@@ -172,7 +172,7 @@
       "# Okay lets load up a file, and see what this silly meta thing gives back\n",
       "filename = '../data/pe/bad/9e42ff1e6f75ae3e60b24e48367c8f26'\n",
       "with open(filename,'rb') as f:\n",
-      "    my_md5 = c.store_sample(filename, f.read(),'pe')\n",
+      "    my_md5 = c.store_sample(filename, f.read(),'exe')\n",
       "output = c.work_request('meta', my_md5)\n",
       "output"
      ],
@@ -193,7 +193,7 @@
         "  'length': 51200,\n",
         "  'md5': '9e42ff1e6f75ae3e60b24e48367c8f26',\n",
         "  'mime_type': 'application/x-dosexec',\n",
-        "  'type_tag': 'pe'}}"
+        "  'type_tag': 'exe'}}"
        ]
       }
      ],
@@ -779,7 +779,7 @@
         "   'md5': '9e42ff1e6f75ae3e60b24e48367c8f26',\n",
         "   'mime_type': 'application/x-dosexec',\n",
         "   'peid_Matches': ['Microsoft Visual C++ v7.0'],\n",
-        "   'type_tag': 'pe'}}}"
+        "   'type_tag': 'exe'}}}"
        ]
       }
      ],
@@ -808,7 +808,7 @@
       "working_set = []\n",
       "for filename in file_list:\n",
       "    with open(filename,'rb') as f:\n",
-      "        md5 = c.store_sample(filename, f.read(), 'pe')\n",
+      "        md5 = c.store_sample(filename, f.read(), 'exe')\n",
       "        working_set.append(md5)"
      ],
      "language": "python",

workbench/notebooks/Generator_Pipelines.ipynb

@@ -71,7 +71,7 @@
       "md5_list = []\n",
       "for filename in file_list:\n",
       "    with open(filename,'rb') as f:\n",
-      "        md5_list.append(c.store_sample(filename, f.read(), 'pe'))\n",
+      "        md5_list.append(c.store_sample(filename, f.read(), 'exe'))\n",
       "print 'Files loaded: %d' % len(md5_list)\n",
       "md5_list[:5]"
      ],

workbench/notebooks/PCAP_DriveBy.ipynb

@@ -1125,7 +1125,7 @@
         "  'sha1': '035db69cc80fc56717a42646911d9aa95b2ff39e',\n",
         "  'sha256': 'e4bbdc8f869502183293797f51d6d64cc6c49d39b82effbcb738abe511054b51',\n",
         "  'ssdeep': '6144:RSxqC+ayi6eWLj622ARbJFMQzynbJDxL3oPlRa:oxqC+ayi6p6EmQz+bf3otA',\n",
-        "  'type_tag': 'pe'}]"
+        "  'type_tag': 'exe'}]"
        ]
       }
      ],