优化

common/database.py

@@ -97,22 +97,21 @@ class Database(object):
         :param list results: results list
         :param str module_name: mo
         """
-        logger.log('TRACE',
-                   f'Saving the subdomain results of {table_name} found by module {module_name} into database')
+        logger.log('TRACE', f'Saving the subdomain results of {table_name} '
+                            f'found by module {module_name} into database')
         table_name = table_name.replace('.', '_')
         if results:
             try:
                 self.conn.bulk_query(
-                    f'insert into "{table_name}" ('
-                    f'id, type, alive, resolve, request, new, url, subdomain,'
-                    f'port, level, cname, content, public, status, reason,'
-                    f'title, banner, header, response, times, ttl, cidr, asn,'
-                    f'ip2region, ip2location, resolver, module, source, elapse, find, brute, valid) '
-                    f'values (:id, :type, :alive, :resolve, :request, :new,'
-                    f':url, :subdomain, :port, :level, :cname, :content,'
-                    f':public, :status, :reason, :title, :banner, :header,'
-                    f':response, :times, :ttl, :cidr, :asn, :ip2region, :ip2location, :resolver,'
-                    f':module, :source, :elapse, :find, :brute, :valid)', results)
+                    f'insert into "{table_name}" (id, type, alive, resolve, request, new,'
+                    f'url, subdomain, port, level, cname, content, public, status, reason,'
+                    f'title, banner, header, response, times, ttl, cidr, asn, ip2region,'
+                    f'ip2location, resolver, module, source, elapse, find, brute, valid) '
+                    f'values (:id, :type, :alive, :resolve, :request, :new, :url, '
+                    f':subdomain, :port, :level, :cname, :content, :public, :status,'
+                    f':reason, :title, :banner, :header, :response, :times, :ttl, :cidr,'
+                    f':asn, :ip2region, :ip2location, :resolver, :module, :source,'
+                    f':elapse, :find, :brute, :valid)', results)
             except Exception as e:
                 logger.log('ERROR', e)
 

common/lookup.py

@@ -20,7 +20,7 @@ class Lookup(Module):
             return None
         for item in answer:
             record = item.to_text()
-            subdomains = self.match_subdomains(self.domain, record)
+            subdomains = self.match_subdomains(record)
             self.subdomains = self.subdomains.union(subdomains)
             self.gen_record(subdomains, record)
         return self.subdomains

common/module.py

@@ -53,7 +53,8 @@ class Module(object):
         """
         begin log
         """
-        logger.log('DEBUG', f'Start {self.source} module to collect subdomains of {self.domain}')
+        logger.log('DEBUG', f'Start {self.source} module to '
+                            f'collect subdomains of {self.domain}')
 
     def finish(self):
         """
@@ -61,7 +62,8 @@ class Module(object):
         """
         self.end = time.time()
         self.elapse = round(self.end - self.start, 1)
-        logger.log('DEBUG', f'Finished {self.source} module to collect {self.domain}\'s subdomains')
+        logger.log('DEBUG', f'Finished {self.source} module to '
+                            f'collect {self.domain}\'s subdomains')
         logger.log('INFOR', f'The {self.source} module took {self.elapse} seconds '
                             f'found {len(self.subdomains)} subdomains')
         logger.log('DEBUG', f'{self.source} module found subdomains of {self.domain}\n'
@@ -209,35 +211,48 @@ class Module(object):
             logger.log('TRACE', f'{module} module does not use proxy')
             return self.proxy
 
-    @staticmethod
-    def match_subdomains(domain, text, distinct=True):
+    def match_subdomains(self, html, distinct=True, fuzzy=True):
         """
         Use regexp to match subdomains
 
-        :param  str domain: domain
-        :param  str text: text
+        :param  str html: response html text
         :param  bool distinct: deduplicate results or not (default True)
