diff --git a/oneforall/aiobrute.py b/oneforall/aiobrute.py index 3262143..3857e37 100644 --- a/oneforall/aiobrute.py +++ b/oneforall/aiobrute.py @@ -28,10 +28,6 @@ from common.database import Database from config import logger -def init_worker(): - signal.signal(signal.SIGINT, signal.SIG_IGN) - - def detect_wildcard(domain): """ 探测域名是否使用泛解析 @@ -131,25 +127,6 @@ def gen_brute_domains(domain, path): return domains -def progress(pr_queue, total): - bar = tqdm.tqdm() - bar.total = total - bar.desc = 'Progress' - bar.ncols = 60 - while True: - done = pr_queue.qsize() - bar.n = done - bar.update() - if done == total: - return - - -async def aiodns_query_a(pr_queue, hostname): - results = await resolve.aiodns_query_a(hostname) - pr_queue.put(1) - return results - - class AIOBrute(Module): """ OneForAll多进程多协程异步子域爆破模块 @@ -250,7 +227,7 @@ class AIOBrute(Module): if self.enable_wildcard and self.wildcard_deal: # 通过对比查询的子域和响应的子域来判断真实子域 # 去掉解析到CDN的情况 - if 'cdn' or 'waf' in name: + if 'cdn' in name or 'waf' in name: continue if not name.endswith(self.domain): continue @@ -273,27 +250,11 @@ class AIOBrute(Module): logger.log('INFOR', f'正在爆破{domain}的域名') # for task in tqdm.tqdm(tasks, total=len(tasks), # desc='Progress'): - m = Manager() - pr_queue = m.Queue() - loop = asyncio.get_event_loop() - loop.run_in_executor(None, progress, pr_queue, len(tasks)) - wrapped_query = functools.partial(aiodns_query_a, pr_queue) - async with aiomp.Pool(processes=self.process, - initializer=init_worker, - childconcurrency=self.coroutine) as pool: - try: - results = await pool.map(wrapped_query, tasks) - except KeyboardInterrupt: - logger.log('ALERT', '爆破终止正在退出') - pool.terminate() # 关闭pool,结束工作进程,不在处理未完成的任务。 - self.save_json() - self.gen_result() - rx_queue.put(self.results) - return - self.deal_results(results) - self.save_json() - self.gen_result() - rx_queue.put(self.results) + results = await resolve.aio_resolve(tasks, self.process, self.coroutine) + self.deal_results(results) + self.save_json() + self.gen_result() + rx_queue.put(self.results) def run(self, rx_queue=None): self.domains = utils.get_domains(self.target) diff --git a/oneforall/common/resolve.py b/oneforall/common/resolve.py index 3a429f1..e9077e8 100644 --- a/oneforall/common/resolve.py +++ b/oneforall/common/resolve.py @@ -1,12 +1,14 @@ +import signal +import socket import asyncio import functools -import socket +from multiprocessing import Manager import tqdm +import aiomultiprocess as aiomp from dns.resolver import Resolver import config -from common import utils from config import logger @@ -21,22 +23,6 @@ def dns_resolver(): return resolver -async def dns_query_a(hostname): - """ - 查询A记录 - - :param str hostname: 主机名 - :return: 查询结果 - """ - resolver = dns_resolver() - try: - answer = resolver.query(hostname, 'A') - except Exception as e: - logger.log('TRACE', e.args) - answer = e - return hostname, answer - - async def aiodns_query_a(hostname): """ 异步查询A记录 @@ -57,56 +43,129 @@ async def aiodns_query_a(hostname): return hostname, answer -def resolve_callback(future, index, datas): +def convert_results(result_list): """ - 解析结果回调处理 + 将结果列表类型转换为结果字典类型 - :param future: future对象 - :param index: 下标 - :param datas: 结果集 + :param result_list: 待转换的结果列表 + :return: 转换后的结果字典 """ - hostname, answer = future.result() - if isinstance(answer, BaseException): - logger.log('TRACE', answer.args) - name = utils.get_classname(answer) - datas[index]['reason'] = name + ' ' + str(answer) - datas[index]['valid'] = 0 - elif isinstance(answer, tuple): - ips = answer[2] - datas[index]['ips'] = str(ips)[1:-1] + result_dict = {} + for result in result_list: + hostname, answer = result + value_dict = {'ips': None, 'reason': None, 'valid': None} + if isinstance(answer, tuple): + value_dict['ips'] = str(answer[2])[1:-1] + result_dict[hostname] = value_dict + elif isinstance(answer, Exception): + value_dict['reason'] = str(answer.args) + value_dict['valid'] = 0 + result_dict[hostname] = value_dict + else: + value_dict['valid'] = 0 + result_dict[hostname] = value_dict + return result_dict -async def bulk_query_a(datas): +def filter_subdomain(data_list): """ - 批量查询A记录 + 过滤出无IPS值的子域到新的子域列表 - :param datas: 待查的数据集 - :return: 查询过得到的结果集 + :param list data_list: 待过滤的数据列表 + :return: 符合条件的子域列表 """ - logger.log('INFOR', '正在异步查询子域的A记录') - tasks = [] - # semaphore = asyncio.Semaphore(config.limit_resolve_conn) - for i, data in enumerate(datas): + subdomains = [] + for data in data_list: if not data.get('ips'): subdomain = data.get('subdomain') - task = asyncio.ensure_future(aiodns_query_a(subdomain)) - wrapped_callback = functools.partial(resolve_callback, - index=i, - datas=datas) - task.add_done_callback(wrapped_callback) # 回调 - tasks.append(task) - if tasks: # 任务列表里有任务不空时才进行解析 - futures = asyncio.as_completed(tasks) - for future in tqdm.tqdm(futures, - total=len(tasks), - desc='Progress', - ncols=60): - await future - # await asyncio.wait(tasks) # 等待所有task完成 + subdomains.append(subdomain) + return subdomains + + +def update_data(data_list, results_dict): + """ + 更新解析结果 + + :param list data_list: 待更新的数据列表 + :param dict results_dict: 解析结果字典 + :return: 更新后的数据列表 + """ + for index, data in enumerate(data_list): + if not data.get('ips'): + subdomain = data.get('subdomain') + value_dict = results_dict.get(subdomain) + data.update(value_dict) + data_list[index] = data + return data_list + + +def init_worker(): + signal.signal(signal.SIGINT, signal.SIG_IGN) + + +def query_progress(pr_queue, total): + bar = tqdm.tqdm() + bar.total = total + bar.desc = 'Resolve Progress' + bar.ncols = 60 + bar.smoothing = 0 + while True: + done = pr_queue.qsize() + bar.n = done + bar.update() + if done == total: + break + bar.close() + + +async def aio_query(pr_queue, hostname): + """ + 异步查询主机名的A记录 + + :param pr_queue: 进度队列 + :param str hostname: 主机名 + :return: 查询结果 + """ + results = await aiodns_query_a(hostname) + pr_queue.put(1) + return results + + +async def aio_resolve(subdomain_list, process_num, coroutine_num): + """ + 异步解析子域A记录 + + :param list subdomain_list: 待解析的子域列表 + :param int process_num: 解析进程数 + :param int coroutine_num: 每个解析进程下的协程数 + :return: 解析结果 + """ + m = Manager() + pr_queue = m.Queue() + loop = asyncio.get_event_loop() + loop.run_in_executor(None, query_progress, pr_queue, len(subdomain_list)) + wrapped_query = functools.partial(aio_query, pr_queue) + async with aiomp.Pool(processes=process_num, + initializer=init_worker, + childconcurrency=coroutine_num) as pool: + results = await pool.map(wrapped_query, subdomain_list) + return results + + +async def bulk_resolve(data_list): + """ + 批量解析A记录并返回解析结果 + + :param list data_list: 待查的数据列表 + :return: 查询过得到的结果列表 + """ + logger.log('INFOR', '正在异步查询子域的A记录') + # semaphore = asyncio.Semaphore(config.limit_resolve_conn) + query_subdomains = filter_subdomain(data_list) + process_num = config.brute_process_num + coroutine_num = config.brute_coroutine_num + results = await aio_resolve(query_subdomains, process_num, coroutine_num) + results_dict = convert_results(results) + data_list = update_data(data_list, results_dict) logger.log('INFOR', '完成异步查询子域的A记录') - return datas - - -def run_bulk_query(datas): - new_datas = asyncio.run(bulk_query_a(datas)) - return new_datas + return data_list diff --git a/oneforall/common/utils.py b/oneforall/common/utils.py index bebae00..e3a1cdb 100644 --- a/oneforall/common/utils.py +++ b/oneforall/common/utils.py @@ -327,3 +327,7 @@ def get_classname(clsobj): def python_version(): return sys.version + + +def count_valid(data): + return len(list(filter(lambda item: item.get('valid') == 1, data))) diff --git a/oneforall/oneforall.py b/oneforall/oneforall.py index 96bfa32..b54a475 100644 --- a/oneforall/oneforall.py +++ b/oneforall/oneforall.py @@ -141,7 +141,7 @@ class OneForAll(object): asyncio.set_event_loop(loop) # 解析子域 - task = resolve.bulk_query_a(self.data) + task = resolve.bulk_resolve(self.data) self.data = loop.run_until_complete(task) # 保存解析结果 @@ -166,8 +166,8 @@ class OneForAll(object): self.datas.extend(self.data) # 在关闭事件循环前加入一小段延迟让底层连接得到关闭的缓冲时间 loop.run_until_complete(asyncio.sleep(0.25)) - valid_count = len(list(filter(lambda item: item.get('valid') == 1, self.data))) - logger.log('INFOR', f'经验证{self.domain}有效子域{valid_count}个') + count = utils.count_valid(self.data) + logger.log('INFOR', f'经验证{self.domain}有效子域{count}个') # 保存请求结果 db.clear_table(self.domain)