#!/usr/bin/env python3
"""ws_prescan.py — deterministic WebSocket / realtime-surface pre-scan.

ARION service artifact. Scans application source for the signatures of
an unsafe realtime surface: upgrade acceptors with no Origin validation
(cross-site WebSocket hijacking), origins explicitly opened to '*',
message payloads flowing into exec/SQL/HTTP sinks, payloads re-broadcast
to every connected client unfiltered, plaintext ws:// endpoints, and
acceptors with no auth handshake or frame-size vocabulary. Line-wise
taint (name assignment only), fully offline, no inference.

Findings (base severity):
  WS-ORIGIN-BYPASS     HIGH    an explicit always-true / wildcard
                               origin acceptance: gorilla CheckOrigin
                               returning literal true, socket.io
                               origins('*:*') or cors origin:'*',
                               allowedOrigins ['*'], verifyClient
                               () => true — any web page can open a
                               socket riding the victim's cookies
  WS-MSG-EXEC          HIGH    a value derived from a socket message
                               (handler param, recv()/ReadMessage()
                               result, event.data, receive_* arg)
                               reaches eval/exec/os.system/subprocess/
                               execSync/Runtime.exec/ProcessBuilder/
                               new Function/vm.runIn*/pickle.loads/
                               __import__ — message bytes become code
  WS-MSG-SQL           HIGH    socket message into .execute/
                               .executemany/.raw/.query/$queryRaw —
                               a frame becomes a query fragment
  WS-NO-ORIGIN-CHECK   MEDIUM  file-level: the file accepts socket
                               connections on a library whose DEFAULT
                               is no origin validation — node `ws`,
                               socket.io, python websockets.serve /
                               starlette websocket.accept / django
                               *WebsocketConsumer, Java @OnOpen —
                               and the file carries zero origin
                               vocabulary (origin/checkOrigin/
                               verifyClient/allowedOrigins/
                               AllowedHostsOriginValidator/
                               allowRequest). CSWSH surface: any site
                               the victim visits can drive their
                               socket. NOT flagged for gorilla
                               Upgrade( or Phoenix join( — both
                               default to host-checked origins
  WS-MSG-FETCH         MEDIUM  socket message into requests.*, fetch,
                               axios, urlopen, http.Get — the frame
                               chooses the outbound target (SSRF
                               channel over the socket)
  WS-INSECURE-URL      MEDIUM  ws:// literal to a non-local host —
                               frames (and any tokens in them) cross
                               the wire unencrypted; localhost/127/
                               0.0.0.0 targets silent
  WS-MSG-REFLECT       LOW     socket message immediately re-sent to
                               other clients (io.emit / broadcast /
                               wss.clients.forEach send /
                               group_send / socket.broadcast) — an
                               unfiltered reflection channel: one
                               client's bytes reach every connected
                               DOM; review what the payload carries
  WS-NO-AUTH           LOW     file-level: acceptor present and zero
                               auth handshake vocabulary (auth/token/
                               jwt/session/cookie/passport/
                               handshake.auth/middleware) — every
                               connection is anonymous; fine for
                               public feeds, worth one line of review
  WS-NO-MSGLIMIT       LOW     file-level: acceptor present and zero
                               frame/message-size vocabulary
                               (maxPayload/maxMessageSize/
                               SetReadLimit/maxHttpBufferSize/
                               read_limit/max_size) — oversized-frame
                               memory pressure surface

Context rules:
  * comment lines and block-comment bodies are stripped first
  * import/require lines never flag
  * test/spec/fixture paths downgrade one tier
  * taint is name-assignment only: handler params, recv()/Read/
    receive_* results, event.data, destructured message fields, and
    names assigned from tainted names; no cross-function or
    cross-file flow
  * all findings require realtime vocabulary somewhere in the file —
    a coincidence `socket.on(` in a non-socket file never gates open
  * file-level codes (NO-ORIGIN-CHECK, NO-AUTH, NO-MSGLIMIT) need a
    server-side acceptor, not a client `new WebSocket(` — a client
    connecting out has nothing to CSWSH

Scope: deterministic pre-scan for human review — flags the line
shapes where the socket surface crosses a trust boundary; it cannot
see proxy-level auth, reverse-proxy origin rules, or middleware
defined in other files. Clean output != a safe realtime surface.

  ws_prescan.py <file-or-dir> [--json]
  ws_prescan.py --selftest
Stdlib only. Exit 0 = no HIGH/MEDIUM findings, 1 = findings, 2 = usage.
"""
import json
import os
import re
import sys

SEV_ORDER = {"HIGH": 0, "MEDIUM": 1, "LOW": 2, "INFO": 3}

TARGET_EXT = (
    ".py", ".js", ".jsx", ".ts", ".tsx", ".mjs", ".cjs",
    ".php", ".rb", ".go", ".java", ".kt", ".cs", ".groovy",
    ".rs", ".scala", ".dart", ".vue", ".svelte", ".ex", ".exs",
)

SKIP_DIRS = {
    "node_modules", ".git", "vendor", "dist", "build", "target",
    "__pycache__", ".venv", "venv", "coverage", ".next", "out",
}

IMPORT_LINE = re.compile(
    r"^\s*(?:from\s+[\w.]+\s+import\b|import\s+[\w.{*]|"
    r"(?:const|let|var)\s+[\w${},\s]*=\s*require\s*\(|"
    r"\buse\s+[\w:]+|require\s*\(['\"][\w./-]+['\"]\s*\)?\s*;?\s*$|"
    r"#include|package\s+[\w.]+;?\s*$|using\s+[\w.]+;|"
    r"alias\s+[\w.]+|defmodule\s)")

