From 476364d2cacab6521cff6fe65167ecd2eabd1914 Mon Sep 17 00:00:00 2001 From: JrD <421273918@qq.com> Date: Tue, 8 Sep 2020 18:56:20 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8Dcheck=E6=A8=A1=E5=9D=97?= =?UTF-8?q?=E4=B8=8B=E5=9C=A8=E8=A7=A3=E6=9E=90=E5=87=BA=E5=A4=9A=E4=B8=AA?= =?UTF-8?q?IP=E4=B8=8D=E9=80=9A=E7=9A=84=E6=83=85=E5=86=B5=E4=B8=8B?= =?UTF-8?q?=E8=B6=85=E6=97=B6=E4=B8=A5=E9=87=8D=E7=9A=84=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- common/check.py | 25 ++++++++++++++++++++----- common/module.py | 21 +++++++++++++-------- modules/check/csp.py | 23 ++++++++++++++++++----- 3 files changed, 51 insertions(+), 18 deletions(-) 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')