99 lines
2.8 KiB
Python
99 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import socket
|
|
import threading
|
|
import unittest
|
|
|
|
from bot_bottle.backend.smolmachines.port_forward import (
|
|
ForwardSpec,
|
|
start_forwarder,
|
|
stop_forwarder,
|
|
)
|
|
|
|
|
|
def _echo_server() -> tuple[socket.socket, int]:
|
|
listener = socket.create_server(("127.0.0.1", 0))
|
|
port = int(listener.getsockname()[1])
|
|
|
|
def run() -> None:
|
|
with listener:
|
|
conn, _ = listener.accept()
|
|
with conn:
|
|
while True:
|
|
data = conn.recv(65536)
|
|
if not data:
|
|
return
|
|
conn.sendall(data)
|
|
|
|
threading.Thread(target=run, daemon=True).start()
|
|
return listener, port
|
|
|
|
|
|
class TestForwarderValidation(unittest.TestCase):
|
|
def test_rejects_generic_localhost_listener(self):
|
|
with self.assertRaises(SystemExit):
|
|
start_forwarder((
|
|
ForwardSpec(
|
|
label="egress",
|
|
listen_host="127.0.0.1",
|
|
listen_port=0,
|
|
target_host="127.0.0.1",
|
|
target_port=9099,
|
|
),
|
|
))
|
|
|
|
def test_rejects_wildcard_listener(self):
|
|
with self.assertRaises(SystemExit):
|
|
start_forwarder((
|
|
ForwardSpec(
|
|
label="egress",
|
|
listen_host="0.0.0.0",
|
|
listen_port=0,
|
|
target_host="127.0.0.1",
|
|
target_port=9099,
|
|
),
|
|
))
|
|
|
|
def test_rejects_non_loopback_target(self):
|
|
with self.assertRaises(SystemExit):
|
|
start_forwarder((
|
|
ForwardSpec(
|
|
label="egress",
|
|
listen_host="127.0.0.16",
|
|
listen_port=0,
|
|
target_host="192.0.2.10",
|
|
target_port=9099,
|
|
),
|
|
))
|
|
|
|
|
|
class TestForwarding(unittest.TestCase):
|
|
def test_forwards_bytes_both_directions(self):
|
|
try:
|
|
probe = socket.create_server(("127.0.0.16", 0))
|
|
except OSError:
|
|
self.skipTest("127.0.0.16 is not bound on this host")
|
|
else:
|
|
probe.close()
|
|
_listener, target_port = _echo_server()
|
|
handle = start_forwarder((
|
|
ForwardSpec(
|
|
label="egress",
|
|
listen_host="127.0.0.16",
|
|
listen_port=0,
|
|
target_host="127.0.0.1",
|
|
target_port=target_port,
|
|
),
|
|
))
|
|
try:
|
|
bound = handle.forwards[0]
|
|
with socket.create_connection((bound.listen_host, bound.listen_port)) as sock:
|
|
sock.sendall(b"hello")
|
|
self.assertEqual(b"hello", sock.recv(5))
|
|
finally:
|
|
stop_forwarder(handle)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|