river_tao/OneForAll-mirror
1
0
Fork 0
mirror of https://github.com/shmilylty/OneForAll.git synced 2026-08-26 12:57:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

移除aiobrute向aiodns_query_a()函数传semaphore参数

Browse Source
This commit is contained in:
shmilylty
2019-08-05 23:38:53 +08:00
parent e3f9d5e149
commit 31666fdd75
2 changed files with 14 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
oneforall/aiobrute.py
+9 -11
View File
@@ -14,7 +14,6 @@ import queue
import secrets import secrets
import signal import signal
import time import time
from functools import partial
import aiomultiprocess import aiomultiprocess
import exrex import exrex
@@ -210,7 +209,6 @@ class AIOBrute(Module):
self.enable_wildcard, self.wildcard_ips, self.wildcard_ttl = detect_wildcard(domain) self.enable_wildcard, self.wildcard_ips, self.wildcard_ttl = detect_wildcard(domain)
tasks = self.gen_tasks(domain) tasks = self.gen_tasks(domain)
logger.log('INFOR', f'正在爆破{domain}的域名') logger.log('INFOR', f'正在爆破{domain}的域名')
sem = asyncio.Semaphore(utils.get_semaphore())
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 aiomultiprocess.Pool(processes=self.processes, initializer=init_worker, async with aiomultiprocess.Pool(processes=self.processes, initializer=init_worker,
childconcurrency=self.coroutine) as pool: childconcurrency=self.coroutine) as pool:
@@ -223,11 +221,10 @@ class AIOBrute(Module):
self.gen_result() self.gen_result()
rx_queue.put(self.results) rx_queue.put(self.results)
return return
else: self.deal_results(results)
self.deal_results(results) self.save_json()
self.save_json() self.gen_result()
self.gen_result() rx_queue.put(self.results)
rx_queue.put(self.results)
def run(self, rx_queue=None): def run(self, rx_queue=None):
self.domains = utils.get_domains(self.target) self.domains = utils.get_domains(self.target)
@@ -257,7 +254,8 @@ class AIOBrute(Module):
loop.run_until_complete(self.main(subdomain, rx_queue)) loop.run_until_complete(self.main(subdomain, rx_queue))
while not rx_queue.empty(): # 队列不空就一直取数据存数据库 while not rx_queue.empty(): # 队列不空就一直取数据存数据库
database.save_db(db_conn, table_name, rx_queue.get()) # 将结果存入数据库中 source, results = rx_queue.get()
database.save_db(db_conn, table_name, results, source) # 将结果存入数据库中
database.copy_table(db_conn, table_name) database.copy_table(db_conn, table_name)
database.deduplicate_subdomain(db_conn, table_name) database.deduplicate_subdomain(db_conn, table_name)
database.remove_invalid(db_conn, table_name) database.remove_invalid(db_conn, table_name)
@@ -281,6 +279,6 @@ def do(domain, result): # 统一入口名字 方便多线程调用
if __name__ == '__main__': if __name__ == '__main__':
fire.Fire(AIOBrute) # fire.Fire(AIOBrute)
# result_queue = queue.Queue() result_queue = queue.Queue()
# do('example.com', result_queue) do('example.com', result_queue)
oneforall/common/resolve.py
+5 -2
View File
@@ -48,11 +48,14 @@ async def aiodns_query_a(hostname, semaphore=None):
:return: 主机名或查询结果或查询异常 :return: 主机名或查询结果或查询异常
""" """
if semaphore is None: if semaphore is None:
semaphore = utils.get_semaphore()
async with semaphore:
resolver = aiodns_resolver() resolver = aiodns_resolver()
answers = await resolver.query(hostname, 'A') answers = await resolver.query(hostname, 'A')
return hostname, answers return hostname, answers
else:
async with semaphore:
resolver = aiodns_resolver()
answers = await resolver.query(hostname, 'A')
return hostname, answers
def resolve_callback(future, index, datas): def resolve_callback(future, index, datas):