1 files changed, 149 insertions, 0 deletions
diff --git a/test/netlib/test_socks.py b/test/netlib/test_socks.py
new file mode 100644
index 00000000..d95dee41
--- /dev/null
+++ b/test/netlib/test_socks.py
@@ -0,0 +1,149 @@
+import ipaddress
+from io import BytesIO
+import socket
+from netlib import socks, tcp, tutils
+
+
+def test_client_greeting():
+    raw = tutils.treader(b"\x05\x02\x00\xBE\xEF")
+    out = BytesIO()
+    msg = socks.ClientGreeting.from_file(raw)
+    msg.assert_socks5()
+    msg.to_file(out)
+
+    assert out.getvalue() == raw.getvalue()[:-1]
+    assert msg.ver == 5
+    assert len(msg.methods) == 2
+    assert 0xBE in msg.methods
+    assert 0xEF not in msg.methods
+
+
+def test_client_greeting_assert_socks5():
+    raw = tutils.treader(b"\x00\x00")
+    msg = socks.ClientGreeting.from_file(raw)
+    tutils.raises(socks.SocksError, msg.assert_socks5)
+
+    raw = tutils.treader(b"HTTP/1.1 200 OK" + b" " * 100)
+    msg = socks.ClientGreeting.from_file(raw)
+    try:
+        msg.assert_socks5()
+    except socks.SocksError as e:
+        assert "Invalid SOCKS version" in str(e)
+        assert "HTTP" not in str(e)
+    else:
+        assert False
+
+    raw = tutils.treader(b"GET / HTTP/1.1" + b" " * 100)
+    msg = socks.ClientGreeting.from_file(raw)
+    try:
+        msg.assert_socks5()
+    except socks.SocksError as e:
+        assert "Invalid SOCKS version" in str(e)
+        assert "HTTP" in str(e)
+    else:
+        assert False
+
+    raw = tutils.treader(b"XX")
+    tutils.raises(
+        socks.SocksError,
+        socks.ClientGreeting.from_file,
+        raw,
+        fail_early=True)
+
+
+def test_server_greeting():
+    raw = tutils.treader(b"\x05\x02")
+    out = BytesIO()
+    msg = socks.ServerGreeting.from_file(raw)
+    msg.assert_socks5()
+    msg.to_file(out)
+
+    assert out.getvalue() == raw.getvalue()
+    assert msg.ver == 5
+    assert msg.method == 0x02
+
+
+def test_server_greeting_assert_socks5():
+    raw = tutils.treader(b"HTTP/1.1 200 OK" + b" " * 100)
+    msg = socks.ServerGreeting.from_file(raw)
+    try:
+        msg.assert_socks5()
+    except socks.SocksError as e:
+        assert "Invalid SOCKS version" in str(e)
+        assert "HTTP" in str(e)
+    else:
+        assert False
+
+    raw = tutils.treader(b"GET / HTTP/1.1" + b" " * 100)
+    msg = socks.ServerGreeting.from_file(raw)
+    try:
+        msg.assert_socks5()
+    except socks.SocksError as e:
+        assert "Invalid SOCKS version" in str(e)
+        assert "HTTP" not in str(e)
+    else:
+        assert False
+
+
+def test_message():
+    raw = tutils.treader(b"\x05\x01\x00\x03\x0bexample.com\xDE\xAD\xBE\xEF")
+    out = BytesIO()
+    msg = socks.Message.from_file(raw)
+    msg.assert_socks5()
+    assert raw.read(2) == b"\xBE\xEF"
+    msg.to_file(out)
+
+    assert out.getvalue() == raw.getvalue()[:-2]
+    assert msg.ver == 5
+    assert msg.msg == 0x01
+    assert msg.atyp == 0x03
+    assert msg.addr == ("example.com", 0xDEAD)
+
+
+def test_message_assert_socks5():
+    raw = tutils.treader(b"\xEE\x01\x00\x03\x0bexample.com\xDE\xAD\xBE\xEF")
+    msg = socks.Message.from_file(raw)
+    tutils.raises(socks.SocksError, msg.assert_socks5)
+
+
+def test_message_ipv4():
+    # Test ATYP=0x01 (IPV4)
+    raw = tutils.treader(b"\x05\x01\x00\x01\x7f\x00\x00\x01\xDE\xAD\xBE\xEF")
+    out = BytesIO()
+    msg = socks.Message.from_file(raw)
+    left = raw.read(2)
+    assert left == b"\xBE\xEF"
+    msg.to_file(out)
+
+    assert out.getvalue() == raw.getvalue()[:-2]
+    assert msg.addr == ("127.0.0.1", 0xDEAD)
+
+
+def test_message_ipv6():
+    # Test ATYP=0x04 (IPV6)
+    ipv6_addr = u"2001:db8:85a3:8d3:1319:8a2e:370:7344"
+
+    raw = tutils.treader(
+        b"\x05\x01\x00\x04" +
+        ipaddress.IPv6Address(ipv6_addr).packed +
+        b"\xDE\xAD\xBE\xEF")
+    out = BytesIO()
+    msg = socks.Message.from_file(raw)
+    assert raw.read(2) == b"\xBE\xEF"
+    msg.to_file(out)
+
+    assert out.getvalue() == raw.getvalue()[:-2]
+    assert msg.addr.host == ipv6_addr
+
+
+def test_message_invalid_rsv():
+    raw = tutils.treader(b"\x05\x01\xFF\x01\x7f\x00\x00\x01\xDE\xAD\xBE\xEF")
+    tutils.raises(socks.SocksError, socks.Message.from_file, raw)
+
+
+def test_message_unknown_atyp():
+    raw = tutils.treader(b"\x05\x02\x00\x02\x7f\x00\x00\x01\xDE\xAD\xBE\xEF")
+    tutils.raises(socks.SocksError, socks.Message.from_file, raw)
+
+    m = socks.Message(5, 1, 0x02, tcp.Address(("example.com", 5050)))
+    tutils.raises(socks.SocksError, m.to_file, BytesIO())