Merge pull request #17 from threat9/long-lived-connection-tcp

Handle multiple command calls to TCPServer within one connection.

README.md

@@ -1,4 +1,4 @@
-# thread9-test-bed
+# threat9-test-bed
 
 ## Installation
 ```bash

tests/service_mocks/test_tcp_service_mock.py

@@ -3,21 +3,18 @@
 from threat9_test_bed.service_mocks import TCPServiceMock
 
 
-def test_tcp_service_mock_add_banner():
+def test_tcp_service_mock_get_command_mock():
     with TCPServiceMock("127.0.0.1", 8023) as target:
         assert target.host == "127.0.0.1"
         assert target.port == 8023
 
+        mocked_scoo = target.get_command_mock(b"scoo")
+        mocked_scoo.return_value = b"bee"
         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/tcp_service/tcp_server.py

@@ -7,6 +7,7 @@
 
 class TCPServer(socketserver.ThreadingTCPServer):
     allow_reuse_address = True
+    daemon_threads = True
 
     def __init__(
             self,
@@ -27,6 +28,7 @@ def get_command_mock(self, command: bytes) -> mock.Mock:
 
 class TCPHandler(socketserver.BaseRequestHandler):
     def handle(self):
-        data = self.request.recv(1024)
-        handler = self.server.handlers[data]
-        self.request.sendall(handler())
+        while True:
+            data = self.request.recv(1024)
+            handler = self.server.handlers.get(data, lambda: b"")
+            self.request.sendall(handler())