procproxy 0.1.0__py3-none-any.whl

proxycore/upstream/socks4.py

@@ -0,0 +1,106 @@
+"""SOCKS4 / SOCKS4a 客户端。
+
+* 认证：SOCKS4 只有 `USERID`，没有口令认证 → `auth: ntlm/basic` 一律视为
+  `unsupported`（非 Windows 平台的 NTLM 必须标注 unsupported）；
+* SOCKS4a：当目标不是 IP 字面量时用 `0.0.0.x`（x≠0）+ USERID 后跟域名的形式传域名
+  （远端解析）。
+"""
+
+from __future__ import annotations
+
+import contextlib
+import ipaddress
+import socket
+from typing import Any
+
+from ..types import Endpoint
+from .base import (
+    BaseDialer,
+    TargetRefused,
+    UnsupportedAuth,
+    UpstreamDown,
+    optional_hostname,
+    recv_exact,
+)
+
+VER = 0x04
+CMD_CONNECT = 0x01
+CMD_BIND = 0x02
+
+# VN=0 时的 CD 码语义
+CD_REQUEST_GRANTED = 0x5A
+CD_REJECTED = 0x5B
+CD_NO_IDENTD = 0x5C
+CD_IDENTD_MISMATCH = 0x5D
+
+CD_MAP: dict[int, str] = {
+    CD_REQUEST_GRANTED: "request granted",
+    CD_REJECTED: "request rejected or failed",
+    CD_NO_IDENTD: "request rejected: identd not reachable",
+    CD_IDENTD_MISMATCH: "request rejected: identd user mismatch",
+}
+
+# 需要"远端解析域名"的标记地址（SOCKS4a）
+SOCKS4A_MARKER = "0.0.0.1"
+
+
+class Socks4Dialer(BaseDialer):
+    kind = "socks4"
+    supports_udp_flag = False
+
+    def __init__(self, upstream: Any, *, timeout: float = 15.0, resolver: Any = None) -> None:
+        super().__init__(upstream, timeout=timeout)
+        self.resolve_alias = resolver
+
+    def _userid(self) -> bytes:
+        return str(getattr(self.upstream, "user", "") or "").encode("utf-8") + b"\x00"
+
+    def _check_auth(self) -> None:
+        password = str(getattr(self.upstream, "password", "") or "")
+        auth = str(getattr(self.upstream, "auth", "none") or "none").lower()
+        if password or auth in {"basic", "ntlm"}:
+            raise UnsupportedAuth(
+                f"上游 {self.alias}: SOCKS4 不支持口令认证（auth={auth}）；"
+                "请改用 socks5，或去掉 password_ref"
+            )
+
+    def _request(self, sock: socket.socket, dst: Endpoint, hostname: str | None) -> None:
+        address = hostname or dst.addr
+        try:
+            ip = ipaddress.ip_address(address)
+            packed = ip.packed
+            tail = b""
+        except ValueError:
+            # SOCKS4a：0.0.0.x + 域名
+            packed = socket.inet_aton(SOCKS4A_MARKER)
+            tail = address.encode("idna") + b"\x00"
+        sock.sendall(bytes([VER, CMD_CONNECT]) + dst.port.to_bytes(2, "big") + packed + self._userid() + tail)
+        reply = recv_exact(sock, 8)
+        if reply[0] not in (0x00, 0x04):
+            raise UpstreamDown(f"上游 SOCKS4 回复版本异常: 0x{reply[0]:02x}")
+        code = reply[1]
+        if code != CD_REQUEST_GRANTED:
+            text = CD_MAP.get(code, f"unknown cd 0x{code:02x}")
+            if code == CD_REJECTED:
+                raise TargetRefused(f"目标不可达（经 SOCKS4 代理）{address}:{dst.port}: {text}")
+            raise UpstreamDown(f"SOCKS4 请求被拒 {address}:{dst.port}: {text}")
+
+    def tunnel_over(self, sock: socket.socket, dst: Endpoint, hostname: str | None = None) -> socket.socket:
+        self._check_auth()
+        sock.settimeout(self.timeout)
+        address = optional_hostname(dst, hostname) or dst.addr
+        self._request(sock, dst, address)
+        return sock
+
+    def dial_tcp_sync(self, dst: Endpoint, hostname: str | None = None) -> socket.socket:
+        self._check_auth()
+        remote = optional_hostname(dst, hostname) or dst.addr
+        sock = self._connect_upstream()
+        try:
+            self._request(sock, dst, remote)
+        except BaseException:
+            with contextlib.suppress(OSError):
+                sock.close()
+            raise
+        sock.settimeout(None)
+        return sock

