From 19ceb7fe5cb3a629c0689494c8d26cf9d664cea3 Mon Sep 17 00:00:00 2001 From: Dominik Rosiek Date: Wed, 5 Jul 2017 16:54:12 +0200 Subject: [PATCH] Fix Aucote HTTP Headers --- .../test_aucote_http_headers/test_tasks.py | 24 +++++++++++++------ tools/aucote_http_headers/tasks.py | 12 ++++++---- 2 files changed, 25 insertions(+), 11 deletions(-) diff --git a/tests/test_tools/test_aucote_http_headers/test_tasks.py b/tests/test_tools/test_aucote_http_headers/test_tasks.py index f0fc1eb7..4f450efd 100644 --- a/tests/test_tools/test_aucote_http_headers/test_tasks.py +++ b/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 diff --git a/tools/aucote_http_headers/tasks.py b/tools/aucote_http_headers/tasks.py index 2c6eba5d..07ed9698 100644 --- a/tools/aucote_http_headers/tasks.py +++ b/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,10 +40,10 @@ 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 @@ -49,6 +51,8 @@ async def __call__(self, *args, **kwargs): 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 = []