+        :param  bool fuzzy: fuzzy match subdomain or not (default True)
         :return set/list: result set or list
         """
         logger.log('TRACE', f'Use regexp to match subdomains in the response body')
-        regexp = r'(?:[a-z0-9](?:[a-z0-9\-]{0,61}[a-z0-9])?\.){0,}' \
-                 + domain.replace('.', r'\.')
-        result = re.findall(regexp, text, re.I)
+        if fuzzy:
+            regexp = r'(?:[a-z0-9](?:[a-z0-9\-]{0,61}[a-z0-9])?\.){0,}' \
+                     + self.domain.replace('.', r'\.')
+            result = re.findall(regexp, html, re.I)
+            if not result:
+                return set()
+            deal = map(lambda s: s.lower(), result)
+            if distinct:
+                return set(deal)
+            else:
+                return list(deal)
+        else:
+            regexp = r'(?:\>|\"|\'|\=|\,)(?:http\:\/\/|https\:\/\/)?' \
+                     r'(?:[a-z0-9](?:[a-z0-9\-]{0,61}[a-z0-9])?\.){0,}' \
+                     + self.domain.replace('.', r'\.')
+            result = re.findall(regexp, html, re.I)
         if not result:
             return set()
-        deal = map(lambda s: s.lower(), result)
+        regexp = r'(?:http://|https://)'
+        deal = map(lambda s: re.sub(regexp, '', s[1:].lower()), result)
         if distinct:
             return set(deal)
         else:
             return list(deal)
 
     @staticmethod
-    def register(domain):
+    def get_maindomain(domain):
         """
-        Get registered domain
+        Get main domain
 
         :param str domain: domain
-        :return: registered domain
+        :return: main domain
         """
         return Domain(domain).registered()
 
@@ -249,7 +264,8 @@ class Module(object):
         """
         if not setting.save_module_result:
             return False
-        logger.log('TRACE', f'Save the subdomain results found by {self.source} module as a json file')
+        logger.log('TRACE', f'Save the subdomain results found by '
+                            f'{self.source} module as a json file')
         path = setting.result_save_dir.joinpath(self.domain, self.module)
         path.mkdir(parents=True, exist_ok=True)
         name = self.source + '.json'

common/request.py

@@ -281,7 +281,7 @@ def urls_request(urls):
     return data
 
 
-def save_data(name, data):
+def save_db(name, data):
     """
     Save request results to database
 

common/resolve.py

@@ -48,7 +48,7 @@ def update_data(data, records):
     return data
 
 
-def save_data(name, data):
+def save_db(name, data):
     """
     保存解析结果到数据库
 

common/search.py

@@ -36,13 +36,12 @@ class Search(Module):
                                                    subdomains_temp[i:i + 2]))))
         return statements_list
 
-    def match_location(self, domain, url):
+    def match_location(self, url):
         """
         匹配跳转之后的url
         针对部分搜索引擎(如百度搜索)搜索展示url时有显示不全的情况
         此函数会向每条结果的链接发送head请求获取响应头的location值并做子域匹配
 
-        :param str domain: 域名
         :param str url: 展示结果的url链接
         :return: 匹配的子域
         :rtype set
@@ -53,28 +52,4 @@ class Search(Module):
         location = resp.headers.get('location')
         if not location:
             return set()
-        return set(self.match_subdomains(domain, location))
-
-    @staticmethod
-    def match_subdomains(domain, html, distinct=True):
-        """
-        Use regexp to match subdomains
-
-        :param  str domain: domain
-        :param  str html: response html text
-        :param  bool distinct: deduplicate results or not (default True)
-        :return set/list: result set or list
-        """
-        logger.log('TRACE', f'Use regexp to match subdomains in the response body')
-        regexp = r'(?:\>|\"|\'|\=|\,)(?:http\:\/\/|https\:\/\/)?' \
-                 r'(?:[a-z0-9](?:[a-z0-9\-]{0,61}[a-z0-9])?\.){0,}' \
-                 + domain.replace('.', r'\.')
-        result = re.findall(regexp, html, re.I)
-        if not result:
-            return set()
-        regexp = r'(?:http://|https://)'
-        deal = map(lambda s: re.sub(regexp, '', s[1:].lower()), result)
-        if distinct:
-            return set(deal)
-        else:
-            return list(deal)
+        return set(self.match_subdomains(location))

common/utils.py

@@ -237,6 +237,21 @@ def save_data(path, data):
         return False
 
 
+def remove_data(path):
+    """
+    删除保存数据的文件
+
+    :param path: 路径
+    :return: 删除成功与否
+    """
+    try:
+        path.unlink()
+    except Exception as e:
+        logger.log('ERROR', e.args)
+        return False
+    return True
+
+
 def check_response(method, resp):
     """
     检查响应 输出非正常响应返回json的信息