Merge pull request #15 from threat9/tcp-service

WIP: TCPServiceMock

setup.py

@@ -32,7 +32,7 @@
     install_requires=[
         "click",
         "Faker",
-        "flask",
+        "flask>=1.0.0",
         "gunicorn",
         "pyopenssl",
         "requests",

tests/service_mocks/test_tcp_service_mock.py

@@ -0,0 +1,23 @@
+import socket
+
+from threat9_test_bed.service_mocks import TCPServiceMock
+
+
+def test_tcp_service_mock_add_banner():
+    with TCPServiceMock("127.0.0.1", 8023) as target:
+        assert target.host == "127.0.0.1"
+        assert target.port == 8023
+
+        mocked_doo = target.get_command_mock(b"doo")
+        mocked_doo.return_value = b"where are you?"
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.connect((target.host, target.port))
+            s.send(b"doo")
+            assert s.recv(1024) == b"where are you?"
+
+        mocked_scoo = target.get_command_mock(b"scoo")
+        mocked_scoo.return_value = b"bee"
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.connect((target.host, target.port))
+            s.send(b"scoo")
+            assert s.recv(1024) == b"bee"

threat9_test_bed/cli.py

@@ -3,8 +3,8 @@
 import click
 
 from .http_service.app import app
-from .http_service.gunicorn_server import GunicornServer
 from .scenarios import HttpScenario, TelnetScenario
+from .service_mocks.http_service_mock import WerkzeugBasedHttpService
 from .telnet_service.protocol import TelnetServerClientProtocol
 from .telnet_service.telnet_server import TelnetServer
 
@@ -30,13 +30,11 @@ def run_http_server(scenario, port):
     app.config.update(
         SCENARIO=HttpScenario[scenario],
     )
-    GunicornServer(
+    WerkzeugBasedHttpService(
         app=app,
-        bind=f"127.0.0.1:{port}",
-        worker_class="gthread",
-        threads=8,
-        accesslog="-",
-    ).run()
+        host=f"127.0.0.1",
+        port=port,
+    ).start()
     logger.debug(f"`http` server has been started on port {port}.")
 
 
@@ -54,14 +52,12 @@ def run_https_server(scenario, port):
     app.config.update(
         SCENARIO=HttpScenario[scenario],
     )
-    GunicornServer(
+    WerkzeugBasedHttpService(
         app=app,
-        bind=f"127.0.0.1:{port}",
-        worker_class="gthread",
-        threads=8,
+        host=f"127.0.0.1",
+        port=port,
         ssl=True,
-        accesslog="-",
-    ).run()
+    ).start()
     logger.debug(f"`https` server has been started on port {port}.")
 
 

threat9_test_bed/http_service/app.py

@@ -64,6 +64,7 @@ def redirect_():
 
 def timeout():
     time.sleep(60 * 60)
+    return found()
 
 
 def error():

threat9_test_bed/service_mocks/__init__.py

@@ -1,2 +1,4 @@
 from .http_scenario_service import HttpScenarioService  # noqa: F401
 from .http_service_mock import HttpServiceMock  # noqa: F401
+from .tcp_service_mock import TCPServiceMock  # noqa: F401
+from .telnet_service_mock import TelnetServiceMock  # noqa: F401

threat9_test_bed/service_mocks/base_service.py

@@ -49,7 +49,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         logger.debug(f"{self} has been terminated.")
 
     @staticmethod
-    def dib_port(port=0) -> (int, socket.socket):
+    def dib_port(port: int = 0) -> (int, socket.socket):
         socket_ = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         socket_.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         socket_.bind(("", port))

threat9_test_bed/service_mocks/http_scenario_service.py

@@ -2,12 +2,12 @@
 
 from ..http_service.app import app
 from ..scenarios import HttpScenario
-from .base_http_service import GunicornBasedHttpService
+from .base_http_service import WerkzeugBasedHttpService
 
 logger = logging.getLogger(__name__)
 
 
-class HttpScenarioService(GunicornBasedHttpService):
+class HttpScenarioService(WerkzeugBasedHttpService):
     def __init__(self, host: str, port: int, scenario: HttpScenario):
         app.config.update(SCENARIO=scenario)
         super().__init__(host, port, app)

threat9_test_bed/service_mocks/tcp_service_mock.py

@@ -0,0 +1,30 @@
+from logging import getLogger
+import threading
+from unittest import mock
+
+from ..tcp_service.tcp_server import TCPHandler, TCPServer
+from .base_service import BaseService
+
+logger = getLogger(__name__)
+
+
+class TCPServiceMock(BaseService):
+
+    def __init__(self, host: str, port: int):
+        super().__init__(host, port)
+        self.server = TCPServer((self.host, self.port), TCPHandler, False)
+        self.server_thread = threading.Thread(target=self.server.serve_forever)
+
+    def start(self):
+        self.server.server_bind()
+        self.server.server_activate()
+        self.server_thread.start()
+
+    def teardown(self):
+        self.server.shutdown()
+        self.server_thread.join()
+        self.server.server_close()
+
+    def get_command_mock(self, command: bytes) -> mock.Mock:
+        logger.debug(f"{self} mock for '{command}' has been added.")
+        return self.server.get_command_mock(command)

threat9_test_bed/service_mocks/telnet_service_mock.py

@@ -5,7 +5,7 @@
 from ..scenarios import TelnetScenario
 from ..telnet_service.protocol import TelnetServerClientProtocol
 from ..telnet_service.telnet_server import TelnetServer
-from .base_http_service import BaseService
+from .base_service import BaseService
 
 logger = logging.getLogger(__name__)
 

threat9_test_bed/tcp_service/tcp_server.py

@@ -0,0 +1,32 @@
+from logging import getLogger
+import socketserver
+from unittest import mock
+
+logger = getLogger(__name__)
+
+
+class TCPServer(socketserver.ThreadingTCPServer):
+    allow_reuse_address = True
+
+    def __init__(
+            self,
+            server_address,
+            request_handler_class,
+            bind_and_activate=True,
+    ):
+        super().__init__(
+            server_address, request_handler_class, bind_and_activate,
+        )
+        self.handlers = {}
+
+    def get_command_mock(self, command: bytes) -> mock.Mock:
+        mocked_handler = mock.MagicMock(name=command)
+        self.handlers[command] = mocked_handler
+        return mocked_handler
+
+
+class TCPHandler(socketserver.BaseRequestHandler):
+    def handle(self):
+        data = self.request.recv(1024)
+        handler = self.server.handlers[data]
+        self.request.sendall(handler())