362 lines
12 KiB
Python
362 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
from collections.abc import Awaitable, Callable
|
|
|
|
import pytest
|
|
|
|
from app.fleet.proxy import ProxyEndpoint, parse_proxy_url, probe_proxy
|
|
|
|
ServerHandler = Callable[[asyncio.StreamReader, asyncio.StreamWriter], Awaitable[None]]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("value", "expected"),
|
|
[
|
|
(
|
|
"socks5://proxy-user:proxy-password@proxy.example.com:10808",
|
|
ProxyEndpoint(
|
|
protocol="socks5",
|
|
host="proxy.example.com",
|
|
port=10808,
|
|
username="proxy-user",
|
|
password="proxy-password",
|
|
),
|
|
),
|
|
(
|
|
"http://proxy-user:proxy-password@proxy.example.com:10808",
|
|
ProxyEndpoint(
|
|
protocol="http",
|
|
host="proxy.example.com",
|
|
port=10808,
|
|
username="proxy-user",
|
|
password="proxy-password",
|
|
),
|
|
),
|
|
(
|
|
"socks5://proxy.example.com:10808@proxy-user:proxy-password",
|
|
ProxyEndpoint(
|
|
protocol="socks5",
|
|
host="proxy.example.com",
|
|
port=10808,
|
|
username="proxy-user",
|
|
password="proxy-password",
|
|
),
|
|
),
|
|
(
|
|
"socks5://proxy.example.com:1080@user:1234",
|
|
ProxyEndpoint(
|
|
protocol="socks5",
|
|
host="proxy.example.com",
|
|
port=1080,
|
|
username="user",
|
|
password="1234",
|
|
),
|
|
),
|
|
(
|
|
"socks5://1.2.3.4:1234@host.example.com:5678",
|
|
ProxyEndpoint(
|
|
protocol="socks5",
|
|
host="1.2.3.4",
|
|
port=1234,
|
|
username="host.example.com",
|
|
password="5678",
|
|
),
|
|
),
|
|
(
|
|
"SOCKS5://user%40team:p%3Aa%2Fss@PROXY.EXAMPLE.COM:1080",
|
|
ProxyEndpoint(
|
|
protocol="socks5",
|
|
host="proxy.example.com",
|
|
port=1080,
|
|
username="user@team",
|
|
password="p:a/ss",
|
|
),
|
|
),
|
|
(
|
|
"http://proxy.example.com:8080",
|
|
ProxyEndpoint(
|
|
protocol="http",
|
|
host="proxy.example.com",
|
|
port=8080,
|
|
username=None,
|
|
password=None,
|
|
),
|
|
),
|
|
],
|
|
)
|
|
def test_parse_proxy_url_supported_formats(value: str, expected: ProxyEndpoint) -> None:
|
|
assert parse_proxy_url(value) == expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"value",
|
|
[
|
|
"",
|
|
"ftp://user:password@proxy.example.com:1080",
|
|
"socks5://user:password@proxy.example.com",
|
|
"socks5://user:password@proxy.example.com:0",
|
|
"socks5://user:password@proxy.example.com:65536",
|
|
"socks5://user:password@proxy.example.com:not-a-port",
|
|
"socks5://user:password@proxy.example.com:1080/path",
|
|
"socks5://user:password@proxy.example.com:1080?mode=test",
|
|
"socks5://user:password@proxy.example.com:1080#fragment",
|
|
"socks5://user:password@proxy.example.com:1080\n",
|
|
"socks5://user:pass%0Dword@proxy.example.com:1080",
|
|
"socks5://user:pass%ZZword@proxy.example.com:1080",
|
|
"socks5://user@proxy.example.com:1080",
|
|
"socks5://user:@proxy.example.com:1080",
|
|
"socks5://user:password@proxy.example.com:1080@extra",
|
|
],
|
|
)
|
|
def test_parse_proxy_url_rejects_invalid_values(value: str) -> None:
|
|
with pytest.raises(ValueError):
|
|
parse_proxy_url(value)
|
|
|
|
|
|
def test_parse_errors_do_not_echo_credentials() -> None:
|
|
username = "sensitive-user"
|
|
password = "sensitive-password"
|
|
with pytest.raises(ValueError) as captured:
|
|
parse_proxy_url(f"socks5://{username}:{password}@proxy.example.com:70000")
|
|
message = str(captured.value)
|
|
assert username not in message
|
|
assert password not in message
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_probe_socks5_performs_username_password_auth_and_connect() -> None:
|
|
observed: asyncio.Future[tuple[str, str, bytes]] = asyncio.get_running_loop().create_future()
|
|
|
|
async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
|
try:
|
|
greeting = await reader.readexactly(3)
|
|
writer.write(b"\x05\x02")
|
|
await writer.drain()
|
|
auth_version, username_length = await reader.readexactly(2)
|
|
username = (await reader.readexactly(username_length)).decode()
|
|
password_length = (await reader.readexactly(1))[0]
|
|
password = (await reader.readexactly(password_length)).decode()
|
|
writer.write(b"\x01\x00")
|
|
await writer.drain()
|
|
connect_request = await reader.readexactly(10)
|
|
writer.write(b"\x05\x00\x00\x01\x00\x00\x00\x00\x00\x00")
|
|
await writer.drain()
|
|
observed.set_result(
|
|
(username, password, greeting + bytes((auth_version,)) + connect_request)
|
|
)
|
|
except Exception as exc: # pragma: no cover - surfaced through the future
|
|
if not observed.done():
|
|
observed.set_exception(exc)
|
|
finally:
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
|
|
server, port = await _start_server(handler)
|
|
try:
|
|
await probe_proxy(
|
|
"127.0.0.1",
|
|
port,
|
|
"socks5",
|
|
username="proxy-user",
|
|
password="proxy-password",
|
|
)
|
|
username, password, protocol_bytes = await observed
|
|
finally:
|
|
await _close_server(server)
|
|
|
|
assert username == "proxy-user"
|
|
assert password == "proxy-password"
|
|
assert protocol_bytes[:3] == b"\x05\x01\x02"
|
|
assert protocol_bytes[3] == 1
|
|
assert protocol_bytes[4:] == b"\x05\x01\x00\x01\x01\x01\x01\x01\x01\xbb"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_probe_socks5_supports_anonymous_method() -> None:
|
|
observed: asyncio.Future[bytes] = asyncio.get_running_loop().create_future()
|
|
|
|
async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
|
try:
|
|
greeting = await reader.readexactly(3)
|
|
writer.write(b"\x05\x00")
|
|
await writer.drain()
|
|
connect_request = await reader.readexactly(10)
|
|
writer.write(b"\x05\x00\x00\x03\x03foo\x00\x50")
|
|
await writer.drain()
|
|
observed.set_result(greeting + connect_request)
|
|
finally:
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
|
|
server, port = await _start_server(handler)
|
|
try:
|
|
await probe_proxy("127.0.0.1", port, "socks5")
|
|
protocol_bytes = await observed
|
|
finally:
|
|
await _close_server(server)
|
|
|
|
assert protocol_bytes[:3] == b"\x05\x01\x00"
|
|
assert protocol_bytes[3:] == b"\x05\x01\x00\x01\x01\x01\x01\x01\x01\xbb"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("method_reply", "auth_reply", "connect_reply"),
|
|
[
|
|
(b"\x05\x00", None, None),
|
|
(b"\x05\x02", b"\x01\x01", None),
|
|
(
|
|
b"\x05\x02",
|
|
b"\x01\x00",
|
|
b"\x05\x05\x00\x01\x00\x00\x00\x00\x00\x00",
|
|
),
|
|
],
|
|
)
|
|
async def test_probe_socks5_rejects_protocol_and_auth_failures_without_secret_leak(
|
|
method_reply: bytes,
|
|
auth_reply: bytes | None,
|
|
connect_reply: bytes | None,
|
|
) -> None:
|
|
async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
|
try:
|
|
await reader.readexactly(3)
|
|
writer.write(method_reply)
|
|
await writer.drain()
|
|
if auth_reply is not None:
|
|
_version, username_length = await reader.readexactly(2)
|
|
await reader.readexactly(username_length)
|
|
password_length = (await reader.readexactly(1))[0]
|
|
await reader.readexactly(password_length)
|
|
writer.write(auth_reply)
|
|
await writer.drain()
|
|
if connect_reply is not None:
|
|
await reader.readexactly(10)
|
|
writer.write(connect_reply)
|
|
await writer.drain()
|
|
finally:
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
|
|
server, port = await _start_server(handler)
|
|
try:
|
|
with pytest.raises(OSError) as captured:
|
|
await probe_proxy(
|
|
"127.0.0.1",
|
|
port,
|
|
"socks5",
|
|
username="secret-user",
|
|
password="secret-password",
|
|
)
|
|
finally:
|
|
await _close_server(server)
|
|
|
|
assert "secret-user" not in str(captured.value)
|
|
assert "secret-password" not in str(captured.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_probe_http_sends_connect_and_basic_authorization() -> None:
|
|
observed: asyncio.Future[bytes] = asyncio.get_running_loop().create_future()
|
|
|
|
async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
|
try:
|
|
request = await reader.readuntil(b"\r\n\r\n")
|
|
observed.set_result(request)
|
|
writer.write(b"HTTP/1.1 200 Connection Established\r\n\r\n")
|
|
await writer.drain()
|
|
finally:
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
|
|
server, port = await _start_server(handler)
|
|
try:
|
|
await probe_proxy(
|
|
"127.0.0.1",
|
|
port,
|
|
"http",
|
|
username="http-user",
|
|
password="http-password",
|
|
)
|
|
request = (await observed).decode("ascii")
|
|
finally:
|
|
await _close_server(server)
|
|
|
|
expected_auth = base64.b64encode(b"http-user:http-password").decode("ascii")
|
|
assert request.startswith("CONNECT 1.1.1.1:443 HTTP/1.1\r\n")
|
|
assert "Host: 1.1.1.1:443\r\n" in request
|
|
assert f"Proxy-Authorization: Basic {expected_auth}\r\n" in request
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("status", [407, 500])
|
|
async def test_probe_http_rejects_non_2xx_without_secret_leak(status: int) -> None:
|
|
async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
|
try:
|
|
await reader.readuntil(b"\r\n\r\n")
|
|
writer.write(f"HTTP/1.1 {status} Failure\r\n\r\n".encode("ascii"))
|
|
await writer.drain()
|
|
finally:
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
|
|
server, port = await _start_server(handler)
|
|
try:
|
|
with pytest.raises(OSError) as captured:
|
|
await probe_proxy(
|
|
"127.0.0.1",
|
|
port,
|
|
"http",
|
|
username="secret-http-user",
|
|
password="secret-http-password",
|
|
)
|
|
finally:
|
|
await _close_server(server)
|
|
|
|
assert "secret-http-user" not in str(captured.value)
|
|
assert "secret-http-password" not in str(captured.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_probe_timeout_is_sanitized() -> None:
|
|
async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
|
try:
|
|
await reader.read()
|
|
finally:
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
|
|
server, port = await _start_server(handler)
|
|
try:
|
|
with pytest.raises(OSError, match="代理探测超时"):
|
|
await probe_proxy("127.0.0.1", port, "socks5", timeout_seconds=0.01)
|
|
finally:
|
|
await _close_server(server)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_probe_rejects_invalid_configuration() -> None:
|
|
with pytest.raises(ValueError):
|
|
await probe_proxy("127.0.0.1", 1080, "https")
|
|
with pytest.raises(ValueError):
|
|
await probe_proxy("127.0.0.1", 0, "socks5")
|
|
with pytest.raises(ValueError):
|
|
await probe_proxy(
|
|
"127.0.0.1",
|
|
1080,
|
|
"socks5",
|
|
username="only-username",
|
|
)
|
|
|
|
|
|
async def _start_server(handler: ServerHandler) -> tuple[asyncio.AbstractServer, int]:
|
|
server = await asyncio.start_server(handler, "127.0.0.1", 0)
|
|
socket_address = server.sockets[0].getsockname()
|
|
return server, int(socket_address[1])
|
|
|
|
|
|
async def _close_server(server: asyncio.AbstractServer) -> None:
|
|
server.close()
|
|
await server.wait_closed()
|