proxycore/upstream/socks5.py

@@ -0,0 +1,266 @@
+"""SOCKS5 客户端（RFC 1928 / RFC 1929），含 UDP ASSOCIATE。
+
+硬约束：
+* `hostname` 非空时必须用 `ATYP=0x03` 传域名（远端解析，防 DNS 泄漏）；
+* 认证失败 → `AuthFailed`；目标不可达 → `TargetRefused`；上游不可达 → `UpstreamDown`；
+* 分片 UDP（`FRAG != 0`）→ 丢弃并记 WARN。
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import socket
+import struct
+import time
+from typing import Any
+
+from ..types import Endpoint, UdpAssociation
+from .base import (
+    AuthFailed,
+    BaseDialer,
+    TargetRefused,
+    UpstreamDown,
+    UpstreamError,
+    optional_hostname,
+    pack_port,
+    parse_port,
+    recv_exact,
+)
+
+VER = 0x05
+AUTH_NONE = 0x00
+AUTH_USERPASS = 0x02
+AUTH_NO_ACCEPTABLE = 0xFF
+
+CMD_CONNECT = 0x01
+CMD_BIND = 0x02
+CMD_UDP_ASSOCIATE = 0x03
+
+ATYP_V4 = 0x01
+ATYP_DOMAIN = 0x03
+ATYP_V6 = 0x04
+
+# RFC 1928 REP 码 → 语义
+REP_SUCCESS = 0x00
+REP_GENERAL_FAILURE = 0x01
+REP_NOT_ALLOWED = 0x02
+REP_NET_UNREACHABLE = 0x03
+REP_HOST_UNREACHABLE = 0x04
+REP_CONN_REFUSED = 0x05
+REP_TTL_EXPIRED = 0x06
+REP_CMD_UNSUPPORTED = 0x07
+REP_ATYP_UNSUPPORTED = 0x08
+
+REP_MAP: dict[int, str] = {
+    REP_SUCCESS: "success",
+    REP_GENERAL_FAILURE: "general failure",
+    REP_NOT_ALLOWED: "connection not allowed by ruleset",
+    REP_NET_UNREACHABLE: "network unreachable",
+    REP_HOST_UNREACHABLE: "host unreachable",
+    REP_CONN_REFUSED: "connection refused",
+    REP_TTL_EXPIRED: "TTL expired",
+    REP_CMD_UNSUPPORTED: "command not supported",
+    REP_ATYP_UNSUPPORTED: "address type not supported",
+}
+
+# 这些 REP 表示"上游本身没问题，是目标的事" → TargetRefused
+_TARGET_REPS = {REP_CONN_REFUSED, REP_HOST_UNREACHABLE, REP_NET_UNREACHABLE, REP_TTL_EXPIRED}
+
+
+def encode_address(host: str, port: int) -> bytes:
+    """按 ATYP 规则编码地址（域名优先透传，交上游解析）。"""
+    kind, raw, atyp = classify_address(host)
+    if kind == "domain":
+        payload = raw.encode("idna") if any(ord(c) > 127 for c in raw) else raw.encode("ascii", "ignore")
+        if len(payload) > 255:
+            raise UpstreamDown(f"域名过长（{len(payload)} > 255）: {raw}")
+        return bytes([ATYP_DOMAIN, len(payload)]) + payload + pack_port(port)
+    return bytes([atyp]) + raw + pack_port(port)
+
+
+def classify_address(host: str) -> tuple[str, Any, int]:
+    """返回 `("ipv4"|"ipv6"|"domain", raw, atyp)`。"""
+    import ipaddress
+
+    try:
+        ip = ipaddress.ip_address(host)
+    except ValueError:
+        return "domain", host, ATYP_DOMAIN
+    if ip.version == 6:
+        return "ipv6", ip.packed, ATYP_V6
+    return "ipv4", ip.packed, ATYP_V4
+
+
+def read_address(sock: socket.socket, atyp: int) -> Endpoint:
+    if atyp == ATYP_V4:
+        raw = recv_exact(sock, 4)
+        port = parse_port(recv_exact(sock, 2))
+        return Endpoint(addr=socket.inet_ntoa(raw), port=port, is_v6=False)
+    if atyp == ATYP_V6:
+        raw = recv_exact(sock, 16)
+        port = parse_port(recv_exact(sock, 2))
+        return Endpoint(addr=socket.inet_ntop(socket.AF_INET6, raw), port=port, is_v6=True)
+    if atyp == ATYP_DOMAIN:
+        length = recv_exact(sock, 1)[0]
+        host = recv_exact(sock, length).decode("latin1")
+        port = parse_port(recv_exact(sock, 2))
+        return Endpoint(addr=host, port=port, is_v6=False)
+    raise UpstreamDown(f"上游返回了不支持的 ATYP=0x{atyp:02x}")
+
+
+def rep_to_error(rep: int, target: str) -> UpstreamError:
+    text = REP_MAP.get(rep, f"unknown rep 0x{rep:02x}")
+    if rep in _TARGET_REPS:
+        return TargetRefused(f"目标不可达（经代理）{target}: {text}")
+    return UpstreamDown(f"上游拒绝请求 {target}: {text}")
+
+
+class Socks5Dialer(BaseDialer):
+    kind = "socks5"
+    supports_udp_flag = True
+
+    def __init__(self, upstream: Any, *, timeout: float = 15.0, resolver: Any = None) -> None:
+        super().__init__(upstream, timeout=timeout)
+        # 供 ChainDialer 使用（把 chain 别名解析成 Upstream 对象）
+        self.resolve_alias = resolver
+
+    # ------------------------------------------------------------ 握手 ----
+    def _greeting(self) -> bytes:
+        user = str(getattr(self.upstream, "user", "") or "")
+        password = str(getattr(self.upstream, "password", "") or "")
+        if user or password:
+            return bytes([VER, 2, AUTH_NONE, AUTH_USERPASS])
+        return bytes([VER, 1, AUTH_NONE])
+
+    def handshake(self, sock: socket.socket) -> None:
+        sock.sendall(self._greeting())
+        reply = recv_exact(sock, 2)
+        if reply[0] != VER:
+            raise UpstreamDown(f"上游 SOCKS 版本不匹配: {reply[0]}")
+        method = reply[1]
+        if method == AUTH_NO_ACCEPTABLE:
+            raise AuthFailed("上游不接受任何认证方式（0xFF）")
+        if method == AUTH_USERPASS:
+            self._userpass_auth(sock)
+        elif method != AUTH_NONE:
+            raise UpstreamDown(f"上游选择了不支持的认证方式 0x{method:02x}")
+
+    def _userpass_auth(self, sock: socket.socket) -> None:
+        user = str(getattr(self.upstream, "user", "") or "").encode("utf-8")
+        password = str(getattr(self.upstream, "password", "") or "").encode("utf-8")
+        if len(user) > 255 or len(password) > 255:
+            raise AuthFailed("用户名/口令超过 255 字节")
+        sock.sendall(bytes([0x01, len(user)]) + user + bytes([len(password)]) + password)
+        reply = recv_exact(sock, 2)
+        if reply[1] != 0x00:
+            raise AuthFailed(f"上游认证失败（SOCKS5 用户名/口令，status=0x{reply[1]:02x}）")
+
+    # ------------------------------------------------------------ 请求 ----
+    def _request(self, sock: socket.socket, cmd: int, dst: Endpoint, hostname: str | None) -> Endpoint:
+        address = hostname or dst.addr
+        sock.sendall(bytes([VER, cmd, 0x00]) + encode_address(address, dst.port))
+        head = recv_exact(sock, 4)
+        if head[0] != VER:
+            raise UpstreamDown(f"上游回复版本异常: {head[0]}")
+        if head[1] != REP_SUCCESS:
+            raise rep_to_error(head[1], f"{address}:{dst.port}")
+        return read_address(sock, head[3])
+
+    def tunnel_over(self, sock: socket.socket, dst: Endpoint, hostname: str | None = None) -> socket.socket:
+        """在已建立的连接之上再开一条 CONNECT（链式代理用）。"""
+        address = optional_hostname(dst, hostname) or dst.addr
+        sock.settimeout(self.timeout)
+        self.handshake(sock)
+        self._request(sock, CMD_CONNECT, dst, address)
+        return sock
+
+    # ------------------------------------------------------------ 拨号 ----
+    def dial_tcp_sync(self, dst: Endpoint, hostname: str | None = None) -> socket.socket:
+        remote = optional_hostname(dst, hostname) or dst.addr
+        sock = self._connect_upstream()
+        try:
+            self.handshake(sock)
+            self._request(sock, CMD_CONNECT, dst, remote)
+        except BaseException:
+            with contextlib.suppress(OSError):
+                sock.close()
+            raise
+        sock.settimeout(None)
+        return sock
+
+    # ------------------------------------------------------ UDP ASSOCIATE --
+    def dial_udp_sync(self, dst: Endpoint) -> UdpAssociation:
+        """建立 UDP 关联；按 RFC 1928 用 `0.0.0.0:0` 表示"还不知道要发去哪儿"。"""
+        if not self.supports_udp():
+            raise UpstreamDown(f"上游 {self.alias} 未启用 UDP（配置 upstream.udp=true 才可用）")
+        sock = self._connect_upstream()
+        try:
+            self.handshake(sock)
+            relay = self._request(sock, CMD_UDP_ASSOCIATE, Endpoint(addr="0.0.0.0", port=0), None)
+        except BaseException:
+            with contextlib.suppress(OSError):
+                sock.close()
+            raise
+        # 上游可能返回 0.0.0.0 / 内网地址 → 按"控制连接的同源地址"处理
+        peer_host = str(sock.getpeername()[0])
+        relay_host = relay.addr
+        if relay_host in {"0.0.0.0", "::"} or relay.port == 0 or relay_host.startswith("127."):
+            relay_host = peer_host
+        return UdpAssociation(
+            relay=Endpoint(addr=relay_host, port=relay.port),
+            control_sock=sock,
+            upstream=self.alias,
+            created_at=time.time(),
+        )
+
+    async def dial_udp(self, dst: Endpoint) -> UdpAssociation:
+        # UDP ASSOCIATE 是完整 TCP 握手 + 认证，不能直接跑在事件循环上。
+        return await asyncio.to_thread(self.dial_udp_sync, dst)
+
+
+# ------------------------------------------------------------- 报文编解码 --
+def build_udp_datagram(dst: Endpoint, payload: bytes, hostname: str | None = None) -> bytes:
+    """`RSV(2) | FRAG(1) | ATYP | ADDR | PORT | DATA`（FRAG 必须为 0）。"""
+    header = b"\x00\x00\x00" + encode_address(optional_hostname(dst, hostname) or dst.addr, dst.port)
+    return header + payload
+
+
+def _require(packet: bytes, offset: int, length: int, what: str) -> None:
+    """回包来自不可信网络，任何字段都必须先确认长度再切。"""
+    if offset + length > len(packet):
+        raise UpstreamDown(f"UDP 回包过短：{what} 需要 {length} 字节，实际只剩 {len(packet) - offset}")
+
+
+def parse_udp_datagram(packet: bytes) -> tuple[Endpoint, bytes]:
+    """解析回包；FRAG != 0 或长度不足直接丢弃（抛 `UpstreamDown`，调用方记 WARN）。
+
+    报文布局：`RSV(2) | FRAG(1) | ATYP(1) | ADDR | PORT(2) | DATA`
+    """
+    if len(packet) < 4:
+        raise UpstreamDown(f"UDP 报文过短（{len(packet)} 字节）")
+    _rsv1, _rsv2, frag, atyp = packet[0], packet[1], packet[2], packet[3]
+    if frag != 0x00:
+        raise UpstreamDown(f"不支持分片 UDP（FRAG=0x{frag:02x}），按协议规定丢弃")
+    offset = 4  # ADDR 从 ATYP 之后开始
+    if atyp == ATYP_V4:
+        _require(packet, offset, 4, "IPv4 地址")
+        addr = socket.inet_ntoa(packet[offset : offset + 4])
+        offset += 4
+    elif atyp == ATYP_V6:
+        _require(packet, offset, 16, "IPv6 地址")
+        addr = socket.inet_ntop(socket.AF_INET6, packet[offset : offset + 16])
+        offset += 16
+    elif atyp == ATYP_DOMAIN:
+        _require(packet, offset, 1, "域名长度")
+        length = packet[offset]
+        offset += 1
+        _require(packet, offset, length, "域名")
+        addr = packet[offset : offset + length].decode("latin1")
+        offset += length
+    else:
+        raise UpstreamDown(f"UDP 回包 ATYP 非法: 0x{atyp:02x}")
+    _require(packet, offset, 2, "端口")
+    port = struct.unpack(">H", packet[offset : offset + 2])[0]
+    offset += 2
+    return Endpoint(addr=addr, port=port, is_v6=atyp == ATYP_V6), packet[offset:]