├── tests └── service_mocks │ ├── __init__.py │ ├── test_udp_service_mock.py │ ├── test_tcp_service_mock.py │ ├── test_http_service_mock.py │ ├── test_http_scenario_service.py │ └── test_telnet_service_mock.py ├── threat9_test_bed ├── http_service │ ├── __init__.py │ ├── gunicorn_server.py │ └── app.py ├── tcp_service │ ├── __init__.py │ └── tcp_server.py ├── telnet_service │ ├── __init__.py │ ├── telnet_server.py │ └── protocol.py ├── udp_service │ ├── __init__.py │ └── udp_server.py ├── __init__.py ├── service_mocks │ ├── __init__.py │ ├── http_scenario_service.py │ ├── http_service_mock.py │ ├── tcp_service_mock.py │ ├── udp_service_mock.py │ ├── telnet_service_mock.py │ ├── base_service.py │ └── base_http_service.py ├── scenarios.py └── cli.py ├── setup.cfg ├── .circleci └── config.yml ├── setup.py ├── .gitignore └── README.md /tests/service_mocks/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /threat9_test_bed/http_service/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /threat9_test_bed/tcp_service/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /threat9_test_bed/telnet_service/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /threat9_test_bed/udp_service/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /threat9_test_bed/__init__.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | logging.getLogger(__name__).addHandler(logging.NullHandler()) 4 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [flake8] 2 | exclude = .cache,.git,.idea,.tox,.eggs,build,docs/*,*/_version.py,*.egg-info 3 | 4 | [isort] 5 | combine_as_imports = 1 6 | include_trailing_comma = 1 7 | known_first_party = 8 | threat9_test_bed 9 | force_sort_within_sections = 1 10 | multi_line_output = 3 11 | not_skip = __init__.py 12 | use_parentheses = 1 13 | -------------------------------------------------------------------------------- /threat9_test_bed/service_mocks/__init__.py: -------------------------------------------------------------------------------- 1 | from .http_scenario_service import HttpScenarioService # noqa: F401 2 | from .http_service_mock import HttpServiceMock # noqa: F401 3 | from .tcp_service_mock import TCPServiceMock # noqa: F401 4 | from .telnet_service_mock import TelnetServiceMock # noqa: F401 5 | from .udp_service_mock import UDPServiceMock # noqa: F401 6 | -------------------------------------------------------------------------------- /threat9_test_bed/service_mocks/http_scenario_service.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | from ..http_service.app import app 4 | from ..scenarios import HttpScenario 5 | from .base_http_service import WerkzeugBasedHttpService 6 | 7 | logger = logging.getLogger(__name__) 8 | 9 | 10 | class HttpScenarioService(WerkzeugBasedHttpService): 11 | def __init__(self, host: str, port: int, scenario: HttpScenario): 12 | app.config.update(SCENARIO=scenario) 13 | super().__init__(host, port, app) 14 | -------------------------------------------------------------------------------- /threat9_test_bed/scenarios.py: -------------------------------------------------------------------------------- 1 | from enum import Enum, unique 2 | 3 | 4 | @unique 5 | class HttpScenario(Enum): 6 | EMPTY_RESPONSE = "empty_response" 7 | TRASH = "trash" 8 | NOT_FOUND = "not_found" 9 | FOUND = "found" 10 | REDIRECT = "redirect" 11 | TIMEOUT = "timeout" 12 | ERROR = "error" 13 | 14 | @staticmethod 15 | def names(): 16 | return [element.name for element in HttpScenario] 17 | 18 | 19 | @unique 20 | class TelnetScenario(Enum): 21 | GENERIC = "generic" 22 | AUTHORIZED = "authorized" 23 | NOT_AUTHORIZED = "not_authorized" 24 | TIMEOUT = "timeout" 25 | 26 | @staticmethod 27 | def names(): 28 | return [element.name for element in TelnetScenario] 29 | -------------------------------------------------------------------------------- /threat9_test_bed/telnet_service/telnet_server.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import logging 3 | 4 | logger = logging.getLogger(__name__) 5 | 6 | 7 | class TelnetServer: 8 | def __init__(self, host: str, port: int, protocol): 9 | self.loop = asyncio.get_event_loop() 10 | 11 | coro = self.loop.create_server(protocol, host, port) 12 | self.server = self.loop.run_until_complete(coro) 13 | 14 | def run(self): 15 | logger.debug(f"Serving on {self.server.sockets[0].getsockname()}") 16 | try: 17 | self.loop.run_forever() 18 | except KeyboardInterrupt: 19 | pass 20 | finally: 21 | self.server.close() 22 | self.loop.run_until_complete(self.server.wait_closed()) 23 | -------------------------------------------------------------------------------- /threat9_test_bed/service_mocks/http_service_mock.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from unittest import mock 3 | 4 | from flask import Flask 5 | 6 | from .base_http_service import WerkzeugBasedHttpService 7 | 8 | logger = logging.getLogger(__name__) 9 | 10 | 11 | class HttpServiceMock(WerkzeugBasedHttpService): 12 | def __init__(self, host: str, port: int, ssl=False): 13 | super().__init__(host, port, Flask("target"), ssl) 14 | 15 | def get_route_mock(self, rule, **options): 16 | mocked_view = mock.MagicMock(name=rule, spec=lambda: None) 17 | self.app.add_url_rule(rule, 18 | endpoint=rule, 19 | view_func=mocked_view, 20 | **options) 21 | logger.debug(f"{self} mock for '{rule}' has been added.") 22 | return mocked_view 23 | -------------------------------------------------------------------------------- /.circleci/config.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | jobs: 3 | build: 4 | machine: true 5 | steps: 6 | - checkout 7 | - run: 8 | name: Install Python3.6 9 | command: | 10 | sudo add-apt-repository ppa:deadsnakes/ppa 11 | sudo apt-get update 12 | sudo apt-get install python3.6 13 | sudo apt-get install python3.6-dev 14 | sudo apt-get install python3.6-venv 15 | python3.6 -m venv ~/venv 16 | - run: 17 | name: Install dependencies 18 | command: | 19 | source ~/venv/bin/activate 20 | pip install .[tests] 21 | - run: 22 | name: Lint 23 | command: | 24 | source ~/venv/bin/activate 25 | flake8 26 | isort --check-only --diff 27 | unify --check-only --quote \" -r . 28 | - run: 29 | name: Run tests 30 | command: | 31 | source ~/venv/bin/activate 32 | py.test -vv tests/ -------------------------------------------------------------------------------- /tests/service_mocks/test_udp_service_mock.py: -------------------------------------------------------------------------------- 1 | import re 2 | import socket 3 | 4 | from threat9_test_bed.service_mocks import UDPServiceMock 5 | 6 | 7 | def test_udp_service_mock_get_command_mock(): 8 | with UDPServiceMock("127.0.0.1", 8023) as target: 9 | assert target.host == "127.0.0.1" 10 | assert target.port == 8023 11 | 12 | mocked_doo = target.get_command_mock(b"doo") 13 | mocked_doo.return_value = b"where are you?" 14 | 15 | mocked_scoo = target.get_command_mock(b"scoo") 16 | mocked_scoo.return_value = b"bee" 17 | 18 | mocked_scoo = target.get_command_mock(re.compile(b"\d\dfoo\d\d")) 19 | mocked_scoo.return_value = b"bar" 20 | 21 | with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s: 22 | s.connect((target.host, target.port)) 23 | s.send(b"doo") 24 | assert s.recv(1024) == b"where are you?" 25 | 26 | s.send(b"scoo") 27 | assert s.recv(1024) == b"bee" 28 | 29 | s.send(b"12foo34") 30 | assert s.recv(1024) == b"bar" 31 | 32 | s.send(b"56foo78") 33 | assert s.recv(1024) == b"bar" 34 | -------------------------------------------------------------------------------- /tests/service_mocks/test_tcp_service_mock.py: -------------------------------------------------------------------------------- 1 | import re 2 | import socket 3 | 4 | from threat9_test_bed.service_mocks import TCPServiceMock 5 | 6 | 7 | def test_tcp_service_mock_get_command_mock(): 8 | with TCPServiceMock("127.0.0.1", 8023) as target: 9 | assert target.host == "127.0.0.1" 10 | assert target.port == 8023 11 | 12 | mocked_scoo = target.get_command_mock(b"scoo") 13 | mocked_scoo.return_value = b"bee" 14 | 15 | mocked_doo = target.get_command_mock(b"doo") 16 | mocked_doo.return_value = b"where are you?" 17 | 18 | mocked_foo = target.get_command_mock(re.compile(b"\d\dfoo\d\d")) 19 | mocked_foo.return_value = b"bar" 20 | 21 | with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: 22 | s.connect((target.host, target.port)) 23 | 24 | s.send(b"doo") 25 | assert s.recv(1024) == b"where are you?" 26 | 27 | s.send(b"scoo") 28 | assert s.recv(1024) == b"bee" 29 | 30 | s.send(b"12foo34") 31 | assert s.recv(1024) == b"bar" 32 | 33 | s.send(b"56foo78") 34 | assert s.recv(1024) == b"bar" 35 | -------------------------------------------------------------------------------- /threat9_test_bed/service_mocks/tcp_service_mock.py: -------------------------------------------------------------------------------- 1 | from logging import getLogger 2 | import threading 3 | from typing import Pattern, Union 4 | from unittest import mock 5 | 6 | from ..tcp_service.tcp_server import TCPHandler, TCPServer 7 | from .base_service import BaseService 8 | 9 | logger = getLogger(__name__) 10 | 11 | 12 | class TCPServiceMock(BaseService): 13 | 14 | def __init__(self, host: str, port: int): 15 | super().__init__(host, port) 16 | self.server = TCPServer((self.host, self.port), TCPHandler, False) 17 | self.server_thread = threading.Thread(target=self.server.serve_forever) 18 | 19 | def start(self): 20 | self.server.server_bind() 21 | self.server.server_activate() 22 | self.server_thread.start() 23 | 24 | def teardown(self): 25 | self.server.shutdown() 26 | self.server_thread.join() 27 | self.server.server_close() 28 | 29 | def get_command_mock( 30 | self, 31 | command: Union[bytes, Pattern[bytes]], 32 | ) -> mock.MagicMock: 33 | logger.debug(f"{self} mock for '{command}' has been added.") 34 | return self.server.get_command_mock(command) 35 | -------------------------------------------------------------------------------- /threat9_test_bed/service_mocks/udp_service_mock.py: -------------------------------------------------------------------------------- 1 | from logging import getLogger 2 | import socket 3 | import threading 4 | from typing import Pattern, Union 5 | from unittest import mock 6 | 7 | from ..udp_service.udp_server import UDPHandler, UDPServer 8 | from .base_service import BaseService 9 | 10 | logger = getLogger(__name__) 11 | 12 | 13 | class UDPServiceMock(BaseService): 14 | 15 | socket_type = socket.SOCK_DGRAM 16 | 17 | def __init__(self, host: str, port: int): 18 | super().__init__(host, port) 19 | self.server = UDPServer((self.host, self.port), UDPHandler, False) 20 | self.server_thread = threading.Thread(target=self.server.serve_forever) 21 | 22 | def start(self): 23 | self.server.server_bind() 24 | self.server.server_activate() 25 | self.server_thread.start() 26 | 27 | def teardown(self): 28 | self.server.shutdown() 29 | self.server_thread.join() 30 | self.server.server_close() 31 | 32 | def get_command_mock( 33 | self, 34 | command: Union[bytes, Pattern[bytes]], 35 | ) -> mock.MagicMock: 36 | logger.debug(f"{self} mock for '{command}' has been added.") 37 | return self.server.get_command_mock(command) 38 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | 3 | from setuptools import find_packages, setup 4 | 5 | HERE = Path(__file__).parent.resolve() 6 | 7 | 8 | with open(str(HERE / "README.md"), encoding="utf-8") as f: 9 | long_description = f.read() 10 | 11 | setup( 12 | name="threat9-test-bed", 13 | use_scm_version={ 14 | "root": str(HERE), 15 | "write_to": str(HERE / "threat9_test_bed" / "_version.py"), 16 | }, 17 | url="http://threat9.com", 18 | author="Mariusz Kupidura", 19 | author_email="f4wkes@gmail.com", 20 | description="Threat9 Test Bed", 21 | long_description=long_description, 22 | packages=find_packages(where=str(HERE)), 23 | include_package_data=True, 24 | entry_points={ 25 | "console_scripts": [ 26 | "test-bed = threat9_test_bed.cli:cli", 27 | ], 28 | }, 29 | setup_requires=[ 30 | "setuptools_scm", 31 | ], 32 | install_requires=[ 33 | "click", 34 | "Faker", 35 | "flask>=1.0.0", 36 | "gunicorn", 37 | "pyopenssl", 38 | "requests", 39 | ], 40 | extras_require={ 41 | "tests": [ 42 | "flake8", 43 | "isort", 44 | "pytest", 45 | "unify", 46 | ] 47 | }, 48 | classifiers=[ 49 | "Operating System :: POSIX", 50 | "Intended Audience :: Developers", 51 | 52 | "Programming Language :: Python", 53 | "Programming Language :: Python :: 3.6", 54 | ], 55 | ) 56 | -------------------------------------------------------------------------------- /threat9_test_bed/http_service/gunicorn_server.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | import tempfile 3 | 4 | from OpenSSL import crypto 5 | from gunicorn.app.base import BaseApplication 6 | from werkzeug.serving import generate_adhoc_ssl_pair 7 | 8 | 9 | class GunicornServer(BaseApplication): 10 | def __init__(self, app, **kwargs): 11 | self.options = kwargs 12 | self.application = app 13 | super().__init__() 14 | 15 | def load_config(self): 16 | if self.options.get("ssl"): 17 | cert_path, pkey_path = self.generate_devel_ssl_pair() 18 | self.options["certfile"] = str(cert_path) 19 | self.options["keyfile"] = str(pkey_path) 20 | 21 | config = { 22 | key: value for key, value in self.options.items() 23 | if key in self.cfg.settings and value is not None 24 | } 25 | for key, value in config.items(): 26 | self.cfg.set(key.lower(), value) 27 | 28 | def load(self): 29 | return self.application 30 | 31 | @staticmethod 32 | def generate_devel_ssl_pair() -> (Path, Path): 33 | cert_path = Path(tempfile.gettempdir()) / "threat9-test-bed.crt" 34 | pkey_path = Path(tempfile.gettempdir()) / "threat9-test-bed.key" 35 | 36 | if not cert_path.exists() or not pkey_path.exists(): 37 | cert, pkey = generate_adhoc_ssl_pair() 38 | with open(cert_path, "wb") as f: 39 | f.write(crypto.dump_certificate(crypto.FILETYPE_PEM, cert)) 40 | with open(pkey_path, "wb") as f: 41 | f.write(crypto.dump_privatekey(crypto.FILETYPE_PEM, pkey)) 42 | 43 | return cert_path, pkey_path 44 | -------------------------------------------------------------------------------- /threat9_test_bed/udp_service/udp_server.py: -------------------------------------------------------------------------------- 1 | from logging import getLogger 2 | import socketserver 3 | from typing import Pattern, Union 4 | from unittest import mock 5 | 6 | logger = getLogger(__name__) 7 | 8 | 9 | class UDPServer(socketserver.ThreadingUDPServer): 10 | allow_reuse_address = True 11 | 12 | def __init__( 13 | self, 14 | server_address, 15 | request_handler_class, 16 | bind_and_activate=True, 17 | ): 18 | super().__init__( 19 | server_address, request_handler_class, bind_and_activate, 20 | ) 21 | self.handlers = {} 22 | 23 | def get_handler( 24 | self, 25 | command: Union[bytes, Pattern[bytes]], 26 | ) -> mock.MagicMock: 27 | handler = self.handlers.get(command) 28 | if handler: 29 | return handler 30 | 31 | for pattern_key in self.handlers: 32 | if isinstance(pattern_key, Pattern): 33 | if pattern_key.match(command): 34 | return self.handlers[pattern_key] 35 | 36 | handler = mock.MagicMock(name="default_handler") 37 | handler.return_value = b"" 38 | 39 | return handler 40 | 41 | def get_command_mock( 42 | self, 43 | command: Union[bytes, Pattern[bytes]], 44 | ) -> mock.MagicMock: 45 | mocked_handler = mock.MagicMock(name=command) 46 | self.handlers[command] = mocked_handler 47 | return mocked_handler 48 | 49 | 50 | class UDPHandler(socketserver.DatagramRequestHandler): 51 | def handle(self): 52 | data = self.rfile.read() 53 | handler = self.server.get_handler(data) 54 | self.wfile.write(handler()) 55 | -------------------------------------------------------------------------------- /threat9_test_bed/service_mocks/telnet_service_mock.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import threading 3 | import typing 4 | from unittest import mock 5 | 6 | from ..scenarios import TelnetScenario 7 | from ..telnet_service.protocol import TelnetServerClientProtocol 8 | from ..telnet_service.telnet_server import TelnetServer 9 | from .base_service import BaseService 10 | 11 | logger = logging.getLogger(__name__) 12 | 13 | 14 | class TelnetServiceMock(BaseService): 15 | def __init__(self, 16 | host: str, port: int, 17 | scenario: TelnetScenario=TelnetScenario.AUTHORIZED): 18 | super().__init__(host, port) 19 | self.protocol = TelnetServerClientProtocol(scenario) 20 | self.server = TelnetServer( 21 | self.host, 22 | self.port, 23 | lambda: self.protocol 24 | ) 25 | self.server_thread = threading.Thread(target=self.server.run) 26 | 27 | def start(self): 28 | self.server_thread.start() 29 | 30 | def teardown(self): 31 | self.server.loop.call_soon_threadsafe(self.server.loop.stop) 32 | self.server_thread.join() 33 | 34 | def get_command_mock( 35 | self, 36 | command: typing.Union[str, typing.Pattern[str]], 37 | ): 38 | command_mock = mock.MagicMock(name=command) 39 | self.protocol.add_command_handler(command, command_mock) 40 | return command_mock 41 | 42 | def add_credentials(self, login: str, password: str): 43 | """ Add custom credentials pair. """ 44 | self.protocol.add_credentials(login, password) 45 | 46 | def add_banner(self, banner: bytes): 47 | """ Add welcoming banner after connection. """ 48 | self.protocol.add_banner(banner) 49 | -------------------------------------------------------------------------------- /threat9_test_bed/tcp_service/tcp_server.py: -------------------------------------------------------------------------------- 1 | from logging import getLogger 2 | import socketserver 3 | from typing import Pattern, Union 4 | from unittest import mock 5 | 6 | logger = getLogger(__name__) 7 | 8 | 9 | class TCPServer(socketserver.ThreadingTCPServer): 10 | allow_reuse_address = True 11 | daemon_threads = True 12 | 13 | def __init__( 14 | self, 15 | server_address, 16 | request_handler_class, 17 | bind_and_activate=True, 18 | ): 19 | super().__init__( 20 | server_address, request_handler_class, bind_and_activate, 21 | ) 22 | self.handlers = {} 23 | 24 | def get_handler( 25 | self, 26 | command: Union[bytes, Pattern[bytes]], 27 | ) -> mock.MagicMock: 28 | handler = self.handlers.get(command) 29 | if handler: 30 | return handler 31 | 32 | for pattern_key in self.handlers: 33 | if isinstance(pattern_key, Pattern): 34 | if pattern_key.match(command): 35 | return self.handlers[pattern_key] 36 | 37 | handler = mock.MagicMock(name="default_handler") 38 | handler.return_value = b"" 39 | 40 | return handler 41 | 42 | def get_command_mock( 43 | self, 44 | command: Union[bytes, Pattern[bytes]], 45 | ) -> mock.MagicMock: 46 | mocked_handler = mock.MagicMock(name=command) 47 | self.handlers[command] = mocked_handler 48 | return mocked_handler 49 | 50 | 51 | class TCPHandler(socketserver.BaseRequestHandler): 52 | def handle(self): 53 | while True: 54 | data = self.request.recv(1024) 55 | handler = self.server.get_handler(data) 56 | self.request.sendall(handler()) 57 | -------------------------------------------------------------------------------- /threat9_test_bed/http_service/app.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import time 3 | 4 | from faker import Faker 5 | from flask import Flask, abort, g, redirect 6 | 7 | from ..scenarios import HttpScenario 8 | 9 | logger = logging.getLogger(__name__) 10 | 11 | 12 | def get_faker(): 13 | faker = g.get("faker", None) 14 | if faker is None: 15 | faker = Faker() 16 | g.user = faker 17 | return faker 18 | 19 | 20 | app = Flask(__name__) 21 | 22 | ALLOWED_METHODS = [ 23 | "GET", 24 | "POST", 25 | "PATCH", 26 | "PUT", 27 | "OPTIONS", 28 | "DELETE", 29 | ] 30 | 31 | 32 | @app.route("/", defaults={"path": ""}, methods=ALLOWED_METHODS) 33 | @app.route("/", methods=ALLOWED_METHODS) 34 | def catch_all(path): 35 | scenario_handler = SCENARIO_TO_HANDLER_MAP.get( 36 | app.config["SCENARIO"], 37 | error, 38 | ) 39 | logger.debug( 40 | f"Executing '{scenario_handler.__name__}' scenario handler..." 41 | ) 42 | return scenario_handler() 43 | 44 | 45 | def empty_response(): 46 | return "", 200 47 | 48 | 49 | def trash(): 50 | return get_faker().paragraph(variable_nb_sentences=True), 200 51 | 52 | 53 | def not_found(): 54 | abort(404) 55 | 56 | 57 | def found(): 58 | return "OK", 200 59 | 60 | 61 | def redirect_(): 62 | return redirect("https://www.youtube.com/watch?v=dQw4w9WgXcQ") 63 | 64 | 65 | def timeout(): 66 | time.sleep(60 * 60) 67 | return found() 68 | 69 | 70 | def error(): 71 | abort(500) 72 | 73 | 74 | SCENARIO_TO_HANDLER_MAP = { 75 | HttpScenario.EMPTY_RESPONSE: empty_response, 76 | HttpScenario.TRASH: trash, 77 | HttpScenario.NOT_FOUND: not_found, 78 | HttpScenario.FOUND: found, 79 | HttpScenario.REDIRECT: redirect_, 80 | HttpScenario.TIMEOUT: timeout, 81 | HttpScenario.ERROR: error, 82 | } 83 | -------------------------------------------------------------------------------- /tests/service_mocks/test_http_service_mock.py: -------------------------------------------------------------------------------- 1 | from unittest.mock import MagicMock 2 | 3 | import requests 4 | 5 | from threat9_test_bed.service_mocks.http_service_mock import HttpServiceMock 6 | 7 | 8 | def test_http_service_mock(): 9 | with HttpServiceMock("127.0.0.1", 8080) as target: 10 | assert target.host == "127.0.0.1" 11 | assert target.port == 8080 12 | mock = target.get_route_mock("/foo", methods=["POST"]) 13 | mock.return_value = "bar", 201 14 | assert isinstance(mock, MagicMock) 15 | response = requests.post(f"http://{target.host}:{target.port}/foo") 16 | assert response.status_code == 201 17 | assert response.content == b"bar" 18 | 19 | 20 | def test_https_service_mock(): 21 | with HttpServiceMock("127.0.0.1", 8080, ssl=True) as target: 22 | assert target.host == "127.0.0.1" 23 | assert target.port == 8080 24 | mock = target.get_route_mock("/foo", methods=["POST"]) 25 | mock.return_value = "bar", 201 26 | assert isinstance(mock, MagicMock) 27 | response = requests.post(f"https://{target.host}:{target.port}/foo", 28 | verify=False) 29 | assert response.status_code == 201 30 | assert response.content == b"bar" 31 | 32 | 33 | def test_http_service_mock_random_port(): 34 | with HttpServiceMock("127.0.0.1", 0) as target: 35 | assert target.host == "127.0.0.1" 36 | assert target.port in range(1024, 65535 + 1) 37 | mock = target.get_route_mock("/foo", methods=["POST"]) 38 | mock.return_value = "bar", 201 39 | assert isinstance(mock, MagicMock) 40 | response = requests.post(f"http://{target.host}:{target.port}/foo", 41 | verify=False) 42 | assert response.status_code == 201 43 | assert response.content == b"bar" 44 | -------------------------------------------------------------------------------- /threat9_test_bed/service_mocks/base_service.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import socket 3 | import time 4 | 5 | logger = logging.getLogger(__name__) 6 | 7 | 8 | class BaseService: 9 | 10 | socket_type = socket.SOCK_STREAM 11 | 12 | def __init__(self, host: str, port: int): 13 | self.host = host 14 | self.port, self.dibbed_port_socket = self.dib_port(port) 15 | 16 | def _wait_for_service(self): 17 | elapsed_time = 0 18 | start_time = time.time() 19 | while elapsed_time < 5: 20 | s = socket.socket(type=self.socket_type) 21 | s.settimeout(1) 22 | try: 23 | s.connect((self.host, self.port)) 24 | except (ConnectionRefusedError, ConnectionAbortedError, 25 | socket.timeout): 26 | elapsed_time = time.time() - start_time 27 | s.close() 28 | else: 29 | s.close() 30 | break 31 | else: 32 | raise TimeoutError(f"{self.__class__.__name__} " 33 | f"couldn't be set up before test.") 34 | 35 | def start(self): 36 | raise NotImplementedError() 37 | 38 | def teardown(self): 39 | raise NotImplementedError() 40 | 41 | def __enter__(self): 42 | logger.debug(f"Starting {self}...") 43 | self.dibbed_port_socket.close() 44 | self.start() 45 | self._wait_for_service() 46 | logger.debug(f"{self} has been started.") 47 | return self 48 | 49 | def __exit__(self, exc_type, exc_val, exc_tb): 50 | logger.debug(f"Terminating {self}...") 51 | self.teardown() 52 | logger.debug(f"{self} has been terminated.") 53 | 54 | @staticmethod 55 | def dib_port(port: int = 0) -> (int, socket.socket): 56 | socket_ = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 57 | socket_.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 58 | socket_.bind(("", port)) 59 | return int(socket_.getsockname()[1]), socket_ 60 | 61 | def __repr__(self): 62 | return ( 63 | f"{self.__class__.__name__}(host='{self.host}', port={self.port})" 64 | ) 65 | -------------------------------------------------------------------------------- /threat9_test_bed/cli.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | import click 4 | 5 | from .http_service.app import app 6 | from .scenarios import HttpScenario, TelnetScenario 7 | from .service_mocks.http_service_mock import WerkzeugBasedHttpService 8 | from .telnet_service.protocol import TelnetServerClientProtocol 9 | from .telnet_service.telnet_server import TelnetServer 10 | 11 | logger = logging.getLogger(__name__) 12 | 13 | 14 | @click.group() 15 | def cli(): 16 | pass 17 | 18 | 19 | @cli.command("http") 20 | @click.option("--port", 21 | default=8080, 22 | show_default=True, 23 | help="HTTP server port.") 24 | @click.option("--scenario", 25 | required=True, 26 | type=click.Choice(HttpScenario.names()), 27 | help="HTTP server behaviour.") 28 | def run_http_server(scenario, port): 29 | logger.debug("Starting `http` server...") 30 | app.config.update( 31 | SCENARIO=HttpScenario[scenario], 32 | ) 33 | WerkzeugBasedHttpService( 34 | app=app, 35 | host=f"127.0.0.1", 36 | port=port, 37 | ).start() 38 | logger.debug(f"`http` server has been started on port {port}.") 39 | 40 | 41 | @cli.command("https") 42 | @click.option("--port", 43 | default=8443, 44 | show_default=True, 45 | help="HTTPS server port.") 46 | @click.option("--scenario", 47 | required=True, 48 | type=click.Choice(HttpScenario.names()), 49 | help="HTTP server behaviour.") 50 | def run_https_server(scenario, port): 51 | logger.debug("Starting `https` server...") 52 | app.config.update( 53 | SCENARIO=HttpScenario[scenario], 54 | ) 55 | WerkzeugBasedHttpService( 56 | app=app, 57 | host=f"127.0.0.1", 58 | port=port, 59 | ssl=True, 60 | ).start() 61 | logger.debug(f"`https` server has been started on port {port}.") 62 | 63 | 64 | @cli.command("telnet") 65 | @click.option("--port", 66 | default=8023, 67 | show_default=True, 68 | help="Telnet server port.") 69 | @click.option("--scenario", 70 | required=True, 71 | type=click.Choice(TelnetScenario.names()), 72 | help="Telnet server behaviour.") 73 | def run_telnet_server(scenario, port): 74 | logger.debug("Starting `telnet` server...") 75 | TelnetServer( 76 | host="127.0.0.1", 77 | port=port, 78 | protocol=lambda: TelnetServerClientProtocol(TelnetScenario[scenario]), 79 | ).run() 80 | logger.debug(f"`telnet` server has been started on port {port}.") 81 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .idea 2 | */_version.py 3 | 4 | ### VisualStudioCode template 5 | .vscode/* 6 | !.vscode/settings.json 7 | !.vscode/tasks.json 8 | !.vscode/launch.json 9 | !.vscode/extensions.json 10 | ### VirtualEnv template 11 | .Python 12 | [Bb]in 13 | [Ii]nclude 14 | [Ll]ib 15 | [Ll]ib64 16 | [Ll]ocal 17 | [Ss]cripts 18 | pyvenv.cfg 19 | .venv 20 | pip-selfcheck.json 21 | ### macOS template 22 | # General 23 | .DS_Store 24 | .AppleDouble 25 | .LSOverride 26 | 27 | # Icon must end with two \r 28 | Icon 29 | 30 | # Thumbnails 31 | ._* 32 | 33 | # Files that might appear in the root of a volume 34 | .DocumentRevisions-V100 35 | .fseventsd 36 | .Spotlight-V100 37 | .TemporaryItems 38 | .Trashes 39 | .VolumeIcon.icns 40 | .com.apple.timemachine.donotpresent 41 | 42 | # Directories potentially created on remote AFP share 43 | .AppleDB 44 | .AppleDesktop 45 | Network Trash Folder 46 | Temporary Items 47 | .apdisk 48 | ### Python template 49 | # Byte-compiled / optimized / DLL files 50 | __pycache__/ 51 | *.py[cod] 52 | *$py.class 53 | 54 | # C extensions 55 | *.so 56 | 57 | # Distribution / packaging 58 | build/ 59 | develop-eggs/ 60 | dist/ 61 | downloads/ 62 | eggs/ 63 | .eggs/ 64 | lib/ 65 | lib64/ 66 | parts/ 67 | sdist/ 68 | var/ 69 | wheels/ 70 | *.egg-info/ 71 | .installed.cfg 72 | *.egg 73 | MANIFEST 74 | 75 | # PyInstaller 76 | # Usually these files are written by a python script from a template 77 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 78 | *.manifest 79 | *.spec 80 | 81 | # Installer logs 82 | pip-log.txt 83 | pip-delete-this-directory.txt 84 | 85 | # Unit test / coverage reports 86 | htmlcov/ 87 | .tox/ 88 | .coverage 89 | .coverage.* 90 | .cache 91 | nosetests.xml 92 | coverage.xml 93 | *.cover 94 | .hypothesis/ 95 | 96 | # Translations 97 | *.mo 98 | *.pot 99 | 100 | # Django stuff: 101 | *.log 102 | local_settings.py 103 | 104 | # Flask stuff: 105 | instance/ 106 | .webassets-cache 107 | 108 | # Scrapy stuff: 109 | .scrapy 110 | 111 | # Sphinx documentation 112 | docs/_build/ 113 | 114 | # PyBuilder 115 | target/ 116 | 117 | # Jupyter Notebook 118 | .ipynb_checkpoints 119 | 120 | # pyenv 121 | .python-version 122 | 123 | # celery beat schedule file 124 | celerybeat-schedule 125 | 126 | # SageMath parsed files 127 | *.sage.py 128 | 129 | # Environments 130 | .env 131 | env/ 132 | venv/ 133 | ENV/ 134 | env.bak/ 135 | venv.bak/ 136 | 137 | # Spyder project settings 138 | .spyderproject 139 | .spyproject 140 | 141 | # Rope project settings 142 | .ropeproject 143 | 144 | # mkdocs documentation 145 | /site 146 | 147 | # mypy 148 | .mypy_cache/ 149 | -------------------------------------------------------------------------------- /threat9_test_bed/service_mocks/base_http_service.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import multiprocessing 3 | import threading 4 | from wsgiref.simple_server import make_server as wsgiref_make_server 5 | 6 | from flask import Flask 7 | from werkzeug.serving import make_server as werkzeug_make_server 8 | 9 | from ..http_service.gunicorn_server import GunicornServer 10 | from .base_service import BaseService 11 | 12 | logger = logging.getLogger(__name__) 13 | 14 | 15 | class GunicornBasedHttpService(BaseService): 16 | """ `gunicorn` based HTTP service 17 | 18 | `Flask` application served using `gunicorn` in separate process using 19 | async workers (threads in this case). 20 | 21 | Application served by this base class suppose to handle unbuffered 22 | requests, `nginx` in this case is no option hence async workers. 23 | """ 24 | def __init__(self, host: str, port: int, app: Flask, ssl=False): 25 | super().__init__(host, port) 26 | self.app = app 27 | self.server = GunicornServer( 28 | app=self.app, 29 | bind=f"{self.host}:{self.port}", 30 | worker_class="gthread", 31 | threads=8, 32 | ssl=ssl, 33 | accesslog="-", 34 | ) 35 | self.server_process = multiprocessing.Process(target=self.server.run) 36 | 37 | def start(self): 38 | self.server_process.start() 39 | 40 | def teardown(self): 41 | self.server_process.terminate() 42 | self.server_process.join() 43 | 44 | 45 | class WSGIRefBasedHttpService(BaseService): 46 | """ `wsgiref` based HTTP service 47 | 48 | `Flask` application served using `wsgiref` in separate thread. 49 | 50 | We can leverage shared state between main thread and thread handling 51 | `wsgiref` server and dynamically attach `Mock` object as view functions. 52 | """ 53 | def __init__(self, host: str, port: int, app: Flask): 54 | super().__init__(host, port) 55 | self.app = app 56 | self.server = wsgiref_make_server(self.host, self.port, self.app) 57 | self.server_thread = threading.Thread(target=self.server.serve_forever) 58 | 59 | def start(self): 60 | self.server_thread.start() 61 | 62 | def teardown(self): 63 | self.server.shutdown() 64 | self.server_thread.join() 65 | self.server.server_close() 66 | 67 | 68 | class WerkzeugBasedHttpService(BaseService): 69 | """ `werkzeug` based HTTP service 70 | 71 | `Flask` application served using `werkzeug` in separate thread. 72 | 73 | We can leverage shared state between main thread and thread handling 74 | `wsgiref` server and dynamically attach `Mock` object as view functions. 75 | """ 76 | def __init__(self, host: str, port: int, app: Flask, ssl=False): 77 | super().__init__(host, port) 78 | self.app = app 79 | self.dibbed_port_socket.close() # werkzeug binds port on server init 80 | self.server = werkzeug_make_server( 81 | self.host, self.port, self.app, 82 | threaded=True, 83 | ssl_context="adhoc" if ssl else None 84 | ) 85 | self.server_thread = threading.Thread(target=self.server.serve_forever) 86 | 87 | def start(self): 88 | self.server_thread.start() 89 | 90 | def teardown(self): 91 | self.server.shutdown() 92 | self.server_thread.join() 93 | self.server.server_close() 94 | -------------------------------------------------------------------------------- /threat9_test_bed/telnet_service/protocol.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import logging 3 | import typing 4 | 5 | from faker import Faker 6 | 7 | from threat9_test_bed.scenarios import TelnetScenario 8 | 9 | logger = logging.getLogger(__name__) 10 | faker = Faker() 11 | 12 | 13 | def authorized(func): 14 | def _wrapper(self, data): 15 | message = data.decode().strip() 16 | if not self.authorized: 17 | if not self.login: 18 | self.login = message 19 | self.transport.write(b"Password: ") 20 | return 21 | else: 22 | self.password = message 23 | 24 | if (self.login, self.password) in self.creds: 25 | self.authorized = True 26 | self.transport.write(self.prompt.encode()) 27 | else: 28 | self.transport.write(b"\r\nLogin incorrect\r\ntarget login: ") 29 | self.login = None 30 | self.password = None 31 | return 32 | else: 33 | func(self, data) 34 | 35 | return _wrapper 36 | 37 | 38 | class GreedyList(list): 39 | def __contains__(self, item): 40 | return True 41 | 42 | 43 | class TelnetServerClientProtocol(asyncio.Protocol): 44 | def __init__(self, scenario: TelnetScenario): 45 | self.transport = None 46 | self.remote_address = None 47 | self.scenario = scenario 48 | 49 | self.login = None 50 | self.password = None 51 | self.authorized = False 52 | 53 | self.banner = b"" 54 | self._command_mocks = {} 55 | self._creds = [ 56 | ("admin", "admin"), 57 | ("kocia", "dupa"), 58 | ] 59 | 60 | @property 61 | def creds(self): 62 | if self.scenario is TelnetScenario.NOT_AUTHORIZED: 63 | return [] 64 | elif self.scenario is TelnetScenario.AUTHORIZED: 65 | return GreedyList() 66 | elif self.scenario is TelnetScenario.GENERIC: 67 | return self._creds 68 | 69 | @property 70 | def prompt(self): 71 | return f"{self.login}@target:~$ " 72 | 73 | def connection_made(self, transport: asyncio.Transport): 74 | self.remote_address = transport.get_extra_info("peername") 75 | logger.debug(f"Connection from {self.remote_address}") 76 | self.transport = transport 77 | if self.banner: 78 | self.transport.write(self.banner + b"\r\n") 79 | self.transport.write(b"Login: ") 80 | 81 | def _get_handler( 82 | self, 83 | command: typing.Union[str, typing.Pattern[str]], 84 | ) -> typing.Callable: 85 | 86 | handler = self._command_mocks.get(command) 87 | if handler: 88 | return handler 89 | 90 | for pattern_key in self._command_mocks: 91 | if isinstance(pattern_key, typing.Pattern): 92 | if pattern_key.match(command): 93 | return self._command_mocks[pattern_key] 94 | 95 | return faker.paragraph 96 | 97 | @authorized 98 | def data_received(self, data: bytes): 99 | logger.debug(f"{self.remote_address} send: {data}") 100 | command = data.decode().strip() 101 | handler = self._get_handler(command) 102 | self.transport.write( 103 | f"{handler()}\r\n"f"{self.prompt}".encode() 104 | ) 105 | 106 | def add_command_handler(self, command: str, handler: typing.Callable): 107 | self._command_mocks[command] = handler 108 | 109 | def add_credentials(self, login: str, password: str): 110 | self._creds.append((login, password)) 111 | 112 | def add_banner(self, banner: bytes): 113 | self.banner = banner 114 | -------------------------------------------------------------------------------- /tests/service_mocks/test_http_scenario_service.py: -------------------------------------------------------------------------------- 1 | import uuid 2 | 3 | import requests 4 | 5 | from threat9_test_bed.scenarios import HttpScenario 6 | from threat9_test_bed.service_mocks.http_scenario_service import ( 7 | HttpScenarioService, 8 | ) 9 | 10 | 11 | def test_http_scenario_service_trash(): 12 | with HttpScenarioService("127.0.0.1", 8080, 13 | scenario=HttpScenario.TRASH) as target: 14 | assert target.host == "127.0.0.1" 15 | assert target.port == 8080 16 | response = requests.get(f"http://{target.host}:{target.port}" 17 | f"/{uuid.uuid4().hex}") 18 | assert response.status_code == 200 19 | assert isinstance(response.content, bytes) 20 | assert len(response.content) > 0 21 | 22 | 23 | def test_http_scenario_service_empty_response(): 24 | with HttpScenarioService("127.0.0.1", 8080, 25 | scenario=HttpScenario.EMPTY_RESPONSE) as target: 26 | assert target.host == "127.0.0.1" 27 | assert target.port == 8080 28 | response = requests.post(f"http://{target.host}:{target.port}" 29 | f"/{uuid.uuid4().hex}") 30 | assert response.status_code == 200 31 | assert isinstance(response.content, bytes) 32 | assert len(response.content) == 0 33 | 34 | 35 | def test_http_scenario_service_error(): 36 | with HttpScenarioService("127.0.0.1", 8080, 37 | scenario=HttpScenario.ERROR) as target: 38 | assert target.host == "127.0.0.1" 39 | assert target.port == 8080 40 | response = requests.put(f"http://{target.host}:{target.port}" 41 | f"/{uuid.uuid4().hex}") 42 | assert response.status_code == 500 43 | assert isinstance(response.content, bytes) 44 | assert len(response.content) > 0 45 | 46 | 47 | def test_http_scenario_service_found(): 48 | with HttpScenarioService("127.0.0.1", 8080, 49 | scenario=HttpScenario.FOUND) as target: 50 | assert target.host == "127.0.0.1" 51 | assert target.port == 8080 52 | response = requests.patch(f"http://{target.host}:{target.port}" 53 | f"/{uuid.uuid4().hex}") 54 | assert response.status_code == 200 55 | assert response.content == b"OK" 56 | 57 | 58 | def test_http_scenario_service_not_found(): 59 | with HttpScenarioService("127.0.0.1", 8080, 60 | scenario=HttpScenario.NOT_FOUND) as target: 61 | assert target.host == "127.0.0.1" 62 | assert target.port == 8080 63 | response = requests.post(f"http://{target.host}:{target.port}" 64 | f"/{uuid.uuid4().hex}") 65 | assert response.status_code == 404 66 | assert isinstance(response.content, bytes) 67 | assert len(response.content) > 0 68 | 69 | 70 | def test_http_scenario_service_redirect(): 71 | with HttpScenarioService("127.0.0.1", 8080, 72 | scenario=HttpScenario.REDIRECT) as target: 73 | assert target.host == "127.0.0.1" 74 | assert target.port == 8080 75 | response = requests.get(f"http://{target.host}:{target.port}" 76 | f"/{uuid.uuid4().hex}", 77 | allow_redirects=False) 78 | assert response.status_code == 302 79 | 80 | 81 | def test_http_scenario_service_random_port(): 82 | with HttpScenarioService("127.0.0.1", 0, 83 | scenario=HttpScenario.FOUND) as target: 84 | assert target.host == "127.0.0.1" 85 | assert target.port in range(1024, 65535 + 1) 86 | response = requests.post(f"http://{target.host}:{target.port}/foo") 87 | assert response.status_code == 200 88 | assert response.content == b"OK" 89 | -------------------------------------------------------------------------------- /tests/service_mocks/test_telnet_service_mock.py: -------------------------------------------------------------------------------- 1 | import re 2 | from telnetlib import Telnet 3 | from unittest.mock import MagicMock 4 | import uuid 5 | 6 | from threat9_test_bed.scenarios import TelnetScenario 7 | from threat9_test_bed.service_mocks.telnet_service_mock import ( 8 | TelnetServiceMock, 9 | ) 10 | 11 | 12 | def test_telnet_service_mock_not_authorized(): 13 | with TelnetServiceMock("127.0.0.1", 8023, 14 | scenario=TelnetScenario.NOT_AUTHORIZED) as target: 15 | assert target.host == "127.0.0.1" 16 | assert target.port == 8023 17 | mock = target.get_command_mock("foo") 18 | mock.return_value = "bar" 19 | assert isinstance(mock, MagicMock) 20 | tn = Telnet(target.host, target.port, timeout=1.0) 21 | 22 | tn.expect([b"Login: ", b"login: "], 1.0) 23 | tn.write(b"admin" + b"\r\n") 24 | 25 | tn.expect([b"Password: ", b"password"], 1.0) 26 | tn.write(b"admin" + b"\r\n") 27 | 28 | _, match_object, _ = tn.expect([b"Login incorrect"], 1.0) 29 | assert match_object 30 | tn.close() 31 | 32 | 33 | def test_telnet_service_mock_generic(): 34 | with TelnetServiceMock("127.0.0.1", 8023, 35 | scenario=TelnetScenario.GENERIC) as target: 36 | assert target.host == "127.0.0.1" 37 | assert target.port == 8023 38 | mock = target.get_command_mock("foo") 39 | mock.return_value = "bar" 40 | assert isinstance(mock, MagicMock) 41 | tn = Telnet(target.host, target.port, timeout=1.0) 42 | 43 | tn.expect([b"Login: ", b"login: "], 1.0) 44 | tn.write(b"admin" + b"\r\n") 45 | 46 | tn.expect([b"Password: ", b"password"], 1.0) 47 | tn.write(b"admin" + b"\r\n") 48 | 49 | tn.expect([b"admin@target:~\$"], 1.0) 50 | tn.write(b"foo" + b"\r\n") 51 | _, match_object, _ = tn.expect([b"bar"], 1.0) 52 | assert match_object 53 | tn.close() 54 | 55 | 56 | def test_telnet_service_mock_authorized(): 57 | with TelnetServiceMock("127.0.0.1", 0, 58 | scenario=TelnetScenario.AUTHORIZED) as target: 59 | assert target.host == "127.0.0.1" 60 | assert target.port in range(1024, 65535 + 1) 61 | mock = target.get_command_mock("foo") 62 | mock.return_value = "bar" 63 | assert isinstance(mock, MagicMock) 64 | tn = Telnet(target.host, target.port, timeout=1.0) 65 | 66 | login = uuid.uuid4().hex.encode() 67 | tn.expect([b"Login: ", b"login: "], 1.0) 68 | tn.write(login + b"\r\n") 69 | 70 | password = uuid.uuid4().hex.encode() 71 | tn.expect([b"Password: ", b"password"], 1.0) 72 | tn.write(password + b"\r\n") 73 | 74 | tn.expect([login + b"@target:~\$"], 1.0) 75 | tn.write(b"foo" + b"\r\n") 76 | _, match_object, _ = tn.expect([b"bar"], 1.0) 77 | assert match_object 78 | tn.close() 79 | 80 | 81 | def test_telnet_service_mock_add_credentials(): 82 | with TelnetServiceMock("127.0.0.1", 8023, 83 | scenario=TelnetScenario.GENERIC) as target: 84 | login = uuid.uuid4().hex.encode() 85 | password = uuid.uuid4().hex.encode() 86 | 87 | assert target.host == "127.0.0.1" 88 | assert target.port == 8023 89 | 90 | tn = Telnet(target.host, target.port, timeout=1.0) 91 | 92 | tn.expect([b"Login: ", b"login: "], 1.0) 93 | tn.write(login + b"\r\n") 94 | 95 | tn.expect([b"Password: ", b"password"], 1.0) 96 | tn.write(password + b"\r\n") 97 | 98 | _, match_object, _ = tn.expect([b"Login incorrect"], 1.0) 99 | assert match_object 100 | 101 | target.add_credentials(login.decode(), password.decode()) 102 | 103 | tn.expect([b"Login: ", b"login: "], 1.0) 104 | tn.write(login + b"\r\n") 105 | 106 | tn.expect([b"Password: ", b"password"], 1.0) 107 | tn.write(password + b"\r\n") 108 | 109 | _, match_object, foo = tn.expect([login + b"@target:~\$"], 1.0) 110 | 111 | assert match_object, foo.decode() 112 | tn.close() 113 | 114 | 115 | def test_telnet_service_mock_add_banner(): 116 | with TelnetServiceMock("127.0.0.1", 8023, 117 | scenario=TelnetScenario.GENERIC) as target: 118 | banner = b"Scoobeedoobeedoo where are you?" 119 | target.add_banner(banner) 120 | 121 | assert target.host == "127.0.0.1" 122 | assert target.port == 8023 123 | 124 | tn = Telnet(target.host, target.port, timeout=1.0) 125 | 126 | _, match_object, _ = tn.expect([banner], 1.0) 127 | assert match_object 128 | 129 | _, match_object, foo = tn.expect([b"Login: ", b"login: "], 1.0) 130 | assert match_object, foo.decode() 131 | tn.close() 132 | 133 | 134 | def test_telnet_service_mock_regexp_mock(): 135 | with TelnetServiceMock("127.0.0.1", 8023, 136 | scenario=TelnetScenario.GENERIC) as target: 137 | assert target.host == "127.0.0.1" 138 | assert target.port == 8023 139 | 140 | mock = target.get_command_mock(re.compile("\d\dscoo\d\d")) 141 | mock.return_value = "bee" 142 | assert isinstance(mock, MagicMock) 143 | 144 | tn = Telnet(target.host, target.port, timeout=1.0) 145 | 146 | tn.expect([b"Login: ", b"login: "], 1.0) 147 | tn.write(b"admin" + b"\r\n") 148 | 149 | tn.expect([b"Password: ", b"password"], 1.0) 150 | tn.write(b"admin" + b"\r\n") 151 | 152 | tn.expect([b"admin@target:~\$"], 1.0) 153 | tn.write(b"12scoo34" + b"\r\n") 154 | _, match_object, _ = tn.expect([b"bee"], 1.0) 155 | assert match_object 156 | 157 | tn.write(b"56scoo78" + b"\r\n") 158 | _, match_object, _ = tn.expect([b"bee"], 1.0) 159 | assert match_object 160 | 161 | tn.close() 162 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # threat9-test-bed 2 | 3 | ## Installation 4 | ```bash 5 | $ pip install git+https://github.com/threat9/threat9-test-bed.git 6 | ``` 7 | 8 | ## Test utilities 9 | 10 | ### `HttpServiceMock` 11 | `HttpServiceMock` is a `flask` application that allows for adding 12 | `unittests.mock` as view functions. This gives us ability to setup dummy 13 | http services on demand for testing purposes. 14 | 15 | ```python 16 | from threat9_test_bed.service_mocks import HttpServiceMock 17 | 18 | from foo import ExploitUnderTest 19 | 20 | 21 | def test_exploit(): 22 | with HttpServiceMock("localhost", 8080) as target: 23 | cgi_mock = target.get_route_mock("/cgi-bin/cgiSrv.cgi", 24 | methods=["POST"]) 25 | cgi_mock.return_value = 'foo status="doing" bar' 26 | check_mock = target.get_route_mock("/routersploit.check", 27 | methods=["GET", "POST"]) 28 | check_mock.return_value = 'root' 29 | 30 | exploit = ExploitUnderTest(f'http://{target.host}', target.port) 31 | assert exploit.check() is True 32 | cgi_mock.assert_called_once() 33 | assert check_mock.call_count == 2 34 | ``` 35 | It is very convenient to use `py.test` library and it's fixture abilities. 36 | Such fixture will perform setup and teardown automatically before each test. 37 | All we have to do is to pass `target` as the test argument. 38 | ```python 39 | import pytest 40 | from threat9_test_bed.service_mocks import HttpServiceMock 41 | 42 | from foo import ExploitUnderTest 43 | 44 | 45 | @pytest.fixture 46 | def target(): 47 | with HttpServiceMock("localhost", 8080) as target_: 48 | yield target_ 49 | 50 | 51 | def test_exploit(target): 52 | cgi_mock = target.get_route_mock("/cgi-bin/cgiSrv.cgi", 53 | methods=["POST"]) 54 | cgi_mock.return_value = 'foo status="doing" bar' 55 | check_mock = target.get_route_mock("/routersploit.check", 56 | methods=["GET", "POST"]) 57 | check_mock.return_value = 'root' 58 | 59 | exploit = ExploitUnderTest(f'http://{target.host}', target.port) 60 | assert exploit.check() is True 61 | cgi_mock.assert_called_once() 62 | assert check_mock.call_count == 2 63 | ``` 64 | #### Adhoc SSL support 65 | You can serve `HttpScenarioService` using adhoc SSL certificate by setting 66 | `ssl` keyword argument to `True`: 67 | 68 | ```python 69 | from threat9_test_bed.service_mocks import HttpServiceMock 70 | 71 | @pytest.fixture 72 | def trash_target(): 73 | with HttpServiceMock("127.0.0.1", 0, ssl=True) as http_service: 74 | yield http_service 75 | ``` 76 | 77 | ### `HttpScenarioService` 78 | `HttpScenarioService` allows for creating test utilities using pre-defined 79 | [scenarios](#http-scenarios) 80 | ```python 81 | import pytest 82 | 83 | from threat9_test_bed.scenarios import HttpScenario 84 | from threat9_test_bed.service_mocks import HttpScenarioService 85 | 86 | 87 | @pytest.fixture(scope="session") 88 | def empty_target(): 89 | with HttpScenarioService("127.0.0.1", 8081, 90 | HttpScenario.EMPTY_RESPONSE) as http_service: 91 | yield http_service 92 | 93 | 94 | @pytest.fixture(scope="session") 95 | def trash_target(): 96 | with HttpScenarioService("127.0.0.1", 8082, 97 | HttpScenario.TRASH) as http_service: 98 | yield http_service 99 | 100 | ``` 101 | 102 | #### Adhoc SSL support 103 | You can serve `HttpScenarioService` using adhoc SSL certificate by setting 104 | `ssl` keyword argument to `True`: 105 | 106 | ```python 107 | from threat9_test_bed.service_mocks import HttpScenarioService 108 | 109 | @pytest.fixture(scope="session") 110 | def trash_target(): 111 | with HttpScenarioService("127.0.0.1", 8443, HttpScenario.TRASH, 112 | ssl=True) as http_service: 113 | yield http_service 114 | ``` 115 | 116 | ### `TelnetServiceMock` 117 | `TelnetServiceMock` allows for creating test utilities using pre-defined 118 | [scenarios](#telnet-scenarios) as well as attaching `unittests.mock` 119 | as command handlers. This gives us ability to setup dummy telnet service 120 | on demand for testing purposes. 121 | ```python 122 | from telnetlib import Telnet 123 | 124 | import pytest 125 | 126 | from threat9_test_bed.service_mocks.telnet_service_mock import TelnetServiceMock 127 | from threat9_test_bed.scenarios import TelnetScenarios 128 | 129 | 130 | @pytest.fixture 131 | def generic_target(): 132 | with TelnetServiceMock("127.0.0.1", 8023, 133 | TelnetScenarios.AUTHORIZED) as telnet_service: 134 | yield telnet_service 135 | 136 | 137 | def test_telnet(generic_target): 138 | command_mock = target.get_command_mock("scoobeedoobeedoo") 139 | command_mock.return_value = "Where are you?" 140 | 141 | tn = Telnet(target.host, target.port, timeout=5) 142 | tn.expect([b"Login: ", b"login: "], 5) 143 | tn.write(b"admin" + b"\r\n") 144 | 145 | tn.expect([b"Password: ", b"password"], 5) 146 | tn.write(b"admin" + b"\r\n") 147 | 148 | tn.expect([b"admin@target:~$"], 5) 149 | tn.write(b"scoobeedoobeedoo" + b"\r\n") 150 | _, match_object, _ = tn.expect([b"Where are you?"], 5) 151 | 152 | tn.close() 153 | 154 | assert match_object 155 | ``` 156 | 157 | ### Random port 158 | To avoid `port` collison during tests you can tell test utilities to set 159 | it for you by passing `0` 160 | ```python 161 | @pytest.fixture(scope="session") 162 | def trash_target(): 163 | with HttpScenarioService("127.0.0.1", 0, 164 | HttpScenario.TRASH) as http_service: 165 | yield http_service 166 | ``` 167 | 168 | ## Services 169 | ### `http` 170 | ```bash 171 | $ test-bed http 172 | ``` 173 | #### `http` scenarios 174 | |Scenario | Behavior | 175 | |-------------------|---------------| 176 | |`EMPTY_RESPONSE` | returns empty response with `200` status code | 177 | |`TRASH` | returns 100 characters long gibberish with `200` status code | 178 | |`NOT_FOUND` | returns `404` status code | 179 | |`FOUND` | returns _OK_ with `200` status code | 180 | |`REDIRECT` | redirects you with `302` status code | 181 | |`TIMEOUT` | sleep the server for 1 hour which effectively times out the request | 182 | |`ERROR` | returns `500` status code | | 183 | 184 | ```bash 185 | $ test-bed http --scenario TRASH 186 | ``` 187 | 188 | ### `https` 189 | ```bash 190 | $ test-bed https 191 | ``` 192 | 193 | #### `https` scenarios 194 | |Scenario | Behavior | 195 | |-------------------|---------------| 196 | |`EMPTY_RESPONSE` | returns empty response with `200` status code | 197 | |`TRASH` | returns 100 characters long gibberish with `200` status code | 198 | |`NOT_FOUND` | returns `404` status code | 199 | |`FOUND` | returns _OK_ with `200` status code | 200 | |`REDIRECT` | redirects you with `302` status code | 201 | |`TIMEOUT` | sleep the server for 1 hour which effectively times out the request | 202 | |`ERROR` | returns `500` status code | 203 | 204 | ```bash 205 | $ test-bed https --scenario FOUND 206 | ``` 207 | 208 | ### `telnet` 209 | After successful authorization elnet service responds with random 210 | _Lorem ipsum..._ for every command 211 | ```bash 212 | $ test-bed telnet 213 | ``` 214 | #### `telnet` scenarios 215 | |Scenario | Behavior | 216 | |-------------------|---------------| 217 | |`AUTHORIZED` | Any authorization attempt ends with success | 218 | |`NOT_AUTHORIZED` | Every authorization attempt ends with failure | 219 | |`GENERIC` | Authorization using `admin/admin` credentials | 220 | |`TIMEOUT` | Server hangs as soon as client has been connected | 221 | 222 | ```bash 223 | $ test-bed telnet --scenario GENERIC 224 | ``` 225 | 226 | ## Troubleshooting 227 | > I can't start my `https` service on port 443 due to `PermissionError` 228 | 229 | Running services on it's default port may need extra privileges thus 230 | prepending command with `sudo` should do the trick e.g. 231 | ```bash 232 | $ sudo test-bed https --scenario TRASH --port 443 233 | [2017-09-16 12:51:18,137: INFO/werkzeug] * Running on https://127.0.0.1:443/ (Press CTRL+C to quit) 234 | ``` 235 | This solution can be applied to other services and it's default ports as well. --------------------------------------------------------------------------------