From 48180b85d17fa1486f692a7ec50abcbc5d2806d9 Mon Sep 17 00:00:00 2001 From: Salad Dais Date: Tue, 15 Jun 2021 18:46:51 +0000 Subject: [PATCH] Export proxy test utils for use in addon test suites --- hippolyzer/lib/proxy/lludp_proxy.py | 2 +- hippolyzer/lib/proxy/socks_proxy.py | 4 +- hippolyzer/lib/proxy/test_utils.py | 80 ++++++++++++++++++++++++++ tests/proxy/__init__.py | 78 ------------------------- tests/proxy/integration/test_addons.py | 9 +-- tests/proxy/integration/test_http.py | 3 +- tests/proxy/integration/test_lludp.py | 9 ++- tests/proxy/test_httpflows.py | 3 +- tests/proxy/test_object_manager.py | 4 +- 9 files changed, 93 insertions(+), 99 deletions(-) create mode 100644 hippolyzer/lib/proxy/test_utils.py diff --git a/hippolyzer/lib/proxy/lludp_proxy.py b/hippolyzer/lib/proxy/lludp_proxy.py index 8e304eb..87d9b69 100644 --- a/hippolyzer/lib/proxy/lludp_proxy.py +++ b/hippolyzer/lib/proxy/lludp_proxy.py @@ -43,7 +43,7 @@ class InterceptingLLUDPProxyProtocol(UDPProxyProtocol): ) raise PermissionError(f"UDPBanned message {msg.name}") - def _handle_proxied_packet(self, packet: UDPPacket): + def handle_proxied_packet(self, packet: UDPPacket): region: Optional[ProxiedRegion] = None # Try to do an initial region lookup so we have it for handle_proxied_packet() if self.session: diff --git a/hippolyzer/lib/proxy/socks_proxy.py b/hippolyzer/lib/proxy/socks_proxy.py index a79387c..6d98ac7 100644 --- a/hippolyzer/lib/proxy/socks_proxy.py +++ b/hippolyzer/lib/proxy/socks_proxy.py @@ -207,12 +207,12 @@ class UDPProxyProtocol(asyncio.DatagramProtocol): ) try: - self._handle_proxied_packet(src_packet) + self.handle_proxied_packet(src_packet) except: logging.exception("Barfed while handling UDP packet!") raise - def _handle_proxied_packet(self, packet): + def handle_proxied_packet(self, packet): self.transport.send_packet(packet) def close(self): diff --git a/hippolyzer/lib/proxy/test_utils.py b/hippolyzer/lib/proxy/test_utils.py new file mode 100644 index 0000000..8231a26 --- /dev/null +++ b/hippolyzer/lib/proxy/test_utils.py @@ -0,0 +1,80 @@ +import asyncio +import unittest +from typing import Any, Optional, List, Tuple + +from hippolyzer.lib.base.datatypes import UUID +from hippolyzer.lib.base.message.message import Message +from hippolyzer.lib.base.message.udpserializer import UDPMessageSerializer +from hippolyzer.lib.base.network.transport import UDPPacket, AbstractUDPTransport, ADDR_TUPLE +from hippolyzer.lib.proxy.lludp_proxy import InterceptingLLUDPProxyProtocol +from hippolyzer.lib.proxy.region import ProxiedRegion +from hippolyzer.lib.proxy.sessions import SessionManager +from hippolyzer.lib.proxy.settings import ProxySettings +from hippolyzer.lib.proxy.transport import SOCKS5UDPTransport + + +class BaseProxyTest(unittest.IsolatedAsyncioTestCase): + def setUp(self) -> None: + self.client_addr = ("127.0.0.1", 1) + self.region_addr = ("127.0.0.1", 3) + self.circuit_code = 1234 + self.session_manager = SessionManager(ProxySettings()) + self.session = self.session_manager.create_session({ + "session_id": UUID.random(), + "secure_session_id": UUID.random(), + "agent_id": UUID.random(), + "circuit_code": self.circuit_code, + "sim_ip": self.region_addr[0], + "sim_port": self.region_addr[1], + "region_x": 0, + "region_y": 123, + "seed_capability": "https://test.localhost:4/foo", + }) + self.transport = MockTransport() + self.protocol = InterceptingLLUDPProxyProtocol( + self.client_addr, self.session_manager) + self.protocol.transport = self.transport + self.serializer = UDPMessageSerializer() + self.session.objects.track_region_objects(123) + + async def _wait_drained(self): + await asyncio.sleep(0.001) + + def _setup_default_circuit(self): + self._setup_region_circuit(self.session.regions[-1]) + self.session.main_region = self.session.regions[-1] + + def _setup_region_circuit(self, region: ProxiedRegion): + # Not going to send a UseCircuitCode, so have to pretend we already did the + # client -> region NAT hole-punching + self.protocol.session = self.session + self.protocol.far_to_near_map[region.circuit_addr] = self.client_addr + self.session_manager.claim_session(self.session.id) + self.session.open_circuit(self.client_addr, region.circuit_addr, + self.protocol.transport) + + def _msg_to_packet(self, msg: Message, src, dst) -> UDPPacket: + return UDPPacket(src_addr=src, dst_addr=dst, data=self.serializer.serialize(msg), + direction=msg.direction) + + def _msg_to_datagram(self, msg: Message, src, dst, socks_header=True): + packet = self._msg_to_packet(msg, src, dst) + return SOCKS5UDPTransport.serialize(packet, force_socks_header=socks_header) + + +class MockTransport(AbstractUDPTransport): + def sendto(self, data: Any, addr: Optional[ADDR_TUPLE] = ...) -> None: + pass + + def abort(self) -> None: + pass + + def close(self) -> None: + pass + + def __init__(self): + super().__init__() + self.packets: List[Tuple[bytes, Tuple[str, int]]] = [] + + def send_packet(self, packet: UDPPacket) -> None: + self.packets.append((packet.data, packet.dst_addr)) diff --git a/tests/proxy/__init__.py b/tests/proxy/__init__.py index f6b9d17..e69de29 100644 --- a/tests/proxy/__init__.py +++ b/tests/proxy/__init__.py @@ -1,78 +0,0 @@ -import asyncio -from typing import * -import unittest - -from hippolyzer.lib.base.datatypes import UUID -from hippolyzer.lib.base.message.udpserializer import UDPMessageSerializer -from hippolyzer.lib.base.network.transport import AbstractUDPTransport, UDPPacket, ADDR_TUPLE -from hippolyzer.lib.proxy.lludp_proxy import InterceptingLLUDPProxyProtocol -from hippolyzer.lib.base.message.message import Message -from hippolyzer.lib.proxy.region import ProxiedRegion -from hippolyzer.lib.proxy.sessions import SessionManager -from hippolyzer.lib.proxy.settings import ProxySettings -from hippolyzer.lib.proxy.transport import SOCKS5UDPTransport - - -class MockTransport(AbstractUDPTransport): - def sendto(self, data: Any, addr: Optional[ADDR_TUPLE] = ...) -> None: - pass - - def abort(self) -> None: - pass - - def close(self) -> None: - pass - - def __init__(self): - super().__init__() - self.packets: List[Tuple[bytes, Tuple[str, int]]] = [] - - def send_packet(self, packet: UDPPacket) -> None: - self.packets.append((packet.data, packet.dst_addr)) - - -class BaseProxyTest(unittest.IsolatedAsyncioTestCase): - def setUp(self) -> None: - self.client_addr = ("127.0.0.1", 1) - self.region_addr = ("127.0.0.1", 3) - self.circuit_code = 1234 - self.session_manager = SessionManager(ProxySettings()) - self.session = self.session_manager.create_session({ - "session_id": UUID.random(), - "secure_session_id": UUID.random(), - "agent_id": UUID.random(), - "circuit_code": self.circuit_code, - "sim_ip": self.region_addr[0], - "sim_port": self.region_addr[1], - "region_x": 0, - "region_y": 123, - "seed_capability": "https://test.localhost:4/foo", - }) - self.transport = MockTransport() - self.protocol = InterceptingLLUDPProxyProtocol( - self.client_addr, self.session_manager) - self.protocol.transport = self.transport - self.serializer = UDPMessageSerializer() - self.session.objects.track_region_objects(123) - - async def _wait_drained(self): - await asyncio.sleep(0.001) - - def _setup_default_circuit(self): - self._setup_region_circuit(self.session.regions[-1]) - self.session.main_region = self.session.regions[-1] - - def _setup_region_circuit(self, region: ProxiedRegion): - # Not going to send a UseCircuitCode, so have to pretend we already did the - # client -> region NAT hole-punching - self.protocol.session = self.session - self.protocol.far_to_near_map[region.circuit_addr] = self.client_addr - self.session_manager.claim_session(self.session.id) - self.session.open_circuit(self.client_addr, region.circuit_addr, - self.protocol.transport) - - def _msg_to_datagram(self, msg: Message, src, dst, direction, socks_header=True): - serialized = self.serializer.serialize(msg) - packet = UDPPacket(src_addr=src, dst_addr=dst, data=serialized, - direction=direction) - return SOCKS5UDPTransport.serialize(packet, force_socks_header=socks_header) diff --git a/tests/proxy/integration/test_addons.py b/tests/proxy/integration/test_addons.py index fd11edb..7beb14b 100644 --- a/tests/proxy/integration/test_addons.py +++ b/tests/proxy/integration/test_addons.py @@ -15,11 +15,9 @@ from hippolyzer.lib.proxy.addon_utils import ( ) from hippolyzer.lib.proxy.addons import AddonManager from hippolyzer.lib.proxy.commands import handle_command -from hippolyzer.lib.base.network.transport import Direction from hippolyzer.lib.proxy.region import ProxiedRegion from hippolyzer.lib.proxy.sessions import Session - -from .. import BaseProxyTest +from hippolyzer.lib.proxy.test_utils import BaseProxyTest class MockAddon(BaseAddon): @@ -82,9 +80,8 @@ class AddonIntegrationTests(BaseProxyTest): Block("AgentData", AgentID=self.session.agent_id, SessionID=self.session.id), Block("ChatData", Message=command, Channel=AddonManager.COMMAND_CHANNEL, fill_missing=True), ) - packet = self._msg_to_datagram(msg, src=self.client_addr, - dst=self.region_addr, direction=Direction.OUT) - self.protocol.datagram_received(packet, self.client_addr) + packet = self._msg_to_packet(msg, src=self.client_addr, dst=self.region_addr) + self.protocol.handle_proxied_packet(packet) async def test_simple_command_setting_params(self): self._setup_default_circuit() diff --git a/tests/proxy/integration/test_http.py b/tests/proxy/integration/test_http.py index e817d35..06e4412 100644 --- a/tests/proxy/integration/test_http.py +++ b/tests/proxy/integration/test_http.py @@ -20,8 +20,7 @@ from hippolyzer.lib.proxy.http_flow import HippoHTTPFlow from hippolyzer.lib.proxy.http_proxy import SerializedCapData from hippolyzer.lib.proxy.message_logger import FilteringMessageLogger from hippolyzer.lib.proxy.sessions import SessionManager - -from .. import BaseProxyTest +from hippolyzer.lib.proxy.test_utils import BaseProxyTest class MockAddon(BaseAddon): diff --git a/tests/proxy/integration/test_lludp.py b/tests/proxy/integration/test_lludp.py index 55d3e0b..e0694d7 100644 --- a/tests/proxy/integration/test_lludp.py +++ b/tests/proxy/integration/test_lludp.py @@ -19,8 +19,7 @@ from hippolyzer.lib.proxy.message_logger import FilteringMessageLogger, LLUDPMes from hippolyzer.lib.base.network.transport import Direction from hippolyzer.lib.proxy.region import ProxiedRegion from hippolyzer.lib.proxy.sessions import Session - -from .. import BaseProxyTest +from hippolyzer.lib.proxy.test_utils import BaseProxyTest class MockAddon(BaseAddon): @@ -98,7 +97,7 @@ class LLUDPIntegrationTests(BaseProxyTest): packet_id=1, ) datagram = self._msg_to_datagram(msg, self.client_addr, self.region_addr, - Direction.OUT, socks_header=True) + socks_header=True) self.protocol.datagram_received(datagram, self.client_addr) await self._wait_drained() self.assertFalse(self.session.pending) @@ -116,7 +115,7 @@ class LLUDPIntegrationTests(BaseProxyTest): packet_id=1, ) datagram = self._msg_to_datagram(msg, self.client_addr, self.region_addr, - Direction.OUT, socks_header=True) + socks_header=True) self.protocol.datagram_received(datagram, source_addr=self.client_addr) await self._wait_drained() # Packet got dropped completely @@ -133,7 +132,7 @@ class LLUDPIntegrationTests(BaseProxyTest): packet_id=1, ) datagram = self._msg_to_datagram(msg, self.client_addr, (self.region_addr[0], 9), - Direction.OUT, socks_header=True) + socks_header=True) self.protocol.datagram_received(datagram, source_addr=self.client_addr) await self._wait_drained() # The session claim will still work diff --git a/tests/proxy/test_httpflows.py b/tests/proxy/test_httpflows.py index 197599a..b44ea61 100644 --- a/tests/proxy/test_httpflows.py +++ b/tests/proxy/test_httpflows.py @@ -2,8 +2,7 @@ from mitmproxy.test import tflow, tutils from hippolyzer.lib.proxy.http_flow import HippoHTTPFlow from hippolyzer.lib.proxy.message_logger import HTTPMessageLogEntry - -from . import BaseProxyTest +from hippolyzer.lib.proxy.test_utils import BaseProxyTest class TestHTTPFlows(BaseProxyTest): diff --git a/tests/proxy/test_object_manager.py b/tests/proxy/test_object_manager.py index be10f25..e777b20 100644 --- a/tests/proxy/test_object_manager.py +++ b/tests/proxy/test_object_manager.py @@ -15,9 +15,7 @@ from hippolyzer.lib.proxy.addons import AddonManager from hippolyzer.lib.proxy.addon_utils import BaseAddon from hippolyzer.lib.proxy.region import ProxiedRegion from hippolyzer.lib.proxy.vocache import RegionViewerObjectCacheChain, RegionViewerObjectCache, ViewerObjectCacheEntry - -from . import BaseProxyTest - +from hippolyzer.lib.proxy.test_utils import BaseProxyTest OBJECT_UPDATE_COMPRESSED_DATA = ( b"\x12\x12\x10\xbf\x16XB~\x8f\xb4\xfb\x00\x1a\xcd\x9b\xe5\xd2\x04\x00\x00\t\x00\xcdG\x00\x00"