From 232014f2215de65b5ebbba3ab80588b321b01de2 Mon Sep 17 00:00:00 2001
From: Hannah Ward <Hannah.ward9001@gmail.com>
Date: Wed, 17 Aug 2016 13:01:11 +0100
Subject: [PATCH] Added virustotal tests

---
 .gitignore                                   |  1 +
 misp_modules/modules/expansion/virustotal.py | 21 ++++++++++++++------
 tests/bodyvirustotal.json.sample             |  1 +
 tests/test.py                                | 13 ++++++++++++
 4 files changed, 30 insertions(+), 6 deletions(-)
 create mode 100644 tests/bodyvirustotal.json.sample

diff --git a/.gitignore b/.gitignore
index 96763964..abab012b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,6 @@
 *.pyc
 *.swp
+test/bodyvirustotal.json
 __pycache__
 build/
 dist/
diff --git a/misp_modules/modules/expansion/virustotal.py b/misp_modules/modules/expansion/virustotal.py
index 930872ed..df7b8b42 100755
--- a/misp_modules/modules/expansion/virustotal.py
+++ b/misp_modules/modules/expansion/virustotal.py
@@ -6,7 +6,7 @@ import base64
 import os
 
 misperrors = {'error': 'Error'}
-mispattributes = {'input': ['domain', "ip-src", "ip-dst"],
+mispattributes = {'input': ['hostname', 'domain', "ip-src", "ip-dst"],
                   'output':['domain', "ip-src", "ip-dst", "text"]
                   }
 
@@ -16,16 +16,19 @@ moduleinfo = {'version': '1', 'author': 'Hannah Ward',
               'module-type': ['expansion']}
 
 # config fields that your code expects from the site admin
-moduleconfig = ["apikey"]
-
+moduleconfig = ["apikey", "event_limit"]
+limit = 5 #Default
 
 def handler(q=False):
+    global limit
     if q is False:
         return False
 
     q = json.loads(q)
 
     key = q["config"]["apikey"]
+    limit = int(q["config"].get("event_limit", 5))
+
     r = {"results": []}
     
     if "ip-src" in q:
@@ -34,6 +37,8 @@ def handler(q=False):
       r["results"] += getIP(q["ip-dst"], key)
     if "domain" in q:
       r["results"] += getDomain(q["domain"], key)
+    if 'hostname' in q:
+      r["results"] += getDomain(q['hostname'], key)
 
     uniq = []
     for res in r["results"]:
@@ -43,6 +48,7 @@ def handler(q=False):
     return r
 
 def getIP(ip, key, do_not_recurse = False):
+    global limit
     print("Getting info for {}".format(ip))
     toReturn = []
     req = requests.get("https://www.virustotal.com/vtapi/v2/ip-address/report", 
@@ -53,7 +59,7 @@ def getIP(ip, key, do_not_recurse = False):
       return []
     
     if "resolutions" in req:
-      for res in req["resolutions"]:
+      for res in req["resolutions"][:limit]:
         toReturn.append( {"types":["domain"], "values":[res["hostname"]]})
         #Pivot from here to find all domain info
         if not do_not_recurse:
@@ -63,6 +69,8 @@ def getIP(ip, key, do_not_recurse = False):
     return toReturn
     
 def getDomain(domain, key, do_not_recurse=False):
+    global limit
+
     print("Getting info for {}".format(domain))
     toReturn = []
     req = requests.get("https://www.virustotal.com/vtapi/v2/domain/report", 
@@ -73,7 +81,7 @@ def getDomain(domain, key, do_not_recurse=False):
       return []
     
     if "resolutions" in req:
-      for res in req["resolutions"]:
+      for res in req["resolutions"][:limit]:
         toReturn.append( {"types":["ip-dst", "ip-src"], "values":[res["ip_address"]]})
         #Pivot from here to find all info on IPs
         if not do_not_recurse:
@@ -103,13 +111,14 @@ def isset(d, key):
     return False
 
 def getMoreInfo(req, key):
+    global limit
     print("Getting extra info for {}".format(req))
     r = []
     #Get all hashes first
     hashes = []
     hashes = findAll(req, ["md5", "sha1", "sha256", "sha512"])
     r.append({"types":["md5", "sha1", "sha256", "sha512"], "values":hashes})
-    for hsh in hashes[:5]:
+    for hsh in hashes[:limit]:
       #Search VT for some juicy info
       data = requests.get("http://www.virustotal.com/vtapi/v2/file/report",
                           params={"allinfo":1, "apikey":key, "resource":hsh}
diff --git a/tests/bodyvirustotal.json.sample b/tests/bodyvirustotal.json.sample
new file mode 100644
index 00000000..43cad6bf
--- /dev/null
+++ b/tests/bodyvirustotal.json.sample
@@ -0,0 +1 @@
+{"module": "virustotal", "ip-dst": "5.104.106.190", "config": {"api_key": "deadbeef"} }
diff --git a/tests/test.py b/tests/test.py
index 1314479a..a8693384 100644
--- a/tests/test.py
+++ b/tests/test.py
@@ -5,6 +5,7 @@ import unittest
 import requests
 import base64
 import json
+import os
 
 class TestModules(unittest.TestCase):
 
@@ -36,5 +37,17 @@ class TestModules(unittest.TestCase):
             response = requests.post(self.url + "query", data=data)
             print(response.json())
 
+    def test_virustotal(self):
+        # This can't actually be tested without disclosing a private
+        # API key. This will attempt to run with a .gitignored keyfile
+        # and pass if it can't find one
+
+        if not os.path.exists("tests/bodyvirustotal.json"):
+          return
+        
+        with open("tests/bodyvirustotal.json", "r") as f:
+          response = requests.post(self.url + "query", data=f.read()).json()
+        assert(response)
+
 if __name__ == '__main__':
   unittest.main()