diff --git a/common/request.py b/common/request.py index c12c446..8b2e286 100644 --- a/common/request.py +++ b/common/request.py @@ -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 diff --git a/common/utils.py b/common/utils.py index de18a77..e2052a9 100644 --- a/common/utils.py +++ b/common/utils.py @@ -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