tqdm配置

This commit is contained in:
shmilylty
2019-08-19 19:45:11 +08:00
parent 9b99b9fd35
commit b60d86fdc3
4 changed files with 19 additions and 9 deletions
+4 -2
View File
@@ -242,8 +242,10 @@ class AIOBrute(Module):
= detect_wildcard(domain) = detect_wildcard(domain)
tasks = self.gen_tasks(domain) tasks = self.gen_tasks(domain)
logger.log('INFOR', f'正在爆破{domain}的域名') logger.log('INFOR', f'正在爆破{domain}的域名')
for task in tqdm.tqdm(tasks, desc='Progress', for task in tqdm.tqdm(tasks,
smoothing=1.0, ncols=True): desc='Progress',
smoothing=1.0,
ncols=True):
async with aiomp.Pool(processes=self.process, async with aiomp.Pool(processes=self.process,
initializer=init_worker, initializer=init_worker,
childconcurrency=self.coroutine) as pool: childconcurrency=self.coroutine) as pool:
+5 -1
View File
@@ -102,7 +102,11 @@ async def bulk_query_a(datas):
tasks.append(task) tasks.append(task)
if tasks: # 任务列表里有任务不空时才进行解析 if tasks: # 任务列表里有任务不空时才进行解析
futures = asyncio.as_completed(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: try:
await future await future
except: except:
+3 -3
View File
@@ -110,7 +110,7 @@ def get_domains(target):
:return: 域名集合 :return: 域名集合
""" """
domains = list() domains = list()
logger.log('INFOR', f'正在获取域名') logger.log('DEBUG', f'正在获取域名')
if isinstance(target, (set, tuple)): if isinstance(target, (set, tuple)):
domains = list(target) domains = list(target)
elif isinstance(target, list): elif isinstance(target, list):
@@ -184,11 +184,11 @@ def save_data(fpath, data):
try: try:
with open(fpath, 'w', encoding="utf-8", newline='') as file: with open(fpath, 'w', encoding="utf-8", newline='') as file:
file.write(data) file.write(data)
logger.log('INFOR', fpath) logger.log('ALERT', fpath)
except TypeError: except TypeError:
with open(fpath, 'wb') as file: with open(fpath, 'wb') as file:
file.write(data) file.write(data)
logger.log('INFOR', fpath) logger.log('ALERT', fpath)
except Exception as e: except Exception as e:
logger.log('ERROR', e) logger.log('ERROR', e)
+7 -3
View File
@@ -115,13 +115,15 @@ class Takeover(Module):
def progress(self): def progress(self):
while not self.subdomainq.empty(): while not self.subdomainq.empty():
time.sleep(0.3)
done = self.bar.total - self.subdomainq.qsize() done = self.bar.total - self.subdomainq.qsize()
self.bar.n = done self.bar.n = done
self.bar.update() self.bar.update()
self.bar.close()
def run(self): def run(self):
start = time.time() start = time.time()
logger.log('INFOR', f'开始执行{self.source}模块') logger.log('DEBUG', f'开始执行{self.source}模块')
self.format = utils.check_format(self.format) self.format = utils.check_format(self.format)
self.dpath = utils.check_dpath(self.dpath) self.dpath = utils.check_dpath(self.dpath)
self.subdomains = utils.get_domains(self.target) self.subdomains = utils.get_domains(self.target)
@@ -132,8 +134,10 @@ class Takeover(Module):
# 创建待检查的子域队列 # 创建待检查的子域队列
for domain in self.subdomains: for domain in self.subdomains:
self.subdomainq.put(domain) self.subdomainq.put(domain)
# 设置进度大小 # 设置进度
self.bar.total = self.subdomainq.qsize() self.bar.total = self.subdomainq.qsize()
self.bar.desc = 'Progress'
self.bar.ncols = True
# 进度线程 # 进度线程
threads = [] threads = []
thread = Thread(target=self.progress, daemon=True) thread = Thread(target=self.progress, daemon=True)
@@ -153,7 +157,7 @@ class Takeover(Module):
elapsed = round(end - start, 1) elapsed = round(end - start, 1)
logger.log('INFOR', f'{self.source}模块耗时{elapsed}' logger.log('INFOR', f'{self.source}模块耗时{elapsed}'
f'发现{len(self.results)}个子域存在接管风险') f'发现{len(self.results)}个子域存在接管风险')
logger.log('INFOR', f'结束执行{self.source}模块') logger.log('DEBUG', f'结束执行{self.source}模块')
if __name__ == '__main__': if __name__ == '__main__':