26063a0e8f
Extract shared DNS wire-format helpers into src/derp/dns.py so both the UDP plugin (dns.py) and the new TCP plugin (tdns.py) share the same encode/decode/build/parse logic. The !tdns command routes queries through the SOCKS5 proxy via derp.http.open_connection, using TCP framing (2-byte length prefix). Default server: 1.1.1.1. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
129 lines
4.1 KiB
Python
129 lines
4.1 KiB
Python
"""Tests for the TCP DNS plugin."""
|
|
|
|
import asyncio
|
|
import importlib.util
|
|
import struct
|
|
import sys
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from derp.dns import encode_name
|
|
|
|
# plugins/ is not a Python package -- load the module from file path
|
|
_spec = importlib.util.spec_from_file_location(
|
|
"plugins.tdns", Path(__file__).resolve().parent.parent / "plugins" / "tdns.py",
|
|
)
|
|
_mod = importlib.util.module_from_spec(_spec)
|
|
sys.modules[_spec.name] = _mod
|
|
_spec.loader.exec_module(_mod)
|
|
|
|
from plugins.tdns import _query_tcp, cmd_tdns # noqa: E402
|
|
|
|
|
|
def _make_a_response(ip_bytes: bytes = b"\x01\x02\x03\x04") -> bytes:
|
|
"""Build a minimal A-record DNS response."""
|
|
tid = b"\x00\x01"
|
|
flags = struct.pack("!H", 0x8180)
|
|
counts = struct.pack("!HHHH", 1, 1, 0, 0)
|
|
qname = encode_name("example.com")
|
|
question = qname + struct.pack("!HH", 1, 1)
|
|
answer = qname + struct.pack("!HHIH", 1, 1, 300, len(ip_bytes)) + ip_bytes
|
|
return tid + flags + counts + question + answer
|
|
|
|
|
|
class TestQueryTcp:
|
|
def test_sends_length_prefixed_query(self):
|
|
response = _make_a_response()
|
|
framed = struct.pack("!H", len(response)) + response
|
|
|
|
reader = AsyncMock()
|
|
reader.readexactly = AsyncMock(side_effect=[
|
|
framed[:2], # length prefix
|
|
framed[2:], # payload
|
|
])
|
|
writer = MagicMock()
|
|
writer.drain = AsyncMock()
|
|
writer.wait_closed = AsyncMock()
|
|
|
|
mock_open = AsyncMock(return_value=(reader, writer))
|
|
with patch.object(_mod, "_open_connection", mock_open):
|
|
rcode, results = asyncio.run(_query_tcp("example.com", 1, "1.1.1.1"))
|
|
|
|
assert rcode == 0
|
|
assert results == ["1.2.3.4"]
|
|
|
|
# Verify the written data has a 2-byte length prefix
|
|
written = writer.write.call_args[0][0]
|
|
pkt_len = struct.unpack("!H", written[:2])[0]
|
|
assert pkt_len == len(written) - 2
|
|
|
|
def test_closes_writer_on_error(self):
|
|
reader = AsyncMock()
|
|
reader.readexactly = AsyncMock(side_effect=asyncio.IncompleteReadError(b"", 2))
|
|
writer = MagicMock()
|
|
writer.drain = AsyncMock()
|
|
writer.wait_closed = AsyncMock()
|
|
|
|
mock_open = AsyncMock(return_value=(reader, writer))
|
|
with patch.object(_mod, "_open_connection", mock_open):
|
|
with pytest.raises(asyncio.IncompleteReadError):
|
|
asyncio.run(_query_tcp("example.com", 1, "1.1.1.1"))
|
|
|
|
writer.close.assert_called_once()
|
|
|
|
|
|
class TestCmdTdns:
|
|
def test_no_args(self):
|
|
bot = AsyncMock()
|
|
msg = MagicMock()
|
|
msg.text = "!tdns"
|
|
asyncio.run(cmd_tdns(bot, msg))
|
|
bot.reply.assert_called_once()
|
|
assert "Usage" in bot.reply.call_args[0][1]
|
|
|
|
def test_ip_auto_ptr(self):
|
|
bot = AsyncMock()
|
|
msg = MagicMock()
|
|
msg.text = "!tdns 1.2.3.4"
|
|
|
|
mock_tcp = AsyncMock(return_value=(0, ["host.example.com"]))
|
|
with patch.object(_mod, "_query_tcp", mock_tcp):
|
|
asyncio.run(cmd_tdns(bot, msg))
|
|
|
|
reply = bot.reply.call_args[0][1]
|
|
assert "PTR" in reply
|
|
assert "host.example.com" in reply
|
|
|
|
def test_explicit_type(self):
|
|
bot = AsyncMock()
|
|
msg = MagicMock()
|
|
msg.text = "!tdns example.com MX"
|
|
|
|
mock_tcp = AsyncMock(return_value=(0, ["10 mail.example.com"]))
|
|
with patch.object(_mod, "_query_tcp", mock_tcp):
|
|
asyncio.run(cmd_tdns(bot, msg))
|
|
|
|
reply = bot.reply.call_args[0][1]
|
|
assert "MX" in reply
|
|
|
|
def test_custom_server(self):
|
|
bot = AsyncMock()
|
|
msg = MagicMock()
|
|
msg.text = "!tdns example.com A @8.8.8.8"
|
|
|
|
mock_tcp = AsyncMock(return_value=(0, ["93.184.216.34"]))
|
|
with patch.object(_mod, "_query_tcp", mock_tcp):
|
|
asyncio.run(cmd_tdns(bot, msg))
|
|
|
|
assert mock_tcp.call_args[0][2] == "8.8.8.8"
|
|
|
|
def test_unknown_type(self):
|
|
bot = AsyncMock()
|
|
msg = MagicMock()
|
|
msg.text = "!tdns example.com BOGUS"
|
|
asyncio.run(cmd_tdns(bot, msg))
|
|
reply = bot.reply.call_args[0][1]
|
|
assert "Unknown type" in reply
|