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

重构解析模块

Browse Source
This commit is contained in:
Jing Ling
2020-02-17 03:39:37 +08:00
parent 246cad2511
commit 6c902544b5
4 changed files with 132 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files
oneforall/aiobrute.py
+2 -41
View File
@@ -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,23 +250,7 @@ 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
results = await resolve.aio_resolve(tasks, self.process, self.coroutine)
self.deal_results(results)
self.save_json()
self.gen_result()
oneforall/common/resolve.py
+119 -60
View File
@@ -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
oneforall/common/utils.py
+4
View File
@@ -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)))
oneforall/oneforall.py
+3 -3
View File
@@ -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)