malwarebazaar dep upgrades + use auth

salt/sensoroni/files/analyzers/malwarebazaar/malwarebazaar.py

@@ -2,12 +2,21 @@ import requests
 import helpers
 import json
 import sys
+import os
+import argparse
 
 # supports querying for hash, gimphash, tlsh, and telfhash
 # usage is as follows:
 # python3 malwarebazaar.py '{"artifactType":"x", "value":"y"}'
 
 
+def checkConfigRequirements(conf):
+    if not conf.get('api_key'):
+        sys.exit(126)
+    else:
+        return True
+
+
 def buildReq(observ_type, observ_value):
     # determine correct query type to send based off of observable type
     unique_types = {'gimphash': 1, 'telfhash': 1, 'tlsh': 1}
@@ -18,10 +27,13 @@ def buildReq(observ_type, observ_value):
     return {'query': qtype, observ_type: observ_value}
 
 
-def sendReq(meta, query):
+def sendReq(conf, meta, query):
     # send a post request with our compiled query to the API
     url = meta['baseUrl']
-    response = requests.post(url, query)
+    headers = {}
+    if conf.get('api_key'):
+        headers['Auth-Key'] = conf['api_key']
+    response = requests.post(url, query, headers=headers)
     return response.json()
 
 
@@ -113,10 +125,11 @@ def prepareResults(raw):
     return {'response': raw, 'summary': summary, 'status': status}
 
 
-def analyze(input):
+def analyze(conf, input):
     # put all of our methods together, pass them input, and return
     # properly formatted json/python dict output
-    data = json.loads(input)
+    checkConfigRequirements(conf)
+    data = helpers.parseArtifact(input)
     meta = helpers.loadMetadata(__file__)
     helpers.checkSupportedType(meta, data["artifactType"])
 
@@ -127,7 +140,7 @@ def analyze(input):
         # twice for the sake of retrieving more specific data.
 
         initialQuery = buildReq(data['artifactType'], data['value'])
-        initialRaw = sendReq(meta, initialQuery)
+        initialRaw = sendReq(conf, meta, initialQuery)
 
         # To prevent double-querying when a tlsh/gimphash is invalid,
         # this if statement is necessary.
@@ -140,16 +153,22 @@ def analyze(input):
             return prepareResults(initialRaw)
 
     query = buildReq(data['artifactType'], data['value'])
-    response = sendReq(meta, query)
+    response = sendReq(conf, meta, query)
     return prepareResults(response)
 
 
 def main():
-    if len(sys.argv) == 2:
-        results = analyze(sys.argv[1])
+    dir = os.path.dirname(os.path.realpath(__file__))
+    parser = argparse.ArgumentParser(
+        description='Search MalwareBazaar for a given artifact')
+    parser.add_argument(
+        'artifact', help='the artifact represented in JSON format')
+    parser.add_argument('-c', '--config', metavar='CONFIG_FILE', default=dir + '/malwarebazaar.yaml',
+                        help='optional config file to use instead of the default config file')
+    args = parser.parse_args()
+    if args.artifact:
+        results = analyze(helpers.loadConfig(args.config), args.artifact)
         print(json.dumps(results))
-    else:
-        print("ERROR: Input is not in proper JSON format")
 
 
 if __name__ == '__main__':