From b60d86fdc30ba9e218dc14f2124c58a241c6bd14 Mon Sep 17 00:00:00 2001
From: shmilylty <admin@hackfun.org>
Date: Mon, 19 Aug 2019 19:45:11 +0800
Subject: [PATCH] =?UTF-8?q?tqdm=E9=85=8D=E7=BD=AE?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 oneforall/aiobrute.py       |  6 ++++--
 oneforall/common/resolve.py |  6 +++++-
 oneforall/common/utils.py   |  6 +++---
 oneforall/takeover.py       | 10 +++++++---
 4 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/oneforall/aiobrute.py b/oneforall/aiobrute.py
index 59ef305..a7580a5 100644
--- a/oneforall/aiobrute.py
+++ b/oneforall/aiobrute.py
@@ -242,8 +242,10 @@ class AIOBrute(Module):
                 = detect_wildcard(domain)
         tasks = self.gen_tasks(domain)
         logger.log('INFOR', f'正在爆破{domain}的域名')
-        for task in tqdm.tqdm(tasks, desc='Progress',
-                              smoothing=1.0, ncols=True):
+        for task in tqdm.tqdm(tasks,
+                              desc='Progress',
+                              smoothing=1.0,
+                              ncols=True):
             async with aiomp.Pool(processes=self.process,
                                   initializer=init_worker,
                                   childconcurrency=self.coroutine) as pool:
diff --git a/oneforall/common/resolve.py b/oneforall/common/resolve.py
index 4589645..46515c9 100644
--- a/oneforall/common/resolve.py
+++ b/oneforall/common/resolve.py
@@ -102,7 +102,11 @@ async def bulk_query_a(datas):
             tasks.append(task)
     if tasks:  # 任务列表里有任务不空时才进行解析
         futures = asyncio.as_completed(tasks)
-        for future in tqdm.tqdm(futures, total=len(tasks)):
+        for future in tqdm.tqdm(futures,
+                                total=len(tasks),
+                                desc='Progress',
+                                smoothing=1.0,
+                                ncols=True):
             try:
                 await future
             except:
diff --git a/oneforall/common/utils.py b/oneforall/common/utils.py
index 3892e2c..8706b0b 100644
--- a/oneforall/common/utils.py
+++ b/oneforall/common/utils.py
@@ -110,7 +110,7 @@ def get_domains(target):
     :return: 域名集合
     """
     domains = list()
-    logger.log('INFOR', f'正在获取域名')
+    logger.log('DEBUG', f'正在获取域名')
     if isinstance(target, (set, tuple)):
         domains = list(target)
     elif isinstance(target, list):
@@ -184,11 +184,11 @@ def save_data(fpath, data):
     try:
         with open(fpath, 'w', encoding="utf-8", newline='') as file:
             file.write(data)
-            logger.log('INFOR', fpath)
+            logger.log('ALERT', fpath)
     except TypeError:
         with open(fpath, 'wb') as file:
             file.write(data)
-            logger.log('INFOR', fpath)
+            logger.log('ALERT', fpath)
     except Exception as e:
         logger.log('ERROR', e)
 
diff --git a/oneforall/takeover.py b/oneforall/takeover.py
index f8a6830..fd6424b 100644
--- a/oneforall/takeover.py
+++ b/oneforall/takeover.py
@@ -115,13 +115,15 @@ class Takeover(Module):
 
     def progress(self):
         while not self.subdomainq.empty():
+            time.sleep(0.3)
             done = self.bar.total - self.subdomainq.qsize()
             self.bar.n = done
             self.bar.update()
+        self.bar.close()
 
     def run(self):
         start = time.time()
-        logger.log('INFOR', f'开始执行{self.source}模块')
+        logger.log('DEBUG', f'开始执行{self.source}模块')
         self.format = utils.check_format(self.format)
         self.dpath = utils.check_dpath(self.dpath)
         self.subdomains = utils.get_domains(self.target)
@@ -132,8 +134,10 @@ class Takeover(Module):
             # 创建待检查的子域队列
             for domain in self.subdomains:
                 self.subdomainq.put(domain)
-            # 设置进度大小
+            # 设置进度
             self.bar.total = self.subdomainq.qsize()
+            self.bar.desc = 'Progress'
+            self.bar.ncols = True
             # 进度线程
             threads = []
             thread = Thread(target=self.progress, daemon=True)
@@ -153,7 +157,7 @@ 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}模块')
+        logger.log('DEBUG', f'结束执行{self.source}模块')
 
 
 if __name__ == '__main__':