From bf8df3bc02a9633869d732f65cd36a5202a6d174 Mon Sep 17 00:00:00 2001
From: shmilylty <admin@hackfun.org>
Date: Mon, 19 Aug 2019 18:10:44 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E8=BF=9B=E5=BA=A6=E6=9D=A1?=
 =?UTF-8?q?=20=E4=BF=AE=E5=A4=8D=E7=BA=BF=E7=A8=8B=E9=80=80=E5=87=BA?=
 =?UTF-8?q?=E5=8D=A1=E6=AD=BBbug?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 oneforall/takeover.py | 40 +++++++++++++++++++++++++++++-----------
 1 file changed, 29 insertions(+), 11 deletions(-)

diff --git a/oneforall/takeover.py b/oneforall/takeover.py
index 331ab43..9a0ebb0 100644
--- a/oneforall/takeover.py
+++ b/oneforall/takeover.py
@@ -13,6 +13,7 @@ from threading import Thread
 from queue import Queue
 import fire
 from tablib import Dataset
+from tqdm import tqdm
 import config
 from common import resolve, utils
 from common.module import Module
@@ -60,7 +61,7 @@ class Takeover(Module):
     :param str format:  导出格式(默认xls)
     :param str dpath:   导出目录(默认None)
     """
-    def __init__(self, target, thread=10, dpath=None, format='xls'):
+    def __init__(self, target, thread=100, dpath=None, format='xls'):
         Module.__init__(self)
         self.subdomains = set()
         self.module = 'Check'
@@ -70,7 +71,8 @@ class Takeover(Module):
         self.dpath = dpath
         self.format = format
         self.fingerprints = None
-        self.domainq = Queue()
+        self.subdomainq = Queue()
+        self.bar = tqdm()
         self.cnames = list()
         self.results = Dataset()
 
@@ -96,12 +98,12 @@ class Takeover(Module):
                 break
 
     def check(self):
-        while not self.domainq.empty():  # 保证域名队列遍历结束后能退出线程
-            subdomain = self.domainq.get()  # 从队列中获取域名
+        while not self.subdomainq.empty():  # 保证域名队列遍历结束后能退出线程
+            subdomain = self.subdomainq.get()  # 从队列中获取域名
             cname = get_cname(subdomain)
             maindomain = get_maindomain(cname)
             if cname is None:
-                return
+                continue
             for fingerprint in self.fingerprints:
                 cnames = fingerprint.get('cname')
                 if maindomain not in cnames:
@@ -109,18 +111,33 @@ class Takeover(Module):
                 responses = fingerprint.get('response')
                 self.compare(subdomain, cname, responses)
 
+    def progress(self):
+        while not self.subdomainq.empty():
+            done = self.bar.total - self.subdomainq.qsize()
+            self.bar.n = done
+            self.bar.update()
+
     def run(self):
         start = time.time()
-        logger.log('INFOR', f'正在检查子域接管风险')
+        logger.log('INFOR', f'开始执行{self.source}模块')
         self.format = utils.check_format(self.format)
         self.dpath = utils.check_dpath(self.dpath)
         self.subdomains = utils.get_domains(self.target)
         if self.subdomains:
+            logger.log('INFOR', f'正在检查子域接管风险')
             self.fingerprints = get_fingerprint()
             self.results.headers = ['subdomain', 'cname']
+            # 创建待检查的子域队列
             for domain in self.subdomains:
-                self.domainq.put(domain)
+                self.subdomainq.put(domain)
+            # 设置进度大小
+            self.bar.total = self.subdomainq.qsize()
+            # 进度线程
             threads = []
+            thread = Thread(target=self.progress, daemon=True)
+            thread.start()
+            threads.append(thread)
+            # 检查线程
             for _ in range(self.thread):
                 thread = Thread(target=self.check, daemon=True)
                 thread.start()
@@ -134,10 +151,11 @@ class Takeover(Module):
         elapsed = round(end - start, 1)
         logger.log('INFOR', f'{self.source}模块耗时{elapsed}秒'
                             f'发现{len(self.results)}个子域存在接管风险')
+        logger.log('INFOR', f'结束执行{self.source}模块')
 
 
 if __name__ == '__main__':
-    fire.Fire(Takeover)
-    # takeover = Takeover('./subdomains.txt')
-    # takeover = Takeover('www.baidu.com')
-    # takeover.run()
+    # fire.Fire(Takeover)
+    # takeover = Takeover('www.example.com')
+    takeover = Takeover('./subdomains.txt')
+    takeover.run()