From 8e2d787036f9b85a90894fa5e6e9834c476c2919 Mon Sep 17 00:00:00 2001
From: Jing Ling <shmilylty@gmail.com>
Date: Tue, 18 Aug 2020 00:45:28 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 modules/intelligence/virustotal.py     | 27 ++++++++------------------
 modules/intelligence/virustotal_api.py | 26 +++++++++++++------------
 2 files changed, 22 insertions(+), 31 deletions(-)

diff --git a/modules/intelligence/virustotal.py b/modules/intelligence/virustotal.py
index 52aeed6..542d87b 100644
--- a/modules/intelligence/virustotal.py
+++ b/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')
diff --git a/modules/intelligence/virustotal_api.py b/modules/intelligence/virustotal_api.py
index 9023e44..b4058a6 100644
--- a/modules/intelligence/virustotal_api.py
+++ b/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):
         """
         向接口查询子域并做子域匹配
         """
-        self.header = self.get_header()
-        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)
+        next_cursor = ''
+        while True:
+            self.header = self.get_header()
+            self.header.update({'x-apikey': self.key})
+            self.proxy = self.get_proxy(self.source)
+            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')