优化

modules/intelligence/virustotal.py

@@ -10,7 +10,6 @@ class VirusTotal(Query):
         Query.__init__(self)
         self.source = 'VirusTotalQuery'
         self.module = 'Intelligence'
-        self.addr = 'https://www.virustotal.com/ui/domains/{}/subdomains'
         self.domain = domain
 
     def query(self):
@@ -24,26 +23,16 @@ class VirusTotal(Query):
                                 'TE': 'Trailers'})
             self.proxy = self.get_proxy(self.source)
             params = {'limit': '40', 'cursor': next_cursor}
-            resp = self.get(url=self.addr.format(self.domain), params=params)
+            addr = f'https://www.virustotal.com/ui/domains/{self.domain}/subdomains'
+            resp = self.get(url=addr, params=params)
             if not resp:
-                return
-            data = resp.json()
-            subdomains = set()
-            datas = data.get('data')
-
-            if datas:
-                for data in datas:
-                    subdomain = data.get('id')
-                    if subdomain:
-                        subdomains.add(subdomain)
-            else:
+                break
+            subdomains = self.match_subdomains(resp)
+            if not subdomains:
                 break
             self.subdomains = self.subdomains.union(subdomains)
-            meta = data.get('meta')
-            if meta:
-                next_cursor = meta.get('cursor')
-            else:
-                break
+            data = resp.json()
+            next_cursor = data.get('meta').get('cursor')
 
     def run(self):
         """
@@ -68,4 +57,4 @@ def run(domain):
 
 
 if __name__ == '__main__':
-    run('example.com')
+    run('mi.com')

modules/intelligence/virustotal_api.py

@@ -8,24 +8,26 @@ class VirusTotalAPI(Query):
         self.domain = domain
         self.module = 'Intelligence'
         self.source = 'VirusTotalAPIQuery'
-        self.addr = 'https://www.virustotal.com/vtapi/v2/domain/report'
         self.key = settings.virustotal_api_key
 
     def query(self):
         """
         向接口查询子域并做子域匹配
         """
+        next_cursor = ''
+        while True:
             self.header = self.get_header()
+            self.header.update({'x-apikey': self.key})
             self.proxy = self.get_proxy(self.source)
-        params = {'apikey': self.key, 'domain': self.domain}
-        resp = self.get(self.addr, params)
-        if not resp:
-            return
-        json = resp.json()
-        data = json.get('subdomains')
-        if data:
-            subdomains = set(data)
+            params = {'limit': '40', 'cursor': next_cursor}
+            addr = f'https://www.virustotal.com/api/v3/domains/{self.domain}/subdomains'
+            resp = self.get(url=addr, params=params)
+            subdomains = self.match_subdomains(resp)
+            if not subdomains:
+                break
             self.subdomains = self.subdomains.union(subdomains)
+            data = resp.json()
+            next_cursor = data.get('meta').get('cursor')
 
     def run(self):
         """
@@ -52,4 +54,4 @@ def run(domain):
 
 
 if __name__ == '__main__':
-    run('example.com')
+    run('mi.com')