Export proxy test utils for use in addon test suites

This commit is contained in:
Salad Dais
2021-06-15 18:46:51 +00:00
parent 77d3bf2fe1
commit 48180b85d1
9 changed files with 93 additions and 99 deletions

View File

@@ -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:

View File

@@ -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):

View 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))

View File

@@ -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)

View File

@@ -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()

View File

@@ -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):

View File

@@ -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

View File

@@ -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):

View File

@@ -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"