Analyzer improvements

salt/sensoroni/files/analyzers/virustotal/virustotal.py

@@ -0,0 +1,82 @@
+import json
+import requests
+import argparse
+import helpers
+import os
+
+
+def buildHeaders(conf):
+    headers = {"x-apikey": conf["api_key"]}
+    return headers
+
+
+def sendReq(conf, meta, payload, headers):
+    url = conf['base_url']
+    response = requests.request('GET', url + payload, headers=headers)
+    return response.json()
+
+
+def prepareResults(raw):
+    malicious = 0
+    harmless = 0
+    undetected = 0
+    suspicious = 0
+    timeout = 0
+
+    if "data" in raw:
+        entries = raw["data"]
+        for data in entries:
+            if "attributes" in data:
+                attrs = data["attributes"]
+                if "last_analysis_stats" in attrs:
+                    stats = attrs["last_analysis_stats"]
+                    if len(stats) > 0:
+                        suspicious += stats["suspicious"]
+                        malicious += stats["malicious"]
+                        harmless += stats["harmless"]
+                        undetected += stats["undetected"]
+                        timeout += stats["timeout"]
+
+    if malicious > 0:
+        summary = "malicious"
+        status = "threat"
+    elif suspicious > 0:
+        summary = "suspicious"
+        status = "caution"
+    elif timeout > 0:
+        summary = "timeout"
+        status = "caution"
+    elif harmless > 0 or undetected > 0:
+        summary = "harmless"
+        status = "ok"
+    else:
+        summary = "internal_failure"
+        status = "caution"
+
+    results = {'response': raw, 'summary': summary, 'status': status}
+    return results
+
+
+def analyze(conf, input):
+    meta = helpers.loadMetadata(__file__)
+    data = helpers.parseArtifact(input)
+    helpers.checkSupportedType(meta, data["artifactType"])
+    headers = buildHeaders(conf)
+    response = sendReq(conf, meta, data["value"], headers)
+    return prepareResults(response)
+
+
+def main():
+    dir = os.path.dirname(os.path.realpath(__file__))
+    parser = argparse.ArgumentParser(description='Search VirusTotal 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 + "/virustotal.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))
+
+
+if __name__ == "__main__":
+    main()