vpn-loopback-firewall/tests/test_rewrite.py

"""rewrite_syn 单测: Linux SYN -> Windows 布局, 反向 Windows 包不破坏。"""
from __future__ import annotations

import struct
import unittest

from masquerader.rewrite import Fingerprint, rewrite_syn

FP = Fingerprint()

TCP_OPT_EOL = 0
TCP_OPT_NOP = 1
TCP_OPT_MSS = 2
TCP_OPT_WS = 3
TCP_OPT_SACK = 4
TCP_OPT_TS = 8


def _opt_mss(value: int = 1460) -> bytes:
    return bytes((TCP_OPT_MSS, 4)) + struct.pack("!H", value)


def _opt_sack() -> bytes:
    return bytes((TCP_OPT_SACK, 2))


def _opt_ts(tsval: int = 1000, tsecr: int = 0) -> bytes:
    return struct.pack("!BBII", TCP_OPT_TS, 10, tsval, tsecr)


def _opt_nop() -> bytes:
    return bytes((TCP_OPT_NOP,))


def _opt_ws(scale: int) -> bytes:
    return bytes((TCP_OPT_WS, 3, scale))


def _checksum_fold(data: bytes) -> int:
    if len(data) & 1:
        data += b"\x00"
    total = sum(struct.unpack(f"!{len(data) // 2}H", data))
    total = (total >> 16) + (total & 0xFFFF)
    total += total >> 16
    return total


def _verify_ip(pkt: bytes) -> bool:
    ihl = (pkt[0] & 0x0F) * 4
    return _checksum_fold(pkt[:ihl]) == 0xFFFF


def _verify_tcp(pkt: bytes) -> bool:
    ihl = (pkt[0] & 0x0F) * 4
    tcp = pkt[ihl:]
    off = (tcp[12] >> 4) * 4
    pseudo = (
        pkt[12:16] + pkt[16:20] + struct.pack("!BBH", 0, 6, off)
    )
    return _checksum_fold(pseudo + tcp[:off]) == 0xFFFF


def _build_syn(
    options: bytes,
    window: int = 29200,
    ttl: int = 64,
    flags: int = 0x02,
) -> bytes:
    opt_pad = options + b"\x00" * ((4 - len(options) % 4) % 4)
    tcp_len = 20 + len(opt_pad)
    data_offset = tcp_len // 4
    tcp = bytearray(20 + len(opt_pad))
    struct.pack_into(
        "!HHIIBBHHH", tcp, 0,
        12345,               # sport
        443,                 # dport
        1,                   # seq
        0,                   # ack
        data_offset << 4,    # byte12: data offset + reserved
        flags,               # byte13: tcp flags
        window,              # byte14-15: window
        0,                   # byte16-17: checksum, 下方填入
        0,                   # byte18-19: urgent ptr
    )
    tcp[20:] = opt_pad
    src = bytes([10, 0, 8, 47])
    dst = bytes([1, 2, 3, 4])
    cs = _checksum_fold(src + dst + struct.pack("!BBH", 0, 6, len(tcp)) + bytes(tcp))
    tcp[16:18] = struct.pack("!H", (~cs) & 0xFFFF)
    ip_len = 20 + len(tcp)
    ip = bytearray(20)
    struct.pack_into("!BBHHHBBH4s4s", ip, 0, 0x45, 0, ip_len, 0x1234, 0, ttl, 6, 0, src, dst)
    ip_cs = _checksum_fold(bytes(ip))
    ip[10:12] = struct.pack("!H", (~ip_cs) & 0xFFFF)
    return bytes(ip) + bytes(tcp)


def _linux_options() -> bytes:
    return _opt_mss(1460) + _opt_sack() + _opt_ts(0, 0) + _opt_nop() + _opt_ws(7)


def _windows_options() -> bytes:
    return _opt_mss(1460) + _opt_sack() + _opt_ts(0, 0) + _opt_nop() + _opt_ws(8)


def _parse_opts(pkt: bytes) -> list[int]:
    ihl = (pkt[0] & 0x0F) * 4
    tcp = pkt[ihl:]
    off = (tcp[12] >> 4) * 4
    kinds: list[int] = []
    i = 20
    while i < off:
        kind = tcp[i]
        if kind == TCP_OPT_EOL:
            break
        if kind == TCP_OPT_NOP:
            kinds.append(kind)
            i += 1
            continue
        length = tcp[i + 1]
        kinds.append(kind)
        i += length
    return kinds


