1.重构请求，边请求边存入数据库，解决内存占用过大问题。

2.format参数改为fmt。
3.添加信息富化模块。
4.优化iscdn模块代码。
5.移除数据库new字段。

modules/altdns.py

@@ -193,9 +193,7 @@ class Altdns(Module):
         self.end = time.time()
         self.elapse = round(self.end - self.start, 1)
         self.gen_result()
-        temp_data = resolve.run_resolve(self.domain, self.results)
-        fina_data = request.run_request(self.domain, temp_data, port)
-        data = data + fina_data
+        resolved_data = resolve.run_resolve(self.domain, self.results)
+        request.run_request(self.domain, resolved_data, port)
         logger.log('INFOR', f'Saving altdns results')
         utils.save_db(self.domain, data, 'altdns')
-        return data

modules/enrich.py

@@ -0,0 +1,72 @@
+from modules import iscdn
+from common import utils
+from common.database import Database
+from common.ipasn import IPAsnInfo
+from common.ipreg import IpRegData
+
+
+def get_ips(info):
+    ip = info.get('ip')
+    if not ip:
+        return None
+    ips = ip.split(',')
+    return ips
+
+
+def enrich_info(data):
+    ip_asn = IPAsnInfo()
+    ip_reg = IpRegData()
+    for index, info in enumerate(data):
+        ips = get_ips(info)
+        if not ips:
+            continue
+        public = list()
+        cidr = list()
+        asn = list()
+        org = list()
+        addr = list()
+        isp = list()
+        for ip in ips:
+            public.append(str(utils.ip_is_public(ip)))
+            asn_info = ip_asn.find(ip)
+            cidr.append(asn_info.get('cidr'))
+            asn.append(asn_info.get('asn'))
+            org.append(asn_info.get('org'))
+            ip_info = ip_reg.query(ip)
+            addr.append(ip_info.get('addr'))
+            isp.append(ip_info.get('isp'))
+        data[index]['public'] = ','.join(public)
+        data[index]['cidr'] = ','.join(cidr)
+        data[index]['asn'] = ','.join(asn)
+        data[index]['org'] = ','.join(org)
+        data[index]['addr'] = ','.join(addr)
+        data[index]['isp'] = ','.join(isp)
+        return data
+
+
+class Enrich(object):
+    def __init__(self, domain):
+        self.domain = domain
+
+    def get_data(self):
+        db = Database()
+        fields = ['url', 'cname', 'ip', 'public', 'cdn', 'header',
+                  'cidr', 'asn', 'org', 'addr', 'isp']
+        results = db.get_data_by_fields(self.domain, fields)
+        return results.as_dict()
+
+    def save_db(self, data):
+        db = Database()
+        for info in data:
+            url = info.pop('url')
+            info.pop('cname')
+            info.pop('ip')
+            info.pop('header')
+            db.update_data_by_url(self.domain, info, url)
+        db.close()
+
+    def run(self):
+        data = self.get_data()
+        data = enrich_info(data)
+        data = iscdn.do_check(data)
+        self.save_db(data)

modules/finder.py

@@ -7,6 +7,7 @@ from common import utils
 from common import resolve
 from common import request
 from common.module import Module
+from common.database import Database
 from config import settings
 from config.log import logger
 
@@ -25,16 +26,11 @@ class Finder(Module):
         new_subdomains = found_subdomains - existing_subdomains
         if not len(new_subdomains):
             self.finish()  # 未发现新的子域就直接返回
-            return data
         self.subdomains = new_subdomains
         self.finish()
         self.gen_result()
-        temp_data = resolve.run_resolve(domain, self.results)
-        fina_data = request.run_request(domain, temp_data, port)
-        data = data + fina_data
-        logger.log('INFOR', f'Saving finder results')
-        utils.save_db(domain, data, 'finder')
-        return data
+        resolved_data = resolve.run_resolve(domain, self.results)
+        request.run_request(domain, resolved_data, port)
 
 
 file_path = settings.data_storage_dir.joinpath('common_js_library.json')
@@ -177,22 +173,31 @@ def find_js_urls(domain, req_url, rsp_html):
     return js_urls
 
 
