153 lines
5.8 KiB
Python
153 lines
5.8 KiB
Python
"""轻量 DNS 观察器，用于把 TUN 里的目标 IP 关联回解析域名。
|
|
|
|
TUN 送到 SOCKS5 的请求通常只带目标 IP。这里观察经 UoT 或直连的
|
|
UDP DNS 回包，保存短期的 ``IP -> 域名`` 映射供分流和日志使用。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ipaddress
|
|
import struct
|
|
import threading
|
|
import time
|
|
from typing import Callable
|
|
|
|
|
|
def _read_name(packet: bytes, offset: int) -> tuple[str, int]:
|
|
labels: list[str] = []
|
|
cursor = offset
|
|
next_offset = None
|
|
jumps = 0
|
|
while True:
|
|
if cursor >= len(packet):
|
|
raise ValueError("DNS 名称超出数据包")
|
|
length = packet[cursor]
|
|
if length == 0:
|
|
cursor += 1
|
|
if next_offset is None:
|
|
next_offset = cursor
|
|
break
|
|
if length & 0xC0 == 0xC0:
|
|
if cursor + 1 >= len(packet):
|
|
raise ValueError("DNS 压缩指针不完整")
|
|
pointer = ((length & 0x3F) << 8) | packet[cursor + 1]
|
|
if next_offset is None:
|
|
next_offset = cursor + 2
|
|
cursor = pointer
|
|
jumps += 1
|
|
if jumps > 16:
|
|
raise ValueError("DNS 名称压缩指针循环")
|
|
continue
|
|
if length > 63 or cursor + 1 + length > len(packet):
|
|
raise ValueError("DNS 标签不完整")
|
|
labels.append(packet[cursor + 1:cursor + 1 + length].decode("idna"))
|
|
cursor += 1 + length
|
|
return ".".join(labels).rstrip(".").lower(), int(next_offset)
|
|
|
|
|
|
def _question(packet: bytes) -> tuple[int, str, int]:
|
|
if len(packet) < 12:
|
|
raise ValueError("DNS 头不完整")
|
|
txid, flags, questions, _answers, _authority, _additional = struct.unpack(
|
|
"!HHHHHH", packet[:12]
|
|
)
|
|
if flags & 0x8000 or questions < 1:
|
|
raise ValueError("不是 DNS 查询")
|
|
name, offset = _read_name(packet, 12)
|
|
if offset + 4 > len(packet):
|
|
raise ValueError("DNS 问题段不完整")
|
|
return txid, name, offset + 4
|
|
|
|
|
|
class DnsObserver:
|
|
def __init__(self, log: Callable[[str], None] | None = None,
|
|
max_ttl: int = 300):
|
|
self.log = log or (lambda _message: None)
|
|
self.max_ttl = max(1, int(max_ttl))
|
|
self._lock = threading.RLock()
|
|
self._domain_to_ips: dict[str, tuple[float, tuple[str, ...]]] = {}
|
|
# 一个 CDN IP 可能同时服务多个域名；不能把最后一次 DNS 回包当成唯一身份。
|
|
self._ip_to_domains: dict[str, tuple[float, set[str]]] = {}
|
|
|
|
def observe_query(self, packet: bytes) -> str | None:
|
|
try:
|
|
_txid, domain, _ = _question(packet)
|
|
return domain or None
|
|
except (ValueError, UnicodeError):
|
|
return None
|
|
|
|
def observe_response(self, packet: bytes) -> tuple[str, tuple[str, ...]] | None:
|
|
try:
|
|
if len(packet) < 12:
|
|
raise ValueError("DNS 头不完整")
|
|
_txid, flags, questions, answers, _authority, _additional = struct.unpack(
|
|
"!HHHHHH", packet[:12]
|
|
)
|
|
if not (flags & 0x8000) or questions < 1:
|
|
raise ValueError("不是 DNS 回包")
|
|
domain, offset = _read_name(packet, 12)
|
|
if offset + 4 > len(packet):
|
|
raise ValueError("DNS 问题段不完整")
|
|
offset += 4
|
|
ips: list[str] = []
|
|
ttl = self.max_ttl
|
|
for _ in range(answers):
|
|
_owner, offset = _read_name(packet, offset)
|
|
if offset + 10 > len(packet):
|
|
raise ValueError("DNS 答案头不完整")
|
|
rtype, rclass, record_ttl, size = struct.unpack(
|
|
"!HHIH", packet[offset:offset + 10]
|
|
)
|
|
offset += 10
|
|
if offset + size > len(packet):
|
|
raise ValueError("DNS 答案数据不完整")
|
|
raw = packet[offset:offset + size]
|
|
offset += size
|
|
if rclass != 1 or rtype not in (1, 28):
|
|
continue
|
|
if (rtype == 1 and size != 4) or (rtype == 28 and size != 16):
|
|
continue
|
|
ips.append(str(ipaddress.ip_address(raw)))
|
|
ttl = min(ttl, max(1, int(record_ttl)))
|
|
if not domain or not ips:
|
|
return None
|
|
expires = time.monotonic() + min(ttl, self.max_ttl)
|
|
unique = tuple(dict.fromkeys(ips))
|
|
with self._lock:
|
|
self._domain_to_ips[domain] = (expires, unique)
|
|
for ip in unique:
|
|
previous = self._ip_to_domains.get(ip)
|
|
domains = set(previous[1]) if previous and previous[0] > time.monotonic() else set()
|
|
domains.add(domain)
|
|
self._ip_to_domains[ip] = (expires, domains)
|
|
return domain, unique
|
|
except (ValueError, UnicodeError, OSError):
|
|
return None
|
|
|
|
def lookup_ip(self, ip: str) -> str | None:
|
|
try:
|
|
normalized = str(ipaddress.ip_address(ip))
|
|
except ValueError:
|
|
return None
|
|
now = time.monotonic()
|
|
with self._lock:
|
|
item = self._ip_to_domains.get(normalized)
|
|
if not item:
|
|
return None
|
|
if item[0] <= now:
|
|
self._ip_to_domains.pop(normalized, None)
|
|
return None
|
|
return next(iter(item[1])) if len(item[1]) == 1 else None
|
|
|
|
def lookup_domain(self, domain: str) -> tuple[str, ...]:
|
|
key = str(domain).rstrip(".").lower()
|
|
now = time.monotonic()
|
|
with self._lock:
|
|
item = self._domain_to_ips.get(key)
|
|
if not item:
|
|
return ()
|
|
if item[0] <= now:
|
|
self._domain_to_ips.pop(key, None)
|
|
return ()
|
|
return item[1]
|