添加环境检查

oneforall/brute.py

@@ -133,6 +133,7 @@ def gen_word_subdomains(expression, path):
 
 
 def query_domain_ns_a(ns_list):
+    logger.log('INFOR', f'正在查询权威DNS名称服务器{ns_list}的A记录')
     if not isinstance(ns_list, list):
         return list()
     ns_ip_list = []
@@ -152,6 +153,7 @@ def query_domain_ns_a(ns_list):
 
 
 def query_domain_ns(domain):
+    logger.log('INFOR', f'正在查询{domain}的NS记录')
     domain = utils.get_maindomain(domain)
     resolver = utils.dns_resolver()
     try:
@@ -482,6 +484,7 @@ class Brute(Module):
         self.enable_wildcard = False  # 当前域名是否使用泛解析
         self.wildcard_check = config.enable_wildcard_check
         self.wildcard_deal = config.enable_wildcard_deal
+        self.check_env = True
 
     def gen_brute_dict(self, domain):
         logger.log('INFOR', f'正在为{domain}生成爆破字典')
@@ -590,6 +593,8 @@ class Brute(Module):
 
     def run(self):
         logger.log('INFOR', f'开始执行{self.source}模块')
+        if self.check_env:
+            utils.check_env()
         self.domains = utils.get_domains(self.target)
         all_subdomains = list()
         for self.domain in self.domains:

oneforall/common/utils.py

@@ -9,15 +9,15 @@ from ipaddress import IPv4Address, ip_address
 from stat import S_IXUSR
 
 import psutil
-
-import config
+import tenacity
+import requests
 from pathlib import Path
 from records import Record, RecordCollection
 from dns.resolver import Resolver
 
+import config
 from common.domain import Domain
 from config import logger
-from common import resolve
 
 user_agents = [
     'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 '
@@ -485,18 +485,52 @@ def delete_file(*paths):
             logger.log('ERROR', e.args)
 
 
-def check_env():
+@tenacity.retry(stop=tenacity.stop_after_attempt(2))
+def check_net():
+    logger.log('INFOR', '正在检查网络环境')
+    url = 'https://www.example.com/'
+    logger.log('INFOR', f'访问地址 {url}')
+    try:
+        rsp = requests.get(url)
+    except Exception as e:
+        logger.log('ERROR', e.args)
+        logger.log('ALERT', '访问外网出错 重新检查中')
+        raise tenacity.TryAgain
+    if rsp.status_code != 200:
+        logger.log('ALERT', f'{rsp.request.method} {rsp.request.url} '
+                            f'{rsp.status_code} {rsp.reason}')
+        logger.log('ALERT', '不能正常访问外网 重新检查中')
+        raise tenacity.TryAgain
+    logger.log('INFOR', '能正常访问外网')
+
+
+def check_pre():
+    logger.log('INFOR', '正在检查依赖环境')
     system = platform.system()
     implementation = platform.python_implementation()
-    if implementation != 'CPython':
-        logger.log('ALERT', f'当前Python是基于{implementation}实现但OneForAll只在CPython下测试通过')
-    if system == 'Windows' and implementation == 'CPython':
     version = platform.python_version()
+    if implementation != 'CPython':
+        logger.log('ALERT', f'当前Python是基于{implementation}实现 '
+                            f'但OneForAll只在CPython下测试通过')
+    if version < '3.6':
+        logger.log('FATAL', 'OneForAll需要Python 3.6以上版本')
+        exit(1)
+    if system == 'Windows' and implementation == 'CPython':
         if version < '3.8':
             logger.log('FATAL', 'OneForAll在Windows系统运行时需要Python 3.8以上版本')
             exit(1)
 
 
+def check_env():
+    try:
+        check_net()
+    except Exception as e:
+        logger.log('DEBUG', e.args)
+        logger.log('FATAL', '不能正常访问外网')
+        exit(1)
+    check_pre()
+
+
 def get_maindomain(domain):
     return Domain(domain).registered()
 

oneforall/oneforall.py

@@ -177,6 +177,7 @@ class OneForAll(object):
         if self.brute:
             # 由于爆破会有大量dns解析请求 并发爆破可能会导致其他任务中的网络请求异常
             brute = Brute(self.domain, word=True, export=False)
+            brute.check_env = False
             brute.run()
 
         # 有关数据库处理