+def convert_to_dict(url_list):
+    url_dict = []
+    for url in url_list:
+        url_dict.append({'url': url})
+    return url_dict
+
 def find_subdomains(domain, data):
     subdomains = set()
     js_urls = set()
+    db = Database()
     for infos in data:
         jump_history = infos.get('history')
         req_url = infos.get('url')
         subdomains.update(find_in_history(domain, req_url, jump_history))
-        rsp_html = infos.get('response')
+        rsp_html = db.get_resp_by_url(domain, req_url)
         if not rsp_html:
             logger.log('DEBUG', f'an abnormal response occurred in the request {req_url}')
             continue
         subdomains.update(find_in_resp(domain, req_url, rsp_html))
         js_urls.update(find_js_urls(domain, req_url, rsp_html))
 
-    resp_data = request.bulk_request(js_urls)
-    for _, resp in resp_data:
+    req_data = convert_to_dict(js_urls)
+    resp_data = request.bulk_request(domain, req_data, ret=True)
+    while not resp_data.empty():
+        _, resp = resp_data.get()
         if not isinstance(resp, Response):
             continue
         text = utils.decode_resp_text(resp)

modules/iscdn.py

@@ -3,6 +3,7 @@ import ipaddress
 
 from config import settings
 from common import utils
+from common.database import Database
 from config.log import logger
 
 data_dir = settings.data_storage_dir
@@ -18,6 +19,8 @@ cdn_header_key = utils.load_json(data_dir.joinpath('cdn_header_keys.json'))
 
 
 def check_cname_keyword(cname):
+    if not cname:
+        return False
     names = cname.lower().split(',')
     for name in names:
         for keyword in cdn_cname_keyword.keys():
@@ -26,14 +29,22 @@ def check_cname_keyword(cname):
 
 
 def check_header_key(header):
-    header = set(map(lambda x: x.lower(), header.keys()))
-    for key in cdn_header_key:
-        if key in header:
-            return True
+    if isinstance(header, str):
+        header = json.loads(header)
+    if isinstance(header, dict):
+        header = set(map(lambda x: x.lower(), header.keys()))
+        for key in cdn_header_key:
+            if key in header:
+                return True
+    else:
+        return False
 
 
-def check_cdn_cidr(content):
-    ips = set(content.split(','))
+def check_cdn_cidr(ips):
+    if isinstance(ips, str):
+        ips = set(ips.split(','))
+    else:
+        return False
     for ip in ips:
         try:
             ip = ipaddress.ip_address(ip)
@@ -46,39 +57,30 @@ def check_cdn_cidr(content):
 
 
 def check_cdn_asn(asn):
-    if str(asn) in cdn_asn_list:
-        return True
+    if isinstance(asn, str):
+        if asn in cdn_asn_list:
+            return True
+    return False
 
 
-def check_cdn(data):
-    logger.log('DEBUG', f'Start cdn check module')
+def do_check(data):
+    logger.log('DEBUG', f'Checking cdn')
     for index, item in enumerate(data):
         cname = item.get('cname')
-        if cname:
-            if check_cname_keyword(cname):
-                data[index]['cdn'] = 1
-                continue
+        if check_cname_keyword(cname):
+            data[index]['cdn'] = 1
+            continue
         header = item.get('header')
-        if header:
-            header = json.loads(header)
-            if check_header_key(header):
-                data[index]['cdn'] = 1
-                continue
+        if check_header_key(header):
+            data[index]['cdn'] = 1
+            continue
         ip = item.get('ip')
-        if ip:
-            if check_cdn_cidr(ip):
-                data[index]['cdn'] = 1
-                continue
+        if check_cdn_cidr(ip):
+            data[index]['cdn'] = 1
+            continue
         asn = item.get('asn')
-        if asn:
-            asn = asn[2:]  # 去除AS
-            if check_cdn_asn(asn):
-                data[index]['cdn'] = 1
-                continue
+        if check_cdn_asn(asn):
+            data[index]['cdn'] = 1
+            continue
         data[index]['cdn'] = 0
     return data
-
-
-def save_db(name, data):
-    logger.log('DEBUG', f'Saving cdn check results')
-    utils.save_db(name, data, 'cdn')