From 3e19062670b24fa8e9708ac861459f354531e3bc Mon Sep 17 00:00:00 2001
From: shmilylty <admin@hackfun.org>
Date: Wed, 6 Nov 2019 01:09:30 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0HEAD=E8=AF=B7=E6=B1=82?=
 =?UTF-8?q?=E5=BC=82=E5=B8=B8=E5=A4=84=E7=90=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 oneforall/common/module.py | 35 +++++++++++++++++++++++++++++++----
 oneforall/common/search.py |  7 +++++--
 2 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/oneforall/common/module.py b/oneforall/common/module.py
index cec295f..2d2b997 100644
--- a/oneforall/common/module.py
+++ b/oneforall/common/module.py
@@ -15,7 +15,6 @@ from . import utils
 from .domain import Domain
 from common.database import Database
 
-
 lock = threading.Lock()
 
 
@@ -62,9 +61,37 @@ class Module(object):
         self.elapsed = round(self.end - self.start, 1)
         logger.log('DEBUG', f'结束执行{self.source}模块收集{self.domain}的子域')
         logger.log('INFOR', f'{self.source}模块耗时{self.elapsed}秒发现子域'
-                   f'{len(self.subdomains)}个')
+                            f'{len(self.subdomains)}个')
         logger.log('DEBUG', f'{self.source}模块发现{self.domain}的子域\n'
-                   f'{self.subdomains}')
+                            f'{self.subdomains}')
+
+    def head(self, url, params=None, check=True, **kwargs):
+        """
+        自定义head请求
+
+        :param str url: 请求地址
+        :param dict params: 请求参数
+        :param bool check: 检查响应
+        :param kwargs: 其他参数
+        :return: requests响应对象
+        """
+        try:
+            resp = requests.head(url,
+                                 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)
+            return None
+        if not check:
+            return resp
+        if utils.check_response('HEAD', resp):
+            return resp
+        return None
 
     def get(self, url, params=None, check=True, **kwargs):
         """
@@ -118,7 +145,7 @@ class Module(object):
             return None
         if not check:
             return resp
-        if utils.check_response('GET', resp):
+        if utils.check_response('POST', resp):
             return resp
         return None
 
diff --git a/oneforall/common/search.py b/oneforall/common/search.py
index d995463..422c505 100644
--- a/oneforall/common/search.py
+++ b/oneforall/common/search.py
@@ -47,7 +47,10 @@ class Search(Module):
         :return: 匹配的子域
         :rtype set
         """
-        resp = requests.head(url, headers=self.header, proxies=self.proxy,
-                             timeout=self.timeout, allow_redirects=False)
+        resp = self.head(url, allow_redirects=False)
+        if not resp:
+            return set()
         location = resp.headers.get('location')
+        if not location:
+            return set()
         return set(utils.match_subdomain(domain, location))