def _tcp_field(pkt: bytes, offset: int, size: int) -> int:
    ihl = (pkt[0] & 0x0F) * 4
    return int.from_bytes(pkt[ihl + offset : ihl + offset + size], "big")


class RewriteSynTest(unittest.TestCase):
    def test_linux_syn_rewritten_to_windows(self) -> None:
        pkt = _build_syn(_linux_options(), window=29200, ttl=64)
        out = rewrite_syn(pkt, FP)
        assert out is not None
        self.assertEqual(_tcp_field(out, 14, 2), 64240)  # window
        self.assertEqual(_tcp_field(out, 12, 1) >> 4, (20 + 12) // 4)  # data offset
        kinds = _parse_opts(out)
        self.assertNotIn(TCP_OPT_TS, kinds)          # TS 已剥
        self.assertIn(TCP_OPT_WS, kinds)
        self.assertTrue(_verify_ip(out))              # IP checksum 有效
        self.assertTrue(_verify_tcp(out))             # TCP checksum 有效

    def test_wscale_value_is_8(self) -> None:
        pkt = _build_syn(_linux_options(), window=29200)
        out = rewrite_syn(pkt, FP)
        assert out is not None
        ihl = (out[0] & 0x0F) * 4
        tcp = out[ihl:]
        off = (tcp[12] >> 4) * 4
        i = 20
        ws_val: int | None = None
        while i < off:
            kind = tcp[i]
            if kind == TCP_OPT_EOL:
                break
            if kind == TCP_OPT_NOP:
                i += 1
                continue
            length = tcp[i + 1]
            if kind == TCP_OPT_WS:
                ws_val = tcp[i + 2]
                break
            i += length
        self.assertEqual(ws_val, 8)

    def test_ip_total_length_shrinks_by_ts_size(self) -> None:
        pkt = _build_syn(_linux_options(), window=29200)
        out = rewrite_syn(pkt, FP)
        assert out is not None
        orig_len = int.from_bytes(pkt[2:4], "big")
        new_len = int.from_bytes(out[2:4], "big")
        # TS(10) 剥掉 + 补 pad 后 options 从 20 -> 12, 总长减 8
        self.assertEqual(orig_len - new_len, 8)

    def test_windows_syn_only_strips_ts(self) -> None:
        pkt = _build_syn(_windows_options(), window=64240)
        out = rewrite_syn(pkt, FP)
        assert out is not None
        self.assertEqual(_tcp_field(out, 14, 2), 64240)  # window 不动
        kinds = _parse_opts(out)
        self.assertNotIn(TCP_OPT_TS, kinds)
        self.assertTrue(_verify_ip(out))
        self.assertTrue(_verify_tcp(out))

    def test_non_syn_packet_untouched(self) -> None:
        pkt = _build_syn(_linux_options(), flags=0x18)  # PSH+ACK
        self.assertIsNone(rewrite_syn(pkt, FP))

    def test_syn_ack_untouched(self) -> None:
        pkt = _build_syn(_linux_options(), flags=0x12)  # SYN+ACK
        self.assertIsNone(rewrite_syn(pkt, FP))

    def test_udp_untouched(self) -> None:
        pkt = bytearray(_build_syn(_linux_options()))
        pkt[9] = 17  # proto UDP
        self.assertIsNone(rewrite_syn(bytes(pkt), FP))

    def test_ipv6_untouched(self) -> None:
        pkt = bytearray(_build_syn(_linux_options()))
        pkt[0] = 0x60
        self.assertIsNone(rewrite_syn(bytes(pkt), FP))

    def test_truncated_packet_untouched(self) -> None:
        self.assertIsNone(rewrite_syn(b"\x45\x00\x00\x10", FP))

    def test_custom_fingerprint_respected(self) -> None:
        fp = Fingerprint(window=65535, wscale=7, strip_timestamps=False, ttl=64)
        pkt = _build_syn(_linux_options(), window=29200)
        out = rewrite_syn(pkt, fp)
        assert out is not None
        self.assertEqual(_tcp_field(out, 14, 2), 65535)
        kinds = _parse_opts(out)
        self.assertIn(TCP_OPT_TS, kinds)  # strip_timestamps=False 保留 TS


if __name__ == "__main__":
    unittest.main()