TEST_PATH = re.compile(
    r"(^|/)(test|tests|spec|specs|fixtures?|mocks?|__tests__|__mocks__|"
    r"e2e|cypress|playwright)(/|$)|(_test|_spec|\.test|\.spec)\.",
    re.I)

# --- vocabulary -----------------------------------------------------------

# strong markers — the file actually wires a realtime/socket layer
WS_STRONG = re.compile(
    r"new\s+WebSocketServer|new\s+WebSocket\.Server|"
    r"WebSocketServer\s*\(|ws\.Server|"
    r"require\(\s*['\"](?:ws|socket\.io|sockjs|websocket|"
    r"@nestjs/websockets)[\w./-]*['\"]\s*\)|"
    r"from\s+['\"](?:ws|socket\.io|socket\.io-client|sockjs)[\w./-]*['\"]|"
    r"\bio\.on\(\s*['\"]connection|socket\.on\(\s*['\"]|"
    r"\bio\.emit\(|socket\.broadcast|SockJS|"
    r"import\s+websockets|from\s+websockets|websockets\.serve|"
    r"WebsocketConsumer|AsyncWebsocketConsumer|JsonWebsocketConsumer|"
    r"websocket_connect|websocket_receive|websocket_disconnect|"
    r"websocket\.accept\s*\(|WebSocketResponse|add_ws_route|"
    r"gorilla/websocket|websocket\.Upgrader|\.Upgrade\s*\(|"
    r"SetReadLimit|javax\.websocket|@OnMessage|@OnOpen|"
    r"@ServerEndpoint|ServerEndpointExporter|"
    r"use\s+Phoenix\.Channel|def\s+join\s*\([^)]*socket|"
    r"ActionCable|:\s*Hub\b|MapHub|Microsoft\.AspNetCore\.SignalR|"
    r"\bws://|wss://|new\s+WebSocket\s*\(|aiohttp\.web_ws|"
    r"WebSocketHandler|@WebSocket|WebSocketGateway|"
    r"Ratchet|swoole_websocket|tungstenite|tokio_tungstenite|"
    r"warp::ws|axum.*WebSocket|ws::Ws|"
    r"group_send|channel_layer", re.I)

# server-side acceptors — a client new WebSocket( never counts
ACCEPTOR = re.compile(
    r"new\s+WebSocketServer|new\s+WebSocket\.Server|"
    r"WebSocketServer\s*\(|ws\.Server|createWebSocketServer|"
    r"\bio\.on\(\s*['\"]connection|\bio\.of\s*\(|"
    r"new\s+Server\s*\([^)]*(?:http|server)|"
    r"websockets\.serve|websocket\.accept\s*\(|"
    r"WebsocketConsumer|AsyncWebsocketConsumer|"
    r"JsonWebsocketConsumer|websocket_connect|"
    r"websocket\.Upgrader|\.Upgrade\s*\(|"
    r"@OnOpen|@ServerEndpoint|def\s+join\s*\([^)]*socket|"
    r":\s*Hub\b|MapHub|WebSocketHandler|@WebSocket|"
    r"WebSocketGateway|Ratchet|swoole_websocket|"
    r"tungstenite|WebSocketUpgrade|axum.*WebSocket|"
    r"\.websocket\s*\(\s*['\"]|add_ws_route|"
    r"registerOnConnection|OnConnectedAsync")

# acceptor classes whose DEFAULT origin posture is unchecked —
# gorilla (host-equals-Origin default) and Phoenix (check_origin
# default) are deliberately absent
UNSAFE_ORIGIN_ACCEPTOR = re.compile(
    r"new\s+WebSocketServer|new\s+WebSocket\.Server|"
    r"WebSocketServer\s*\(|ws\.Server|createWebSocketServer|"
    r"\bio\.on\(\s*['\"]connection|\bio\.of\s*\(|"
    r"new\s+Server\s*\([^)]*(?:http|server)|"
    r"websockets\.serve|websocket\.accept\s*\(|"
    r"WebsocketConsumer|AsyncWebsocketConsumer|"
    r"JsonWebsocketConsumer|websocket_connect|"
    r"@OnOpen|@ServerEndpoint|WebSocketGateway|"
    r"Ratchet|swoole_websocket|"
    r"\.websocket\s*\(\s*['\"]|add_ws_route|"
    r"registerOnConnection|OnConnectedAsync")

# explicit always-true / wildcard origin acceptance
ORIGIN_BYPASS = re.compile(
    r"CheckOrigin\s*[:=]\s*(?:func|func\s*\([^)]*\)\s*bool\s*\{?)"
    r"[^}\n]{0,120}return\s+true|"
    r"origins\s*\(\s*['\"]\*:\*['\"]|origins\s*[:=]\s*"
    r"[\['\"]\s*\*|allow_origins\s*=\s*\[\s*['\"]\*['\"]|"
    r"origin\s*[:=]\s*['\"]\*['\"]|"
    r"allowedOrigins\s*[:=]\s*\[\s*['\"]\*['\"]|"
    r"AllowedOrigins\s*[:=]\s*\[\]\s*string\s*\{\s*['\"]\*|"
    r"verifyClient\s*[:=]\s*(?:\([^)]*\)|function[^(]*\([^)]*\))"
    r"\s*(?:=>|\{)\s*(?:return\s+)?true|"
    r"allowRequest\s*[:=]\s*(?:\([^)]*\)|function[^(]*\([^)]*\))"
    r"\s*(?:=>|\{)\s*(?:return\s+)?true|"
    r"cors\s*[:=]\s*\{[^}]{0,80}origin\s*[:=]\s*['\"]\*['\"]|"
    r"setAllowedOrigins\s*\(\s*['\"]\*['\"]|"
    r"AllowedHosts\s*[:=]\s*\[\s*['\"]\*['\"]", re.I)

