Merge pull request #894 from TheHive-Project/feature/vtdown_fix

fix case id, error if type is not guessed
TheHive-Project · Nov 11, 2020 · 3a7dbfc · 3a7dbfc
2 parents 1f0a04a + 890b40c
commit 3a7dbfc
Showing 1 changed file with 56 additions and 20 deletions.
diff --git a/responders/VirustotalDownloader/VirustotalDownloader.py b/responders/VirustotalDownloader/VirustotalDownloader.py
@@ -11,23 +11,33 @@
 from thehive4py.api import TheHiveApi
 from thehive4py.models import Case, CaseObservable
 
+
 class VirustotalDownloader(Responder):
     def __init__(self):
         Responder.__init__(self)
-        self.virustotal_apikey = self.get_param('config.virustotal_apikey', None, "Virustotal API key missing!")
-        self.thehive_url = self.get_param('config.thehive_url', None, "TheHive URL missing!")
-        self.thehive_apikey = self.get_param('config.thehive_apikey', None, "TheHive API key missing!")
+        self.virustotal_apikey = self.get_param(
+            "config.virustotal_apikey", None, "Virustotal API key missing!"
+        )
+        self.thehive_url = self.get_param(
+            "config.thehive_url", None, "TheHive URL missing!"
+        )
+        self.thehive_apikey = self.get_param(
+            "config.thehive_apikey", None, "TheHive API key missing!"
+        )
 
     def run(self):
         Responder.run(self)
 
-        data_type = self.get_param('data.dataType')
-        case_id = self.get_param('data._parent')
+        data_type = self.get_param("data.dataType")
+        case_id = self.get_param("data.case._id")
         ioc_types = ["hash"]
 
         if data_type in ioc_types:
-            url = 'https://www.virustotal.com/vtapi/v2/file/download'
-            params = {'apikey': self.virustotal_apikey, 'hash': self.get_param('data.data')}
+            url = "https://www.virustotal.com/vtapi/v2/file/download"
+            params = {
+                "apikey": self.virustotal_apikey,
+                "hash": self.get_param("data.data"),
+            }
 
             response = requests.get(url, params=params)
 
@@ -36,38 +46,64 @@ def run(self):
                 downloaded_file = response.content
 
                 tempdir = tempfile.gettempdir()
-                f = open(tempdir + "/" +  self.get_param('data.data'), 'wb')
+                f = open(tempdir + "/" + self.get_param("data.data"), "wb")
                 f.write(downloaded_file)
                 f.close()
                 filename = f.name
 
                 kind = filetype.guess(f.name)
 
-                if kind.extension != None:
+                api = TheHiveApi(self.thehive_url, self.thehive_apikey)
+
+                if kind and kind.extension != None:
                     os.rename(f.name, f.name + "." + kind.extension)
                     filename = f.name + "." + kind.extension
 
-                api = TheHiveApi(self.thehive_url, self.thehive_apikey)
-
-                file_observable = CaseObservable(dataType='file',
+                    file_observable = CaseObservable(
+                        dataType="file",
                         data=[filename],
-                        tlp=self.get_param('data.tlp'),
+                        tlp=self.get_param("data.tlp"),
+                        ioc=True,
+                        tags=[
+                            "src:VirusTotal",
+                            str(kind.mime),
+                            str(kind.extension),
+                            "parent:" + self.get_param("data.data"),
+                        ],
+                        message="",
+                    )
+                else:
+                    file_observable = CaseObservable(
+                        dataType="file",
+                        data=[f.name],
+                        tlp=self.get_param("data.tlp"),
                         ioc=True,
-                        tags=['src:VirusTotal', str(kind.mime), str(kind.extension), 'parent:' + self.get_param('data.data')],
-                        message=''
-                        )
+                        tags=[
+                            "src:VirusTotal",
+                            "parent:" + self.get_param("data.data"),
+                        ],
+                        message="",
+                    )
 
                 response = api.create_case_observable(case_id, file_observable)
 
-                self.report({'message': str(response.status_code) + " " + response.text})
+                self.report(
+                    {"message": str(response.status_code) + " " + response.text}
+                )
             else:
-               self.report({'message': 'Virustotal returned the following error code: ' + str(response.status_code) + ". If you receive 403 this means that you are using a public API key but this responder needs a private Virustotal API key!"})
+                self.report(
+                    {
+                        "message": "Virustotal returned the following error code: "
+                        + str(response.status_code)
+                        + ". If you receive 403 this means that you are using a public API key but this responder needs a private Virustotal API key!"
+                    }
+                )
         else:
             self.error('Incorrect dataType. "Hash" expected.')
 
     def operations(self, raw):
-        return [self.build_operation('AddTagToArtifact', tag='Virustotal:Downloaded')]
+        return [self.build_operation("AddTagToArtifact", tag="Virustotal:Downloaded")]
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     VirustotalDownloader().run()