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

添加urls批量请求函数

Browse Source
This commit is contained in:
Jing Ling
2020-07-09 17:56:16 +08:00
parent d31ade7a40
commit 0093cfdd90
2 changed files with 63 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/request.py
+52 -22
View File
@@ -22,7 +22,7 @@ def get_limit_conn():
def get_ports(port):
logger.log('DEBUG', f'Getting port range')
logger.log('DEBUG', 'Getting port range')
ports = set()
if isinstance(port, (set, list, tuple)):
ports = port
@@ -30,17 +30,17 @@ def get_ports(port):
if 0 <= port <= 65535:
ports = {port}
elif port in {'default', 'small', 'large'}:
logger.log('DEBUG', f'{port} port range')
logger.log('DEBUG', '{port} port range')
ports = setting.ports.get(port)
if not ports: # 意外情况
logger.log('ERROR', f'The specified request port range is incorrect')
logger.log('ERROR', 'The specified request port range is incorrect')
ports = {80}
logger.log('INFOR', f'Port range:{ports}')
logger.log('INFOR', 'Port range:{ports}')
return set(ports)
def gen_req_data(data, ports):
logger.log('INFOR', f'Generating request urls')
logger.log('INFOR', 'Generating request urls')
new_data = []
for data in data:
resolve = data.get('resolve')
@@ -70,15 +70,15 @@ def gen_req_data(data, ports):
return new_data
async def fetch(session, url):
async def fetch(session, method, url):
"""
请求
:param session: session对象
:param method: 请求方法
:param str url: url地址
:return: 响应对象和响应文本
"""
method = setting.request_method.upper()
timeout = aiohttp.ClientTimeout(total=None,
connect=None,
sock_read=setting.sockread_timeout,
@@ -104,12 +104,10 @@ async def fetch(session, url):
except UnicodeError:
try:
# 再尝试用gb18030解码
text = await resp.text(encoding='gb18030',
errors='strict')
text = await resp.text(encoding='gb18030', errors='strict')
except UnicodeError:
# 最后尝试自动解码
text = await resp.text(encoding=None,
errors='ignore')
text = await resp.text(encoding=None, errors='ignore')
return resp, text
except Exception as e:
return e
@@ -150,7 +148,7 @@ def get_title(markup):
text = soup.text
if len(text) <= 200:
return text
return repr(text)
return 'None'
@@ -190,27 +188,45 @@ def get_connector():
limit_per_host=setting.limit_per_host)
async def async_request(urls):
results = list()
connector = get_connector()
header = utils.get_random_header()
async with ClientSession(connector=connector, headers=header) as session:
tasks = []
for i, url in enumerate(urls):
task = asyncio.ensure_future(fetch(session, 'GET', url))
tasks.append(task)
if tasks:
futures = asyncio.as_completed(tasks)
for future in tqdm.tqdm(futures,
total=len(tasks),
desc='Request Progress',
ncols=80):
result = await future
results.append(result)
return results
async def bulk_request(data, port):
ports = get_ports(port)
no_req_data = utils.get_filtered_data(data)
to_req_data = gen_req_data(data, ports)
method = setting.request_method
method = setting.request_method.upper()
logger.log('INFOR', f'Use {method} method to request')
logger.log('INFOR', f'Async subdomains request in progress')
logger.log('INFOR', 'Async subdomains request in progress')
connector = get_connector()
header = utils.get_random_header()
async with ClientSession(connector=connector, headers=header) as session:
tasks = []
for i, data in enumerate(to_req_data):
url = data.get('url')
task = asyncio.ensure_future(fetch(session, url))
task = asyncio.ensure_future(fetch(session, method, url))
task.add_done_callback(functools.partial(request_callback,
index=i,
datas=to_req_data))
tasks.append(task)
# 任务列表里有任务不空时才进行解析
if tasks:
# 等待所有task完成 错误聚合到结果列表里
futures = asyncio.as_completed(tasks)
for future in tqdm.tqdm(futures,
total=len(tasks),
@@ -229,6 +245,13 @@ def set_loop_policy():
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
def set_loop():
set_loop_policy()
loop = asyncio.get_event_loop()
asyncio.set_event_loop(loop)
return loop
def run_request(domain, data, port):
"""
HTTP request entrance
@@ -238,20 +261,27 @@ def run_request(domain, data, port):
:param str port: range of ports to be requested
:return list: result
"""
logger.log('INFOR', f'Start subdomain request module')
set_loop_policy()
loop = asyncio.get_event_loop()
asyncio.set_event_loop(loop)
logger.log('INFOR', 'Start subdomain request module')
loop = set_loop()
data = utils.set_id_none(data)
request_coroutine = bulk_request(data, port)
data = loop.run_until_complete(request_coroutine)
# 在关闭事件循环前加入一小段延迟让底层连接得到关闭的缓冲时间
loop.run_until_complete(asyncio.sleep(0.25))
count = utils.count_alive(data)
logger.log('INFOR', f'Request module found {domain} have {count} alive subdomains')
return data
def urls_request(urls):
logger.log('INFOR', 'Start urls request module')
loop = set_loop()
request_coroutine = async_request(urls)
data = loop.run_until_complete(request_coroutine)
loop.run_until_complete(asyncio.sleep(0.25))
logger.log('INFOR', 'End urls request module')
return data
def save_data(name, data):
"""
Save request results to database
common/utils.py
+11 -7
View File
@@ -55,7 +55,6 @@ def gen_fake_header():
'Accept-Encoding': 'gzip, deflate',
'Accept-Language': 'en-US,en;q=0.9,zh-CN;q=0.8,zh;q=0.7',
'Cache-Control': 'max-age=0',
'Connection': 'close',
'DNT': '1',
'Referer': 'https://www.google.com/',
'Upgrade-Insecure-Requests': '1',
@@ -540,13 +539,15 @@ def check_pre():
exit(1)
if system == 'Windows' and implementation == 'CPython':
if version < '3.8':
logger.log('FATAL', 'OneForAll requires Python 3.8 or higher when running on Windows')
logger.log('FATAL', 'OneForAll requires Python 3.8 '
'or higher when running on Windows')
exit(1)
if system in {"Linux", "Darwin"}:
try:
import uvloop
except ImportError:
logger.log('ALERT', f'Please install the uvloop library manually to accelerate subdomain requests')
logger.log('ALERT', f'Please install the uvloop library manually '
f'to accelerate subdomain requests')
def check_env():
@@ -578,8 +579,10 @@ def check_version(local):
return
if latest > local:
change = json.get("body")
logger.log('ALERT', f'The current version is {local} but the latest version is {latest}')
logger.log('ALERT', f'The {latest} version mainly has the following changes\n{change}')
logger.log('ALERT', f'The current version is {local} '
f'but the latest version is {latest}')
logger.log('ALERT', f'The {latest} version mainly has the following changes')
logger.log('ALERT', change)
else:
logger.log('INFOR', f'The current version {local} is already the latest version')
@@ -591,7 +594,7 @@ def get_maindomain(domain):
def call_massdns(massdns_path, dict_path, ns_path, output_path, log_path,
query_type='A', process_num=1, concurrent_num=10000,
quiet_mode=False):
logger.log('DEBUG', f'Start running massdns')
logger.log('DEBUG', 'Start running massdns')
quiet = ''
if quiet_mode:
quiet = '--quiet'
@@ -626,7 +629,8 @@ def get_massdns_path(massdns_dir):
path.chmod(S_IXUSR)
if not path.exists():
logger.log('FATAL', 'There is no massdns for this platform or architecture')
logger.log('INFOR', 'Please try to compile massdns yourself and specify the path in the configuration')
logger.log('INFOR', 'Please try to compile massdns yourself '
'and specify the path in the configuration')
exit(0)
return path