# origin-check vocabulary — presence silences NO-ORIGIN-CHECK
ORIGIN_VOCAB = re.compile(
    r"origin|checkOrigin|verifyClient|allowedOrigin|"
    r"AllowedHostsOriginValidator|allowRequest|"
    r"check_origin|allowed_hosts|setAllowedOrigins|\bcors\b|"
    r"handshake\.headers|allowedHosts|same.?origin|"
    r"Access-Control-Allow-Origin", re.I)

# auth handshake vocabulary — presence silences NO-AUTH
AUTH_VOCAB = re.compile(
    r"\bauth\b|authenticat|authori[sz]|\btoken\b|jwt|"
    r"session[._]|sessionId|req\.session|express-session|"
    r"sessionStore|cookie|credential|passport|"
    r"handshake\.auth|handshake\.query|"
    r"Sec-WebSocket-Protocol.*(?:auth|token)|"
    r"subprotocol|middleware|guards?|"
    r"current_user|context\.user|scope\[|"
    r"principal|identity|login|"
    r"verify\w*\s*\(|isAuthenticated|"
    r"Sec-WebSocket-Key.*header", re.I)

# frame/message-size vocabulary — presence silences NO-MSGLIMIT
LIMIT_VOCAB = re.compile(
    r"maxPayload|max_payload|maxMessageSize|max_message_size|"
    r"SetReadLimit|setReadLimit|readLimit|read_limit|"
    r"maxHttpBufferSize|maxFrameSize|max_frame_size|"
    r"maxSize|max_size|messageSize|message_size|"
    r"maxLength|max_length|payloadLimit|maxReceived|"
    r"maxIncomingMessage|max_message|maxFramePayload|"
    r"max_ws_message|MaxMessageBytes|perMessageDeflate.{0,30}"
    r"max", re.I)

# insecure plaintext endpoint — ws:// to a routable host
WS_URL = re.compile(r"\bws://([^\s'\"`/]+)")
LOCAL_HOST = re.compile(
    r"^(?:localhost|127\.|0\.0\.0\.0|\[::1\]|::1|"
    r"192\.168\.|10\.|172\.(?:1[6-9]|2\d|3[01])\.)")

# exec / code-eval sinks
EXEC_SINK = re.compile(
    r"\beval\s*\(|\bexec\s*\(|execSync\s*\(|execFileSync\s*\(|"
    r"os\.system\s*\(|os\.popen\s*\(|"
    r"subprocess\.(?:call|run|Popen|check_output|check_call)\s*\(|"
    r"child_process|spawnSync\s*\(|"
    r"Runtime\.getRuntime\(\)\.exec|ProcessBuilder\s*\(|"
    r"new\s+Function\s*\(|vm\.runIn\w*\s*\(|"
    r"pickle\.loads\s*\(|marshal\.loads\s*\(|"
    r"__import__\s*\(|importlib\.import_module\s*\(|"
    r"shell_exec\s*\(|passthru\s*\(|pcntl_exec\s*\(|"
    r"System\.cmd\s*\(|Port\.open\s*\(|:os\.cmd\s*\(")

SQL_SINK = re.compile(
    r"\.execute(?:many|script)?\s*\(|\.raw\s*\(|\.query\s*\(|"
    r"\$queryRaw|\$executeRaw|executeQuery\s*\(|"
    r"sequelize\.query|knex\.raw|db\.run\s*\(|conn\.exec\s*\(|"
    r"execute_query\s*\(|pg_query|mysql_query")

FETCH_SINK = re.compile(
    r"requests\.(?:get|post|put|head|request)\s*\(|"
    r"urllib\.request\.urlopen|urlopen\s*\(|\bfetch\s*\(|"
    r"axios\.(?:get|post|put|request)\s*\(|got\.(?:get|post)\s*\(|"
    r"http\.Get\s*\(|http\.NewRequest|http\.Post\s*\(|"
    r"needle\.(?:get|post)|superagent\.|wp_remote_get|"
    r"file_get_contents\s*\(|curl_init|RestTemplate|"
    r"\.getForObject|WebClient|HTTPoison|Finch|"
    r"Net::HTTP|Faraday")

# reflection sinks — payload re-sent to other clients
REFLECT_SINK = re.compile(
    r"\bio\.emit\s*\(|socket\.broadcast\.emit\s*\(|"
    r"\bbroadcast\s*\(|wss\.clients\.forEach|"
    r"\.clients\.forEach.*send|server\.publish\s*\(|"
    r"group_send\s*\(|channel_layer\.group_send|"
    r"hub\.broadcast|\.WriteMessage\s*\(|"
    r"emitter\.emit\s*\(|pubsub\.publish\s*\(")

# --- taint sources --------------------------------------------------------

