Merge pull request #76 from FCG-LLC/fix/aucote_http_headers

Fix Aucote HTTP Headers

tests/test_tools/test_aucote_http_headers/test_tasks.py

@@ -1,12 +1,9 @@
-from collections import KeysView
-from unittest import TestCase
 from unittest.mock import MagicMock, patch
 
 from tornado.concurrent import Future
-from tornado.httpclient import HTTPClient
+from tornado.httpclient import HTTPClient, HTTPError, HTTPResponse, HTTPRequest
 from tornado.testing import gen_test, AsyncTestCase
 
-from fixtures.exploits import Exploit
 from structs import Port, Scan
 from tools.aucote_http_headers.structs import HeaderDefinition, AucoteHttpHeaderResult
 from tools.aucote_http_headers.tasks import AucoteHttpHeadersTask
@@ -146,15 +143,28 @@ async def test_with_requests_os_error(self, http_client):
     @patch('tools.aucote_http_headers.tasks.cfg.get', MagicMock(return_value='test'))
     @gen_test
     async def test_server_reponse_403_logging(self, mock_log, http_client):
-        future = Future()
-        future.set_result(MagicMock(code=403))
-        http_client.instance().head.return_value = future
+        request = HTTPRequest(url='url')
+        response = HTTPResponse(code=403, request=request)
+        http_client.instance().head.side_effect = HTTPError(code=403, response=response)
         self.task.store_vulnerability = MagicMock()
 
         await self.task()
 
         self.assertTrue(mock_log.warning.called)
 
+    @patch('tools.aucote_http_headers.tasks.HTTPClient')
+    @patch('tools.aucote_http_headers.tasks.log')
+    @patch('tools.aucote_http_headers.tasks.cfg.get', MagicMock(return_value='test'))
+    @gen_test
+    async def test_server_reponse_599(self, mock_log, http_client):
+        http_client.instance().head.side_effect = HTTPError(code=403, response=None)
+        self.task.store_vulnerability = MagicMock()
+
+        result = await self.task()
+        expected = None
+
+        self.assertEqual(result, expected)
+
     @patch('tools.aucote_http_headers.tasks.HTTPClient')
     @patch('tools.aucote_http_headers.tasks.cfg.get', MagicMock(side_effect=(None, 'test')))
     @gen_test

tools/aucote_http_headers/tasks.py

@@ -5,6 +5,8 @@
 import time
 import logging as log
 
+from tornado.httpclient import HTTPError
+
 from aucote_cfg import cfg
 from structs import Vulnerability
 from tools.aucote_http_headers.structs import AucoteHttpHeaderResult as Result
@@ -38,17 +40,19 @@ async def __call__(self, *args, **kwargs):
 
         try:
             response = await HTTPClient.instance().head(url=self._port.url, headers=custom_headers, validate_cert=False)
-
-            if response.code != 200:
-                log.warning("Server replied with status code: %i", response.code)
-
+        except HTTPError as exception:
+            if exception.response is None:
+                return
+            response = exception.response
         except ConnectionError:
             log.exception("Cannot connect to %s", self._port.url)
             return
         except OSError as exception:
             log.warning("%s for %s", str(exception), self._port.url)
             return
 
+        if response.code != 200:
+            log.warning("Server replied with status code: %i", response.code)
         headers = response.headers
 
         results = []