Repository: yolothreat/utilitybelt

Issue:
`vt_name_check()` and `vt_api_check()` do not handle errors properly
If VirusTotal returns an error (such as exceeding rate limits), these functions throw a `JSONDecodeError`.

Discussion:

Patch:
diff --git a/.gitignore b/.gitignore
index c26136f..601e383 100644
--- a/.gitignore
+++ b/.gitignore
@@ -57,6 +57,7 @@ venv
 
 ## Data
 data/*
+*.json
 
 # swap files
 *.swp
diff --git a/utilitybelt/utilitybelt.py b/utilitybelt/utilitybelt.py
index 180865c..8307e7a 100644
--- a/utilitybelt/utilitybelt.py
+++ b/utilitybelt/utilitybelt.py
@@ -143,6 +143,24 @@ def is_url(url):
     return bool(re.match(re_url, url))
 
 
+def is_hash(fhash):
+    """Returns true for valid hashes, false for invalid."""
+
+    # Intentionally doing if/else statement for ease of testing and reading
+    if re.match(re_md5, fhash):
+        return True
+    elif re.match(re_sha1, fhash):
+        return True
+    elif re.match(re_sha256, fhash):
+        return True
+    elif re.match(re_sha512, fhash):
+        return True
+    elif re.match(re_ssdeep, fhash):
+        return True
+    else:
+        return False
+
+
 def ip_to_geo(ipaddress):
     """Convert IP to Geographic Information"""
 
@@ -245,7 +263,10 @@ def vt_ip_check(ip, vt_api):
     url = 'https://www.virustotal.com/vtapi/v2/ip-address/report'
     parameters = {'ip': ip, 'apikey': vt_api}
     response = requests.get(url, params=parameters)
-    return response.json()
+    try:
+        return response.json()
+    except ValueError:
+        return None
 
 
 def vt_name_check(domain, vt_api):
@@ -256,7 +277,24 @@ def vt_name_check(domain, vt_api):
     url = 'https://www.virustotal.com/vtapi/v2/domain/report'
     parameters = {'domain': domain, 'apikey': vt_api}
     response = requests.get(url, params=parameters)
-    return response.json()
+    try:
+        return response.json()
+    except ValueError:
+        return None
+
+
+def vt_hash_check(fhash, vt_api):
+    """Checks VirusTotal for occurrences of a file hash"""
+    if not is_hash(fhash):
+        return None
+
+    url = 'https://www.virustotal.com/vtapi/v2/file/report'
+    parameters = {'resource': fhash, 'apikey': vt_api}
+    response = requests.get(url, params=parameters)
+    try:
+        return response.json()
+    except ValueError:
+        return None
 
 
 def ipinfo_ip_check(ip):


Test patch:
diff --git a/utilitybelt/tests/tests.py b/utilitybelt/tests/tests.py
index b826966..ea1f525 100644
--- a/utilitybelt/tests/tests.py
+++ b/utilitybelt/tests/tests.py
@@ -1,4 +1,5 @@
 import os
+import time
 import unittest
 
 from utilitybelt import utilitybelt as ub
@@ -94,6 +95,21 @@ class TestUB(unittest.TestCase):
         self.assertFalse(ub.is_reserved("8.8.4.4"))
         self.assertFalse(ub.is_reserved("192.30.252.131"))
 
+    def test_is_hash(self):
+        # all hashes of the empty string
+        self.assertIsInstance(ub.is_hash("d41d8cd98f00b204e9800998ecf8427e"), bool)
+        # MD5
+        self.assertTrue(ub.is_hash("d41d8cd98f00b204e9800998ecf8427e"))
+        # SHA1
+        self.assertTrue(ub.is_hash("da39a3ee5e6b4b0d3255bfef95601890afd80709"))
+        # SHA256
+        self.assertTrue(ub.is_hash("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"))
+        # SHA512
+        self.assertTrue(ub.is_hash("cf83e1357eefb8bdf1542850d66d8007d620e4050b5715dc83f4a921d36ce9ce47d0d13c5d85f2b0ff8318d2877eec2f63b931bd47417a81a538327af927da3e"))
+        # ssdeep
+        self.assertTrue(ub.is_hash("96:EQOJvOl4ab3hhiNFXc4wwcweomr0cNJDBoqXjmAHKX8dEt001nfEhVIuX0dDcs:3mzpAsZpprbshfu3oujjdENdp21"))
+        self.assertFalse(ub.is_hash("KilroyWasHere"))
+
     def test_vt_ip_check(self):
         vt_api = os.environ["VT_API"]
         self.assertIsNone(ub.vt_ip_check('asdf', vt_api))
@@ -108,6 +124,7 @@ class TestUB(unittest.TestCase):
             if resolution['hostname'] == "github.com":
                 is_gh = True
         self.assertTrue(is_gh)
+        time.sleep(15)  # VT rate limiting
 
     def test_vt_name_check(self):
         vt_api = os.environ["VT_API"]
@@ -120,6 +137,37 @@ class TestUB(unittest.TestCase):
             if resolution['ip_address'] == '192.30.252.130':
                 is_gh = True
         self.assertTrue(is_gh)
+        time.sleep(15)  # VT rate limiting
+
+    def test_vt_hash_check(self):
+        vt_api = os.environ["VT_API"]
+        self.assertIsNone(ub.vt_hash_check('asdf', vt_api))
+        vt_hash_data = ub.vt_hash_check("fe03b4181707f1ea1f3c69dc0a9904181c6fce91", vt_api)
+        self.assertIsInstance(vt_hash_data, dict)
+        self.assertIn('resource', vt_hash_data)
+        self.assertIn('positives', vt_hash_data)
+        self.assertGreater(vt_hash_data['positives'], 0)
+        time.sleep(15)  # VT rate limiting
+        vt_hash_data = ub.vt_hash_check("d41d8cd98f00b204e9800998ecf8427e", vt_api)
+        self.assertIn('positives', vt_hash_data)
+        self.assertEqual(vt_hash_data['positives'], 0)
+        time.sleep(15)  # VT rate limiting
+
+    def test_vt_rate_limiting(self):
+        vt_api = os.environ["VT_API"]
+        # Exceed 4x in 60 seconds
+        data = ub.vt_hash_check("d41d8cd98f00b204e9800998ecf8427e", vt_api)
+        self.assertIsInstance(data, dict)
+        data = ub.vt_hash_check("d41d8cd98f00b204e9800998ecf8427e", vt_api)
+        data = ub.vt_hash_check("d41d8cd98f00b204e9800998ecf8427e", vt_api)
+        data = ub.vt_hash_check("d41d8cd98f00b204e9800998ecf8427e", vt_api)
+        data = ub.vt_name_check("example.org", vt_api)
+        self.assertIsNone(data)
+        data = ub.vt_ip_check("192.30.252.130", vt_api)
+        self.assertIsNone(data)
+        data = ub.vt_hash_check("d41d8cd98f00b204e9800998ecf8427e", vt_api)
+        self.assertIsNone(data)
+        time.sleep(15)
 
     def test_ipinfo(self):
         self.assertIsNone(ub.ipinfo_ip_check('asdf'))