diff --git a/common/check.py b/common/check.py index fe0f6f6..14ef6ef 100644 --- a/common/check.py +++ b/common/check.py @@ -1,3 +1,5 @@ +import requests +from config.log import logger from common.module import Module @@ -12,15 +14,28 @@ class Check(Module): def to_check(self, filenames): urls = set() + urls_www = set() for filename in filenames: - urls.update((f'http://{self.domain}/{filename}', - f'https://{self.domain}/{filename}', - f'http://www.{self.domain}/{filename}', - f'https://www.{self.domain}/{filename}')) + urls.update(( + f'http://{self.domain}/{filename}', + f'https://{self.domain}/{filename}', + )) + urls_www.update(( + f'http://www.{self.domain}/{filename}', + f'https://www.{self.domain}/{filename}' + )) + self.check_loop(urls) + self.check_loop(urls_www) + + def check_loop(self, urls): for url in urls: self.header = self.get_header() self.proxy = self.get_proxy(self.source) - resp = self.get(url, check=False, ignore=True) + try: + resp = self.get(url, check=False, ignore=True, raise_error=True) + except requests.exceptions.ConnectTimeout: + logger.log('DEBUG', f'Connection to {url} timed out, so break check') + break self.subdomains = self.collect_subdomains(resp) if self.subdomains: break diff --git a/common/module.py b/common/module.py index 6a47a96..5f3aaba 100644 --- a/common/module.py +++ b/common/module.py @@ -79,13 +79,13 @@ class Module(object): session.trust_env = False try: resp = session.head(url, - params=params, - cookies=self.cookie, - headers=self.header, - proxies=self.proxy, - timeout=self.timeout, - verify=self.verify, - **kwargs) + params=params, + cookies=self.cookie, + headers=self.header, + proxies=self.proxy, + timeout=self.timeout, + verify=self.verify, + **kwargs) except Exception as e: logger.log('ERROR', e.args) return None @@ -95,7 +95,7 @@ class Module(object): return resp return None - def get(self, url, params=None, check=True, ignore=False, **kwargs): + def get(self, url, params=None, check=True, ignore=False,raise_error=False, **kwargs): """ Custom get request @@ -103,6 +103,7 @@ class Module(object): :param dict params: request parameters :param bool check: check response :param bool ignore: ignore error + :param bool raise_error: raise error or not :param kwargs: other params :return: response object """ @@ -121,6 +122,10 @@ class Module(object): verify=self.verify, **kwargs) except Exception as e: + if raise_error: + if isinstance(e, requests.exceptions.ConnectTimeout): + logger.log(level, e.args) + raise e logger.log(level, e.args) return None if not check: diff --git a/modules/check/csp.py b/modules/check/csp.py index c61a935..d8eb8d0 100644 --- a/modules/check/csp.py +++ b/modules/check/csp.py @@ -11,6 +11,7 @@ class CSP(Check): """ Collect subdomains from ContentSecurityPolicy """ + def __init__(self, domain, header): Check.__init__(self) self.domain = domain @@ -18,6 +19,7 @@ class CSP(Check): self.source = 'CSPCheck' self.csp_header = header + @property def grab_header(self): """ Get header @@ -26,13 +28,24 @@ class CSP(Check): """ csp_header = dict() urls = [f'http://{self.domain}', - f'https://{self.domain}', - f'http://www.{self.domain}', - f'https://www.{self.domain}'] + f'https://{self.domain}'] + urls_www = [f'http://www.{self.domain}', + f'https://www.{self.domain}'] + header = self.grab_loop(csp_header, urls) + if header: + return header + header = self.grab_loop(csp_header, urls_www) + return header + + def grab_loop(self, csp_header, urls): for url in urls: self.header = self.get_header() self.proxy = self.get_proxy(self.source) - response = self.get(url, check=False) + try: + response = self.get(url, check=False, ignore=True, raise_error=True) + except requests.exceptions.ConnectTimeout: + logger.log('DEBUG', f'Connection to {url} timed out, so break check') + break if response: return response.headers return csp_header @@ -42,7 +55,7 @@ class CSP(Check): 正则匹配响应头中的内容安全策略字段以发现子域名 """ if not self.csp_header: - self.csp_header = self.grab_header() + self.csp_header = self.grab_header csp = self.csp_header.get('Content-Security-Policy') if not self.csp_header: logger.log('DEBUG', f'Failed to get header of {self.domain} domain')