Agent: Fix false negatives in HTTPFingerprinter

monkey/infection_monkey/network_scanning/http_fingerprinter.py

@@ -1,8 +1,8 @@
 import logging
 from contextlib import closing
-from typing import Dict, Iterable, Optional, Set, Tuple
+from typing import Dict, Iterable, Optional, Set, Tuple, Any
 
-from requests import head
+from requests import head, Response
 from requests.exceptions import ConnectionError, Timeout
 
 from infection_monkey.i_puppet import (
@@ -55,22 +55,27 @@ def _query_potential_http_server(host: str, port: int) -> Tuple[Optional[str], O
     https = f"https://{host}:{port}"
 
     for url, ssl in ((https, True), (http, False)):  # start with https and downgrade
-        server_header_contents = _get_server_from_headers(url)
+        server_header = _get_server_from_headers(url)
 
-        if server_header_contents is not None:
+        if server_header is not None:
-            return (server_header_contents, ssl)
+            return server_header, ssl
 
-    return (None, None)
+    return None, None
 
 
 def _get_server_from_headers(url: str) -> Optional[str]:
+    headers = _get_http_headers(url)
+    if headers:
+        return headers.get("Server", "")
+
+    return None
+
+
+def _get_http_headers(url: str) -> Optional[Dict[str, Any]]:
     try:
         logger.debug(f"Sending request for headers to {url}")
-        with closing(head(url, verify=False, timeout=1)) as req:  # noqa: DUO123
+        with closing(head(url, verify=False, timeout=1)) as response:  # noqa: DUO123
-            server = req.headers.get("Server")
+            return response.headers
-
-            logger.debug(f'Got server string "{server}" from {url}')
-            return server
     except Timeout:
         logger.debug(f"Timeout while requesting headers from {url}")
     except ConnectionError:  # Someone doesn't like us

monkey/tests/unit_tests/infection_monkey/network_scanning/test_http_fingerprinter.py

@@ -7,25 +7,27 @@ from infection_monkey.network_scanning.http_fingerprinter import HTTPFingerprint
 
 OPTIONS = {"http_ports": [80, 443, 8080, 9200]}
 
-PYTHON_SERVER_HEADER = "SimpleHTTP/0.6 Python/3.6.9"
+PYTHON_SERVER_HEADER = {"Server": "SimpleHTTP/0.6 Python/3.6.9"}
-APACHE_SERVER_HEADER = "Apache/Server/Header"
+APACHE_SERVER_HEADER = {"Server": "Apache/Server/Header"}
+NO_SERVER_HEADER = {"Not_Server": "No Header for you"}
 
 SERVER_HEADERS = {
     "https://127.0.0.1:443": PYTHON_SERVER_HEADER,
     "http://127.0.0.1:8080": APACHE_SERVER_HEADER,
+    "http://127.0.0.1:1080": NO_SERVER_HEADER,
 }
 
 
 @pytest.fixture
-def mock_get_server_from_headers():
+def mock_get_http_headers():
-    return MagicMock(side_effect=lambda port: SERVER_HEADERS.get(port, None))
+    return MagicMock(side_effect=lambda url: SERVER_HEADERS.get(url, None))
 
 
 @pytest.fixture(autouse=True)
-def patch_get_server_from_headers(monkeypatch, mock_get_server_from_headers):
+def patch_get_http_headers(monkeypatch, mock_get_http_headers):
     monkeypatch.setattr(
         "infection_monkey.network_scanning.http_fingerprinter._get_server_from_headers",
-        mock_get_server_from_headers,
+        mock_get_http_headers,
     )