Export proxy test utils for use in addon test suites
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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):
|
||||
|
||||
80
hippolyzer/lib/proxy/test_utils.py
Normal file
80
hippolyzer/lib/proxy/test_utils.py
Normal file
@@ -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))
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user