added more descriptive exception

Xen0ph0n · Oct 9, 2017 · 8044af1 · 8044af1
1 parent 401ed9b
commit 8044af1
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 10 deletions.
diff --git a/vt.py b/vt.py
@@ -11,11 +11,14 @@
 
 class vtAPI():
     def __init__(self):
-        your_private_api_key = '<--------------PRIVATE-API-KEY-GOES-HERE----->'
-        assert "KEY-GOES-HERE" not in your_private_api_key
-        self.api = your_private_api_key
-        self.base = 'https://www.virustotal.com/vtapi/v2/'
-
+        try:
+            your_private_api_key = '<--------------PRIVATE-API-KEY-GOES-HERE----->'
+            assert "KEY-GOES-HERE" not in your_private_api_key
+            self.api = your_private_api_key
+            self.base = 'https://www.virustotal.com/vtapi/v2/'
+        except AssertionError:
+            raise ValueError("Private API key not provided.")
+
     def getReport(self,md5):
         param = {'resource':md5,'apikey':self.api,'allinfo': '1'}
         url = self.base + "file/report"

diff --git a/vtlite.py b/vtlite.py
@@ -11,11 +11,14 @@
 
 class vtAPI():
     def __init__(self):
-        your_public_api_key = '<--------------PUBLIC-API-KEY-GOES-HERE----->'
-        assert "KEY-GOES-HERE" not in your_public_api_key
-        self.api = your_public_api_key
-        self.base = 'https://www.virustotal.com/vtapi/v2/'
-
+        try:
+            your_public_api_key = '<--------------PUBLIC-API-KEY-GOES-HERE----->'
+            assert "KEY-GOES-HERE" not in your_public_api_key
+            self.api = your_public_api_key
+            self.base = 'https://www.virustotal.com/vtapi/v2/'
+        except AssertionError:
+            raise ValueError("Private API key not provided.")
+
     def getReport(self,md5):
         param = {'resource':md5,'apikey':self.api}
         url = self.base + "file/report"