# .on('<event>', (param, ...) => / function (param, ...)) — first arg
# of a socket event handler is remote-controlled
ON_HANDLER = re.compile(
    r"\.on\(\s*['\"]\w[\w -]*['\"]\s*,\s*(?:async\s+)?"
    r"(?:\(\s*([A-Za-z_][\w$]*)|\(\s*\{\s*([^}]*)\}|"
    r"function\s*\(\s*([A-Za-z_][\w$]*)|"
    r"function\s*\(\s*\{\s*([^}]*)\}|"
    r"([A-Za-z_][\w$]*)\s*=>)")

# message-data expressions on any line (event.data / e.data etc.)
EVENT_DATA = re.compile(
    r"\b(?:event|evt|e|ev|frame|packet|message|msg)\.data\b")

# python handler defs — def receive(self, text_data): /
# def on_message(ws, message): / async def handler(websocket, ...)
PY_HANDLER = re.compile(
    r"def\s+(?:receive|receive_json|receive_text|receive_bytes|"
    r"on_message|on_open|on_connect|websocket_receive|handle|"
    r"dispatch|consume)_?\w*\s*\(([^)]*)\)")

# java handler — @OnMessage on prior line handled in pass via flag
JAVA_HANDLER = re.compile(
    r"(?:public|private|protected)\s+\w[\w<>\[\]]*\s+"
    r"(?:on\w+|handle\w*|receive\w*)\s*\(([^)]*)\)")

# recv / read results — name(s) = await ws.recv() /
# mt, msg, err := conn.ReadMessage() / = websocket.receive_text()
RECV_ASSIGN = re.compile(
    r"=\s*(?:await\s+)?\w*\.(?:recv|receive|receive_text|"
    r"receive_json|receive_bytes|ReadMessage|ReadJSON|"
    r"readMessage|read)\s*\(|"
    r":?=\s*(?:await\s+)?\w*\.ReadMessage\s*\(")

# `async for message in websocket` / `for msg in ws` iteration
FOR_RECV = re.compile(
    r"(?:async\s+)?for\s+([A-Za-z_]\w{0,39})\s+in\s+"
    r"(?:websocket|ws|sock|conn|connection|socket|client)\b")

PER_CODE_CAP = 40


class Finding:
    def __init__(self, sev, code, where, line, msg):
        self.sev, self.code, self.where = sev, code, where
        self.line, self.msg = line, msg

    def as_dict(self):
        return {"severity": self.sev, "code": self.code,
                "file": self.where, "line": self.line, "msg": self.msg}


# --- helpers --------------------------------------------------------------
def strip_comments(lines):
    """Remove full-line comments and /* */ block bodies. Keeps count."""
    out, in_block = [], False
    for ln in lines:
        s = ln
        if in_block:
            if "*/" in s:
                s = s.split("*/", 1)[1]
                in_block = False
            else:
                out.append("")
                continue
        while "/*" in s:
            if "*/" in s.split("/*", 1)[1]:
                s = s.replace(s[s.index("/*"):s.index("*/") + 2], " ", 1)
            else:
                s = s.split("/*", 1)[0]
                in_block = True
        st = s.strip()
        if st.startswith(("//", "#", "*", "<!--")):
            out.append("")
            continue
        out.append(s)
    return out


def downgrade(sev):
    return "MEDIUM" if sev == "HIGH" else "LOW"


ASSIGN_RX = re.compile(
    r"^\s*(?:const|let|var|final|static|auto|val|String|self\.)?"
    r"\s*\.?\s*([A-Za-z_]\w{0,39})\s*"
    r"(?::\s*[\w<>\[\],. |]+\s*)?=\s*(.+?);?\s*$")
SELF_ASSIGN = re.compile(r"^\s*self\.([A-Za-z_]\w{0,39})\s*=\s*(.+)$")
# go multi-assign: a, b, c := expr
GO_MULTI = re.compile(
    r"^\s*([A-Za-z_][\w, ]*?)\s*:?=\s*(.+?);?\s*$")


def rhs_taints(rhs, tainted):
    for name in tainted:
        if re.search(r"\b%s\b" % re.escape(name), rhs):
            return True
    return False


def names_in_line(ln, tainted):
    for name in tainted:
        if re.search(r"\b%s\b" % re.escape(name), ln):
            yield name


