Add simple web server mocking example

osint-dev-team · manmolecular · Jul 22, 2020 · Jul 19, 2020 · Jul 19, 2020 · Jul 19, 2020
commit efa1d0362e28a877b2be31e8cecf543461ebcaed
diff --git a/src/scripts/other/network_usage/__init__.py b/src/scripts/other/network_usage/__init__.py
@@ -0,0 +1,5 @@
+import sys
+from pathlib import Path
+
+__root_dir = Path(__file__).parents[4]
+sys.path.append(str(__root_dir))
diff --git a/src/scripts/other/network_usage/__main__.py b/src/scripts/other/network_usage/__main__.py
@@ -0,0 +1,12 @@
+#!/usr/bin/env python3
+
+from pprint import pprint
+from sys import argv
+
+from src.core.utils.module import run_module
+from .module import Runner
+
+result = run_module(
+    Runner, args=argv, arg_name="hostname", arg_default="https://www.example.com/"
+)
+pprint(result)
diff --git a/src/scripts/other/network_usage/module.py b/src/scripts/other/network_usage/module.py
@@ -0,0 +1,34 @@
+#!/usr/bin/env python3
+
+from src.core.base.recon import ReconRunner, PossibleKeys
+from src.core.utils.response import ScriptResponse
+from src.core.utils.validators import validate_kwargs
+from requests import get
+from requests.utils import CaseInsensitiveDict
+
+
+class Runner(ReconRunner):
+    def __init__(self, logger: str = __name__):
+        super(Runner, self).__init__(logger)
+
+    @staticmethod
+    def __get_host(hostname: str) -> CaseInsensitiveDict:
+        """
+        Get host headers
+        :param hostname: hostname to check
+        :return: dictionary with headers
+        """
+        return get(hostname).headers
+
+    @validate_kwargs(PossibleKeys.KEYS)
+    def run(self, *args, **kwargs) -> ScriptResponse.success or ScriptResponse.error:
+        """
+        Main runner function for the script
+        :param args: args from core runner
+        :param kwargs: kwargs from core runner
+        :return: ScriptResponse message
+        """
+        hostname = kwargs.get("hostname")
+        return ScriptResponse.success(
+            result=self.__get_host(hostname), message="Success"
+        )
diff --git a/src/scripts/other/network_usage/test_module.py b/src/scripts/other/network_usage/test_module.py
@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+
+from unittest import TestCase
+from .module import Runner
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from threading import Thread
+
+
+class DefaultValues:
+    """
+    Set default localhost values
+    """
+
+    HOST = "127.0.0.1"
+    PORT = 1337
+
+
+class TestClassHTTPRequestHandler(BaseHTTPRequestHandler):
+    """
+    Defines mocking server class
+    """
+
+    def do_GET(self):
+        """
+        Defines basic handlers
+        :return:
+        """
+        if self.path == "/":
+            self.mock_endpoint()
+
+    def mock_endpoint(self) -> None:
+        """
+        Mock '/' endpoint. Return HTML with custom headers.
+        :return: None
+        """
+        self.send_response(200)
+        self.send_header("Content-type", "text/html; charset=utf-8")
+        self.end_headers()
+        self.wfile.write(
+            b"<html><body><i>Mock page from Mocking Class</i></body></html>"
+        )
+
+
+class NetworkUsageTest(TestCase):
+    """
+    Defines basic tests for the network usage
+    """
+
+    server: Thread or None = None
+
+    def setUp(self):
+        """
+        Setup something before each test function
+        :return: None
+        """
+        self.runner = Runner()
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        """
+        Start up server on class initialization in daemon thread
+        :return: None
+        """
+        super(NetworkUsageTest, cls).setUpClass()
+        cls.server = HTTPServer(
+            server_address=(DefaultValues.HOST, DefaultValues.PORT),
+            RequestHandlerClass=TestClassHTTPRequestHandler,
+        )
+        server_daemon = Thread(target=cls.server.serve_forever, daemon=True)
+        server_daemon.start()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        super(NetworkUsageTest, cls).tearDownClass()
+        cls.server.shutdown()
+
+    def test_create_runner(self) -> None:
+        """
+        Test creation of the class instance
+        :return: None
+        """
+        self.assertIsNotNone(self.runner)
+        self.assertIsInstance(self.runner, Runner)
+
+    def test_local_server(self) -> None:
+        """
+        Test values from the local server
+        :return: None
+        """
+        response = self.runner.run(
+            hostname=f"http://{DefaultValues.HOST}:{DefaultValues.PORT}"
+        )
+        self.assertEqual(response.get("message"), "Success")
+        self.assertEqual(response.get("status"), "success")
+        result = response.get("result")
+        self.assertIn("BaseHTTP", result.get("Server"))
+        self.assertIn("Python", result.get("Server"))
+        self.assertEqual(result.get("Content-type"), "text/html; charset=utf-8")
+
+    def test_remote_server(self) -> None:
+        """
+        Test values from the remote server
+        :return: None
+        """
+        response = self.runner.run(hostname=f"http://www.example.com/")
+        self.assertEqual(response.get("message"), "Success")
+        self.assertEqual(response.get("status"), "success")
+        result = response.get("result")
+        self.assertIn("ECS", result.get("Server"))
+        self.assertEqual(result.get("Content-type"), "text/html; charset=UTF-8")