1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151
|
import pytest
from socksio import (
ProtocolError,
SOCKS4ARequest,
SOCKS4Command,
SOCKS4Connection,
SOCKS4Reply,
SOCKS4ReplyCode,
SOCKS4Request,
SOCKSError,
)
@pytest.mark.parametrize(
"address,expected_address,expected_port",
[
(("127.0.0.1", 3080), b"\x7f\x00\x00\x01", 3080),
(("127.0.0.1", "3080"), b"\x7f\x00\x00\x01", 3080),
("127.0.0.1:8080", b"\x7f\x00\x00\x01", 8080),
((b"127.0.0.1", 3080), b"\x7f\x00\x00\x01", 3080),
((b"127.0.0.1", b"3080"), b"\x7f\x00\x00\x01", 3080),
(b"127.0.0.1:8080", b"\x7f\x00\x00\x01", 8080),
],
)
def test_socks4request_from_address(address, expected_address, expected_port) -> None:
req = SOCKS4Request.from_address(SOCKS4Command.CONNECT, address, user_id=b"socksio")
assert req.command == SOCKS4Command.CONNECT
assert req.addr == expected_address
assert req.port == expected_port
assert req.user_id == b"socksio"
@pytest.mark.parametrize(
"address,user_id",
[
(("::1", 3080), b"socksio"), # IPV6
("localhost:3080", b"socksio"), # Domain names
],
)
def test_socks4request_from_address_errors(address, user_id) -> None:
with pytest.raises(SOCKSError):
SOCKS4Request.from_address(
command=SOCKS4Command.BIND, address=address, user_id=user_id
)
def test_socks4request_from_address_dump_raises_if_no_user_id():
req = SOCKS4Request.from_address(SOCKS4Command.CONNECT, "127.0.0.1:8080")
with pytest.raises(SOCKSError):
req.dumps()
@pytest.mark.parametrize(
"address,expected_address,expected_port",
[
(("127.0.0.1", 3080), b"\x7f\x00\x00\x01", 3080),
(("127.0.0.1", "3080"), b"\x7f\x00\x00\x01", 3080),
("127.0.0.1:8080", b"\x7f\x00\x00\x01", 8080),
((b"127.0.0.1", 3080), b"\x7f\x00\x00\x01", 3080),
((b"127.0.0.1", b"3080"), b"\x7f\x00\x00\x01", 3080),
(b"127.0.0.1:8080", b"\x7f\x00\x00\x01", 8080),
],
)
def test_socks4arequest_from_address(address, expected_address, expected_port) -> None:
req = SOCKS4ARequest.from_address(
SOCKS4Command.CONNECT, address, user_id=b"socksio"
)
assert req.command == SOCKS4Command.CONNECT
assert req.addr == expected_address
assert req.port == expected_port
assert req.user_id == b"socksio"
def test_socks4arequest_from_address_dump_raises_if_no_user_id():
req = SOCKS4ARequest.from_address(SOCKS4Command.CONNECT, "127.0.0.1:8080")
with pytest.raises(SOCKSError):
req.dumps()
@pytest.mark.parametrize("command", [SOCKS4Command.BIND, SOCKS4Command.CONNECT])
def test_socks4_connection_request(command: SOCKS4Command) -> None:
conn = SOCKS4Connection(user_id=b"socks")
request = SOCKS4Request.from_address(command=command, address=("127.0.0.1", 8080))
conn.send(request)
data = conn.data_to_send()
assert len(data) == 9 + 5
assert data[0:1] == b"\x04"
assert data[1:2] == command
assert data[2:4] == (8080).to_bytes(2, byteorder="big")
assert data[4:8] == b"\x7f\x00\x00\x01"
assert data[8:13] == b"socks"
assert data[13] == 0
@pytest.mark.parametrize("request_reply_code", list(SOCKS4ReplyCode))
def test_socks4_receive_data(request_reply_code: bytes) -> None:
conn = SOCKS4Connection(user_id=b"socks")
reply = conn.receive_data(
b"".join(
[
b"\x00",
request_reply_code,
(8080).to_bytes(2, byteorder="big"),
b"\x7f\x00\x00\x01",
]
)
)
assert reply == SOCKS4Reply(
reply_code=SOCKS4ReplyCode(request_reply_code), port=8080, addr="127.0.0.1"
)
@pytest.mark.parametrize(
"data",
[
b"\x00Z\x1f\x90\x7f\x00\x00", # missing one byte
b"\x0FZ\x1f\x90\x7f\x00\x00\x01", # not starting with 0
b"\x00\xFF\x1f\x90\x7f\x00\x00\x01", # incorrect reply code
],
)
def test_socks4_receive_malformed_data(data: bytes) -> None:
conn = SOCKS4Connection(user_id=b"socks")
with pytest.raises(ProtocolError):
conn.receive_data(data)
@pytest.mark.parametrize("command", [SOCKS4Command.BIND, SOCKS4Command.CONNECT])
def test_SOCKS4A_connection_request(command: SOCKS4Command) -> None:
conn = SOCKS4Connection(user_id=b"socks")
request = SOCKS4ARequest.from_address(
command=command, address=("proxy.example.com", 8080)
)
conn.send(request)
data = conn.data_to_send()
assert len(data) == 32
assert data[0:1] == b"\x04"
assert data[1:2] == command
assert data[2:4] == (8080).to_bytes(2, byteorder="big")
assert data[4:8] == b"\x00\x00\x00\xFF"
assert data[8:14] == b"socks\x00"
assert data[14:] == b"proxy.example.com\x00"
|