def scan_file(path, rel):
    try:
        raw = open(path, encoding="utf-8", errors="replace").read()
    except OSError:
        return []
    is_test = bool(TEST_PATH.search(rel))
    lines = strip_comments(raw.splitlines())
    text = "\n".join(lines)

    # file-level gate — no realtime wiring: not a socket file
    if not WS_STRONG.search(text):
        return []

    file_has_acceptor = bool(ACCEPTOR.search(text))
    file_has_unsafe_acceptor = bool(UNSAFE_ORIGIN_ACCEPTOR.search(text))
    file_has_origin = bool(ORIGIN_VOCAB.search(text))
    file_has_auth = bool(AUTH_VOCAB.search(text))
    file_has_limit = bool(LIMIT_VOCAB.search(text))
    findings = []
    counts = {}
    tainted = set()
    on_msg_pending = False

    def emit(sev, code, ln, msg):
        if is_test:
            sev = downgrade(sev)
        if counts.get(code, 0) >= PER_CODE_CAP:
            return
        counts[code] = counts.get(code, 0) + 1
        findings.append(Finding(sev, code, rel, ln, msg))

    # ---- pass 1: taint propagation + per-line findings ------------------
    for i, ln in enumerate(lines, 1):
        st = ln.strip()
        if not st or IMPORT_LINE.match(ln):
            continue

        # -- socket event handler params ----------------------------------
        for hm in ON_HANDLER.finditer(ln):
            grp = hm.group(1) or hm.group(3) or hm.group(5)
            if grp:
                tainted.add(grp)
            for dgrp in (hm.group(2), hm.group(4)):
                if dgrp:
                    for nm in re.findall(r"[A-Za-z_]\w{0,39}", dgrp):
                        tainted.add(nm)

        # event.data / msg.data style references are remote input
        if EVENT_DATA.search(ln):
            line_flags_remote = True
        else:
            line_flags_remote = False

        # python handler defs — taint params after self/ws/websocket
        pm = PY_HANDLER.search(ln)
        if pm:
            params = [p.split("=")[0].split(":")[0].strip()
                      for p in pm.group(1).split(",")]
            for nm in params[1:]:
                if re.match(r"^[A-Za-z_]\w{0,39}$", nm):
                    tainted.add(nm)

        # @OnMessage annotation arms the next java method's params
        if "@OnMessage" in ln or "@OnOpen" in ln:
            on_msg_pending = True
            continue
        if on_msg_pending:
            jm = JAVA_HANDLER.search(ln)
            if jm:
                for p in jm.group(1).split(","):
                    pp = p.strip().split()
                    if pp and pp[-1] not in ("Session", "session"):
                        nm = pp[-1].lstrip("@")
                        if re.match(r"^[A-Za-z_]\w{0,39}$", nm):
                            tainted.add(nm)
                on_msg_pending = False
            elif st and not st.startswith("@"):
                on_msg_pending = False

        # recv/read results taint all LHS names
        rm = RECV_ASSIGN.search(ln)
        if rm:
            lhs = ln[:rm.start()]
            for nm in re.findall(r"[A-Za-z_]\w{0,39}", lhs):
                if nm not in ("const", "let", "var", "err", "error",
                              "_", "mt", "op", "self", "return"):
                    tainted.add(nm)

        fm = FOR_RECV.search(ln)
        if fm:
            tainted.add(fm.group(1))

        # plain name = <tainted expr> / name = event.data propagation
        m = ASSIGN_RX.match(ln) or SELF_ASSIGN.match(ln)
        if m:
            name, rhs = m.group(1), m.group(2)
            if line_flags_remote or rhs_taints(rhs, tainted) \
                    or EVENT_DATA.search(rhs):
                tainted.add(name)
        else:
            gm = GO_MULTI.match(ln)
            if gm and RECV_ASSIGN.search(gm.group(2)):
                for nm in re.findall(r"[A-Za-z_]\w{0,39}", gm.group(1)):
                    if nm not in ("err", "error", "_"):
                        tainted.add(nm)

        line_taints = list(names_in_line(ln, tainted))
        if line_flags_remote:
            line_taints.append("event.data")

        # -- sinks consuming socket payloads --
        if line_taints:
            src = ", ".join(sorted(set(line_taints))[:3])
            if EXEC_SINK.search(ln):
                emit("HIGH", "WS-MSG-EXEC", i,
                     "socket message (%s) reaches an exec/eval sink — "
                     "frame bytes become executable code" % src)
            if SQL_SINK.search(ln):
                emit("HIGH", "WS-MSG-SQL", i,
                     "socket message (%s) into a SQL sink — a frame "
                     "becomes a query fragment" % src)
            if FETCH_SINK.search(ln):
                emit("MEDIUM", "WS-MSG-FETCH", i,
                     "socket message (%s) into an HTTP fetch — the "
                     "frame chooses the outbound target (SSRF "
                     "channel over the socket)" % src)
            if REFLECT_SINK.search(ln):
                emit("LOW", "WS-MSG-REFLECT", i,
                     "socket message (%s) re-sent to other clients "
                     "unfiltered — one client's bytes reach every "
                     "connected DOM" % src)

        # -- per-line surface findings --
        if ORIGIN_BYPASS.search(ln):
            emit("HIGH", "WS-ORIGIN-BYPASS", i,
                 "explicit always-true / wildcard origin acceptance — "
                 "any web page can open a socket riding the victim's "
                 "cookies (CSWSH)")
        um = WS_URL.search(ln)
        if um and not LOCAL_HOST.match(um.group(1)):
            emit("MEDIUM", "WS-INSECURE-URL", i,
                 "plaintext ws:// endpoint to a routable host "
                 "(%s) — frames cross the wire unencrypted"
                 % um.group(1)[:60])

    # ---- file-level: acceptor with zero origin vocabulary ---------------
    if file_has_unsafe_acceptor and not file_has_origin:
        emit_at = None
        for i, ln in enumerate(lines, 1):
            if not ln.strip() or IMPORT_LINE.match(ln):
                continue
            if UNSAFE_ORIGIN_ACCEPTOR.search(ln):
                emit_at = i
                break
            if emit_at is None:
                emit_at = i
        emit("MEDIUM", "WS-NO-ORIGIN-CHECK", emit_at or 1,
             "socket acceptor with zero origin vocabulary — this "
             "library's default accepts any Origin, so any site the "
             "victim visits can drive their socket (CSWSH surface; "
             "verify upstream proxy/gateway rules)")

    # ---- file-level: acceptor with zero auth vocabulary -----------------
    if file_has_acceptor and not file_has_auth:
        emit_at = None
        for i, ln in enumerate(lines, 1):
            if not ln.strip() or IMPORT_LINE.match(ln):
                continue
            if ACCEPTOR.search(ln):
                emit_at = i
                break
            if emit_at is None:
                emit_at = i
        emit("LOW", "WS-NO-AUTH", emit_at or 1,
             "socket acceptor with zero auth-handshake vocabulary — "
             "every connection is anonymous (fine for public feeds, "
             "worth one line of review)")

    # ---- file-level: acceptor with zero frame-size vocabulary -----------
    if file_has_acceptor and not file_has_limit:
        emit_at = None
        for i, ln in enumerate(lines, 1):
            if not ln.strip() or IMPORT_LINE.match(ln):
                continue
            if ACCEPTOR.search(ln):
                emit_at = i
                break
            if emit_at is None:
                emit_at = i
        emit("LOW", "WS-NO-MSGLIMIT", emit_at or 1,
             "socket acceptor with zero frame/message-size "
             "vocabulary — oversized-frame memory pressure surface")

    return findings


