arblant
railway-clean-ip-scanner-pc / tests / mock_server.py
mock_server.py114 lines · 3.8 KB
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114
"""A mock VLESS-over-WebSocket endpoint for offline end-to-end tests.

It speaks just enough of the scanner's probe protocol to be verified
as a "clean IP":

  1. TLS with a self-signed certificate (the scanner trusts all certs,
     exactly like the Android original's trust-all manager)
  2. reads the WebSocket upgrade request, answers "101 Switching
     Protocols"
  3. reads the masked VLESS frame the probe sends
  4. answers with the raw bytes "HTTP/1.1 204 No Content" — the exact
     substring the engine's success check looks for

Bound to 0.0.0.0 on an ephemeral port, so every 127.0.0.x loopback
alias reaches it: a generated config set pointing at several
127.0.0.N addresses can be verified against this one server.

The certificate is generated with the openssl CLI (present on the
test machine) into a temporary directory; nothing touches the
network beyond loopback.
"""
import os
import socket
import ssl
import subprocess
import tempfile
import threading


class MockVlessServer:
    def __init__(self):
        self._tmp = tempfile.TemporaryDirectory()
        cert = os.path.join(self._tmp.name, "cert.pem")
        key = os.path.join(self._tmp.name, "key.pem")
        subprocess.run(
            ["openssl", "req", "-x509", "-newkey", "rsa:2048",
             "-keyout", key, "-out", cert, "-days", "2", "-nodes",
             "-subj", "/CN=localhost"],
            check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        self._ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        self._ctx.load_cert_chain(cert, key)
        self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._sock.bind(("0.0.0.0", 0))
        self._sock.listen(16)
        self._sock.settimeout(0.25)
        self.port = self._sock.getsockname()[1]
        self.accepted = 0          # completed 204 handshakes
        self._stop = threading.Event()
        self._thread = None

    def _handle(self, raw):
        raw.settimeout(5)
        try:
            conn = self._ctx.wrap_socket(raw, server_side=True)
        except Exception:
            raw.close()
            return
        try:
            conn.settimeout(5)
            buf = b""
            while b"\r\n\r\n" not in buf and len(buf) < 8192:
                chunk = conn.recv(4096)
                if not chunk:
                    return
                buf += chunk
            if b"upgrade" not in buf.lower():
                return
            conn.sendall(
                b"HTTP/1.1 101 Switching Protocols\r\n"
                b"Upgrade: websocket\r\n"
                b"Connection: Upgrade\r\n"
                b"Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\n\r\n")
            # the probe's masked VLESS frame; content is not inspected
            try:
                conn.recv(4096)
            except Exception:
                return
            conn.sendall(b"HTTP/1.1 204 No Content\r\n"
                         b"Content-Length: 0\r\n\r\n")
            self.accepted += 1
        except Exception:
            pass
        finally:
            try:
                conn.close()
            except Exception:
                pass

    def _serve(self):
        while not self._stop.is_set():
            try:
                raw, _ = self._sock.accept()
            except socket.timeout:
                continue
            except OSError:
                break
            self._handle(raw)

    def __enter__(self):
        self._thread = threading.Thread(target=self._serve, daemon=True)
        self._thread.start()
        return self

    def __exit__(self, *exc):
        self._stop.set()
        try:
            self._sock.close()
        except Exception:
            pass
        if self._thread:
            self._thread.join(timeout=5)
        self._tmp.cleanup()
        return False