优化模块

oneforall/modules/check/cdx.py

@@ -24,16 +24,15 @@ class CheckCDX(Module):
                 f'https://{self.domain}/crossdomain.xml',
                 f'http://www.{self.domain}/crossdomain.xml',
                 f'https://www.{self.domain}/crossdomain.xml']
-        response = None
         for url in urls:
             self.header = self.get_header()
             self.proxy = self.get_proxy(self.source)
-            response = self.get(url)
+            response = self.get(url, check=False)
-            if response:
-                break
             if not response:
                 return
-        self.subdomains = utils.match_subdomain(self.domain, response.text)
+            if response and len(response.content):
+                self.subdomains = utils.match_subdomain(self.domain,
+                                                        response.text)
 
     def run(self):
         """

oneforall/modules/check/csp.py

@@ -12,38 +12,46 @@ class CheckCSP(Module):
     """
     检查内容安全策略收集子域名
     """
-
     def __init__(self, domain, header):
         Module.__init__(self)
         self.domain = self.register(domain)
         self.module = 'Check'
         self.source = 'ContentSecurityPolicy'
-        self.header = header
+        self.csp_header = header
+
+    def grab_header(self):
+        """
+        抓取请求头
+
+        :return: 请求头
+        """
+        csp_header = dict()
+        urls = [f'http://{self.domain}',
+                f'https://{self.domain}',
+                f'http://www.{self.domain}',
+                f'https://www.{self.domain}']
+        for url in urls:
+            self.header = self.get_header()
+            self.proxy = self.get_proxy(self.source)
+            response = self.get(url, check=False)
+            if response:
+                csp_header = response.headers
+                break
+        return csp_header
 
     def check(self):
         """
         正则匹配响应头中的内容安全策略字段以发现子域名
         """
-        if not self.header:
+        if not self.csp_header:
-            urls = [f'http://{self.domain}',
+            self.csp_header = self.grab_header()
-                    f'https://{self.domain}',
-                    f'http://www.{self.domain}',
-                    f'https://www.{self.domain}']
-            response = None
-            for url in urls:
-                self.header = self.get_header()
-                self.proxy = self.get_proxy(self.source)
-                response = self.get(url)
-                if response:
-                    break
-            if not response:
-                return
-            self.header = response.headers
         csp = self.header.get('Content-Security-Policy')
+        if not self.csp_header:
+            logger.log('DEBUG', f'获取{self.domain}域的请求头失败')
+            return
         if not csp:
             logger.log('DEBUG', f'{self.domain}域的响应头不存在内容安全策略字段')
             return
-        logger.log('DEBUG', f'{self.domain}域的响应头存在内容安全策略字段')
         self.subdomains = utils.match_subdomain(self.domain, csp)
 
     def run(self):

oneforall/modules/check/robots.py

@@ -1,18 +1,14 @@
 """
 检查内容安全策略收集子域名收集子域名
 """
-import requests
-
 from common.module import Module
 from common import utils
-from config import logger
 
 
 class CheckRobots(Module):
     """
     检查robots.txt收集子域名
     """
-
     def __init__(self, domain):
         Module.__init__(self)
         self.domain = self.register(domain)
@@ -27,16 +23,15 @@ class CheckRobots(Module):
                 f'https://{self.domain}/robots.txt',
                 f'http://www.{self.domain}/robots.txt',
                 f'https://www.{self.domain}/robots.txt']
-        response = None
         for url in urls:
             self.header = self.get_header()
             self.proxy = self.get_proxy(self.source)
-            response = self.get(url, allow_redirects=False)
+            response = self.get(url, check=False, allow_redirects=False)
-            if response:
-                break
             if not response:
                 return
-        self.subdomains = utils.match_subdomain(self.domain, response.text)
+            if response and len(response.content):
+                self.subdomains = utils.match_subdomain(self.domain,
+                                                        response.text)
 
     def run(self):
         """

oneforall/modules/check/sitemap.py

@@ -1,18 +1,14 @@
 """
 检查内容安全策略收集子域名收集子域名
 """
-import requests
-
 from common.module import Module
 from common import utils
-from config import logger
 
 
 class CheckRobots(Module):
     """
     检查sitemap收集子域名
     """
-
     def __init__(self, domain):
         Module.__init__(self)
         self.domain = self.register(domain)
@@ -39,17 +35,17 @@ class CheckRobots(Module):
                 f'https://{self.domain}/sitemap_index.xml',
                 f'http://www.{self.domain}/sitemap_index.xml',
                 f'https://www.{self.domain}/sitemap_index.xml']
-        response = None
         for url in urls:
             self.header = self.get_header()
             self.proxy = self.get_proxy(self.source)
             self.timeout = 10
-            response = self.get(url, allow_redirects=False)
+            response = self.get(url, check=False, allow_redirects=False)
-            if response:
-                break
             if not response:
                 return
-        self.subdomains = utils.match_subdomain(self.domain, response.text)
+            if response and len(response.content):
+                self.subdomains = utils.match_subdomain(self.domain,
+                                                        response.text)
+
 
     def run(self):
         """