def find_targets(root):
    if os.path.isfile(root):
        return [root]
    out = []
    for dp, dns, fns in os.walk(root):
        dns[:] = [d for d in dns if d not in SKIP_DIRS]
        for fn in fns:
            if fn.endswith(TARGET_EXT):
                out.append(os.path.join(dp, fn))
    return sorted(out)


def render(findings, targets):
    findings.sort(key=lambda f: (SEV_ORDER[f.sev], f.where, f.line))
    rows = ["# ws_prescan report", "",
            "files scanned: %d — findings: %d" % (len(targets),
                                                  len(findings)), ""]
    cur = None
    for f in findings:
        if f.sev != cur:
            cur = f.sev
            rows.append("## %s" % cur)
        rows.append("- `%s` %s:%d — %s" % (f.code, f.where, f.line,
                                           f.msg))
    if not findings:
        rows.append("no HIGH/MEDIUM/LOW signatures found — this is a "
                    "pre-scan; clean output != a safe realtime "
                    "surface")
    return "\n".join(rows)


# --- selftest -------------------------------------------------------------
SELFTEST = [
    # --- message -> exec / sql / fetch sinks ---
    ("a.js",
     "const { WebSocketServer } = require('ws');\n"
     "const wss = new WebSocketServer({ port: 8080 });\n"
     "wss.on('connection', (ws) => {\n"
     "  ws.on('message', (data) => {\n"
     "    const cmd = data.toString();\n"
     "    execSync(cmd);\n"
     "  });\n"
     "});",
     {"WS-MSG-EXEC", "WS-NO-ORIGIN-CHECK", "WS-NO-AUTH",
      "WS-NO-MSGLIMIT"}),
    ("a.js",
     "const io = require('socket.io')(server);\n"
     "io.on('connection', (socket) => {\n"
     "  socket.on('run', (payload) => eval(payload.code));\n"
     "});",
     {"WS-MSG-EXEC", "WS-NO-ORIGIN-CHECK", "WS-NO-AUTH",
      "WS-NO-MSGLIMIT"}),
    ("a.py",
     "import websockets\n"
     "async def handler(ws):\n"
     "    async for message in ws:\n"
     "        os.system(message)\n"
     "websockets.serve(handler, 'localhost', 9000)",
     {"WS-MSG-EXEC", "WS-NO-ORIGIN-CHECK", "WS-NO-AUTH",
      "WS-NO-MSGLIMIT"}),
    ("a.js",
     "const { WebSocketServer } = require('ws');\n"
     "const wss = new WebSocketServer({ port: 1 });\n"
     "wss.on('connection', (ws) => ws.on('message', (data) => {\n"
     "  db.query('DELETE WHERE id=' + data);\n"
     "}));",
     {"WS-MSG-SQL", "WS-NO-ORIGIN-CHECK", "WS-NO-AUTH",
      "WS-NO-MSGLIMIT"}),
    ("a.py",
     "from starlette.websockets import WebSocket\n"
     "async def ep(websocket: WebSocket):\n"
     "    await websocket.accept()\n"
     "    body = await websocket.receive_text()\n"
     "    db.execute('INSERT INTO t VALUES (' + body + ')')",
     {"WS-MSG-SQL", "WS-NO-ORIGIN-CHECK", "WS-NO-AUTH",
      "WS-NO-MSGLIMIT"}),
    ("a.js",
     "const { WebSocketServer } = require('ws');\n"
     "const wss = new WebSocketServer({ port: 1 });\n"
     "wss.on('connection', (ws) => ws.on('message', (msg) => {\n"
     "  fetch(msg.url);\n"
     "}));",
     {"WS-MSG-FETCH", "WS-NO-ORIGIN-CHECK", "WS-NO-AUTH",
      "WS-NO-MSGLIMIT"}),
    ("a.go",
     "package main\n"
     "import \"github.com/gorilla/websocket\"\n"
     "func h(c *websocket.Conn) {\n"
     "  _, msg, _ := c.ReadMessage()\n"
     "  http.Get(string(msg))\n"
     "}",
     {"WS-MSG-FETCH"}),
    ("a.js",
     "const io = require('socket.io')(server);\n"
     "io.on('connection', (socket) => {\n"
     "  socket.on('chat', (m) => io.emit('chat', m));\n"
     "});",
     {"WS-MSG-REFLECT", "WS-NO-ORIGIN-CHECK", "WS-NO-AUTH",
      "WS-NO-MSGLIMIT"}),
    # --- origin posture ---
    ("a.go",
     "package main\n"
     "import \"github.com/gorilla/websocket\"\n"
     "var up = websocket.Upgrader{\n"
     "  CheckOrigin: func(r *http.Request) bool { return true },\n"
     "}\n"
     "up.Upgrade(w, r, nil)",
     {"WS-ORIGIN-BYPASS", "WS-NO-AUTH", "WS-NO-MSGLIMIT"}),
    ("a.js",
     "const io = require('socket.io')(server, {\n"
     "  cors: { origin: '*' }\n"
     "});\n"
     "io.on('connection', (s) => s.on('x', (d) => save(d)));",
     {"WS-ORIGIN-BYPASS", "WS-NO-AUTH", "WS-NO-MSGLIMIT"}),
    ("a.js",
     "const io = require('socket.io')(server);\n"
     "io.origins('*:*');\n"
     "io.on('connection', (s) => work(s));",
     {"WS-ORIGIN-BYPASS", "WS-NO-AUTH", "WS-NO-MSGLIMIT"}),
    ("a.js",
     "const { WebSocketServer } = require('ws');\n"
     "const wss = new WebSocketServer({\n"
     "  verifyClient: (info) => true, port: 9 });\n"
     "wss.on('connection', (ws) => go(ws));",
     {"WS-ORIGIN-BYPASS", "WS-NO-AUTH", "WS-NO-MSGLIMIT"}),
    ("a.js",
     "const { WebSocketServer } = require('ws');\n"
     "const wss = new WebSocketServer({ port: 9,\n"
     "  verifyClient: (info) => info.origin === 'https://ok.example' });\n"
     "wss.on('connection', (ws) => go(ws));",
     {"WS-NO-AUTH", "WS-NO-MSGLIMIT"}),
    # gorilla default = host-checked origin -> NO-ORIGIN silent
    ("a.go",
     "package main\n"
     "import \"github.com/gorilla/websocket\"\n"
     "var up = websocket.Upgrader{}\n"
     "func h(w http.ResponseWriter, r *http.Request) {\n"
     "  c, _ := up.Upgrade(w, r, nil)\n"
     "  _, m, _ := c.ReadMessage()\n"
     "  use(m)\n"
     "}",
     {"WS-NO-AUTH", "WS-NO-MSGLIMIT"}),  # gorilla default origin ok;
                                       # no auth/limit vocab honest
    # --- insecure url ---
    ("a.js",
     "const ws = new WebSocket('ws://api.example.com/feed');\n"
     "ws.onmessage = (e) => render(e.data);",
     {"WS-INSECURE-URL"}),
    ("a.js",
     "const ws = new WebSocket('ws://localhost:9000/dev');",
     set()),
    ("a.js",
     "const ws = new WebSocket('wss://api.example.com/feed');",
     set()),
    # --- event.data taint ---
    ("a.js",
     "const wss = new WebSocketServer({ port: 2 });\n"
     "wss.on('connection', (ws) => {\n"
     "  ws.onmessage = (e) => {\n"
     "    eval(e.data);\n"
     "  };\n"
     "});",
     {"WS-MSG-EXEC", "WS-NO-ORIGIN-CHECK", "WS-NO-AUTH",
      "WS-NO-MSGLIMIT"}),
    # --- java @OnMessage ---
    ("a.java",
     "import javax.websocket.OnMessage;\n"
     "@ServerEndpoint(\"/x\")\n"
     "public class X {\n"
     "  @OnMessage\n"
     "  public void onMsg(String body, Session s) {\n"
     "    Runtime.getRuntime().exec(body);\n"
     "  }\n"
     "}",
     {"WS-MSG-EXEC", "WS-NO-ORIGIN-CHECK", "WS-NO-AUTH",
      "WS-NO-MSGLIMIT"}),
    # --- destructured handler param ---
    ("a.js",
     "const io = require('socket.io')(server);\n"
     "io.on('connection', (socket) => {\n"
     "  socket.on('q', ({ sql }) => db.query(sql));\n"
     "});",
     {"WS-MSG-SQL", "WS-NO-ORIGIN-CHECK", "WS-NO-AUTH",
      "WS-NO-MSGLIMIT"}),
    # --- hardened file: auth+origin+limits vocabulary ---
    ("a.js",
     "const { WebSocketServer } = require('ws');\n"
     "const wss = new WebSocketServer({ port: 9,\n"
     "  maxPayload: 4096,\n"
     "  verifyClient: (info) => checkOrigin(info.origin) &&\n"
     "    verifyToken(info.req.headers['authorization']) });\n"
     "wss.on('connection', (ws) => ws.on('message', (d) => ok(d)));",
     set()),
    # --- django channels consumer ---
    ("a.py",
     "from channels.generic.websocket import AsyncWebsocketConsumer\n"
     "class C(AsyncWebsocketConsumer):\n"
     "    async def receive(self, text_data):\n"
     "        eval(text_data)",
     {"WS-MSG-EXEC", "WS-NO-ORIGIN-CHECK", "WS-NO-AUTH",
      "WS-NO-MSGLIMIT"}),
    # --- FP guards ---
    ("a.js",
     "// const wss = new WebSocketServer({ port: 1 });\n"
     "console.log('no sockets here');",
     set()),
    ("a.js",
     "const net = require('net');\n"
     "const data = fs.readFileSync('x');\n"
     "eval(data);",
     set()),
    ("a.js",
     "const { WebSocketServer } = require('ws');\n"
     "const wss = new WebSocketServer({ port: 1 });\n"
     "wss.on('connection', (ws) => ws.on('message', (data) => {\n"
     "  log.info('got', data.length);\n"
     "}));",
     {"WS-NO-ORIGIN-CHECK", "WS-NO-AUTH", "WS-NO-MSGLIMIT"}),
    ("a.js",
     "const io = require('socket.io-client');\n"
     "const s = io('https://api.example.com');\n"
     "s.on('update', (row) => render(row));",
     set()),
    ("test/test_ws.js",
     "const { WebSocketServer } = require('ws');\n"
     "const wss = new WebSocketServer({ port: 1 });\n"
     "wss.on('connection', (ws) => ws.on('message', (d) => eval(d)));",
     {"WS-MSG-EXEC", "WS-NO-ORIGIN-CHECK", "WS-NO-AUTH",
      "WS-NO-MSGLIMIT"}),
    ("a.py",
     "import socket\n"
     "s = socket.socket()\n"
     "data = s.recv(1024)\n"
     "eval(data)",
     set()),
    ("a.js",
     "const { Server } = require('socket.io');\n"
     "const io = new Server(httpServer, {\n"
     "  cors: { origin: 'https://app.example.com' },\n"
     "  maxHttpBufferSize: 1e6,\n"
     "  allowRequest: (req, cb) => cb(null, req.headers.auth === t) });\n"
     "io.on('connection', (s) => s.on('m', (d) => save(d)));",
     set()),
    # python recv_assign propagation
    ("a.py",
     "import websockets\n"
     "async def h(ws):\n"
     "    raw = await ws.recv()\n"
     "    cmd = raw.strip()\n"
     "    subprocess.run(cmd)\n"
     "websockets.serve(h, 'localhost', 1)",
     {"WS-MSG-EXEC", "WS-NO-ORIGIN-CHECK", "WS-NO-AUTH",
      "WS-NO-MSGLIMIT"}),
    # broadcast send to all via clients.forEach
    ("a.js",
     "const { WebSocketServer } = require('ws');\n"
     "const wss = new WebSocketServer({ port: 5 });\n"
     "wss.on('connection', (ws) => ws.on('message', (m) => {\n"
     "  wss.clients.forEach((c) => c.send(m));\n"
     "}));",
     {"WS-MSG-REFLECT", "WS-NO-ORIGIN-CHECK", "WS-NO-AUTH",
      "WS-NO-MSGLIMIT"}),
    # django consumer hardened (auth vocab) still gets NO-ORIGIN/limit
    ("a.py",
     "from channels.generic.websocket import WebsocketConsumer\n"
     "class C(WebsocketConsumer):\n"
     "    def connect(self):\n"
     "        if not self.scope['user'].is_authenticated:\n"
     "            self.close()\n"
     "    def receive(self, text_data):\n"
     "        save(text_data)",
     {"WS-NO-ORIGIN-CHECK", "WS-NO-MSGLIMIT"}),
    # phoenix channel — join acceptor, origin default host-checked:
    # no NO-ORIGIN; auth vocab absent -> NO-AUTH, no limit vocab
    ("a.ex",
     "defmodule App.RoomChannel do\n"
     "  use Phoenix.Channel\n"
     "  def join(\"room:\" <> id, _p, socket) do\n"
     "    {:ok, socket}\n"
     "  end\n"
     "end",
     {"WS-NO-AUTH", "WS-NO-MSGLIMIT"}),
    # plain dns name 'socket' with no socket vocab never opens
    ("a.py",
     "data = request.body\n"
     "eval(data)",
     set()),
]


def run_selftest():
    import tempfile
    fails = []
    for i, (name, src, expect) in enumerate(SELFTEST):
        with tempfile.TemporaryDirectory() as td:
            p = os.path.join(td, name)
            os.makedirs(os.path.dirname(p), exist_ok=True)
            open(p, "w").write(src)
            got = {f.code for f in scan_file(p, name)}
            missing = expect - got
            extra = got - expect
            if missing or extra:
                fails.append(("v%d %s" % (i, name),
                              "missing=%s extra=%s" % (missing, extra)))
    if fails:
        for f_ in fails:
            print("SELFTEST FAIL", f_)
        return 1
    print("SELFTEST PASS (%d vectors)" % len(SELFTEST))
    return 0


def main():
    args = [a for a in sys.argv[1:] if a != "--json"]
    want_json = "--json" in sys.argv
    if "--selftest" in sys.argv:
        return run_selftest()
    if not args:
        print(__doc__)
        return 2
    root = args[0]
    targets = find_targets(root)
    findings = []
    for p in targets:
        rel = os.path.relpath(p, root) if os.path.isdir(root) \
            else os.path.basename(p)
        findings += scan_file(p, rel)
    if want_json:
        print(json.dumps({"findings": [f.as_dict() for f in findings],
                          "scanned": len(targets)}, indent=1))
    else:
        print(render(findings, targets))
    return 1 if any(f.sev in ("HIGH", "MEDIUM") for f in findings) else 0


if __name__ == "__main__":
    sys.exit(main())