xas-test-library/console/native_io.py

"""Bounded IPC and streaming redaction for original-environment runs."""
import json
import os
from pathlib import Path
import re
import shlex
import socket
import subprocess

ANSI = re.compile(rb"\x1b(?:\[[0-?]*[ -/]*[@-~]|\][^\x07\x1b]*(?:\x07|\x1b\\)|[()][0-2A-Z]|[@-_])")
SENSITIVE_NAME = re.compile(r"(?i)(?:api_?key|password|passwd|secret|access_token|auth_token|credential)")


def strip_terminal(data, state="text"):
    output = bytearray()
    for value in data:
        if state == "text":
            if value == 27:state = "escape"
            else:output.append(value)
        elif state == "escape":
            state = "csi" if value == 91 else "osc" if value == 93 else "charset" if value in (40, 41) else "text"
        elif state == "csi":
            if 64 <= value <= 126:state = "text"
            elif value == 27:state = "escape"
        elif state == "osc":
            if value == 7:state = "text"
            elif value == 27:state = "osc_escape"
        elif state == "osc_escape":
            state = "text" if value in (92, 7) else "osc"
        else:state = "text"
    return bytes(output), state


class SecretFilter:
    def __init__(self, secrets):
        values = set()
        for value in secrets:
            values.update(v.encode() for v in [value, *value.splitlines()] if len(v) >= 4)
        self.secrets = sorted(values, key=len, reverse=True)
        self.keep = max([256, *[len(v) for v in self.secrets]])
        self.pending = b""
        self.terminal_state = "text"

    def clean(self, data):
        data, _ = strip_terminal(data)
        for secret in self.secrets:
            data = data.replace(secret, b"[REDACTED]")
        return data

    def feed(self, data):
        plain, self.terminal_state = strip_terminal(data, self.terminal_state)
        self.pending += plain
        for secret in self.secrets:self.pending = self.pending.replace(secret, b"[REDACTED]")
        # Emit prompts without requiring a newline, retaining only a possible secret prefix.
        hold = 0
        for secret in self.secrets:
            for size in range(min(len(secret) - 1, len(self.pending)), hold, -1):
                if self.pending.endswith(secret[:size]):
                    hold = size
                    break
        cut = len(self.pending) - hold
        result, self.pending = self.pending[:cut], self.pending[cut:]
        return result

    def finish(self):
        partial = len(self.pending) >= 4 and any(s.startswith(self.pending) for s in self.secrets)
        result = b"[REDACTED]" if partial else self.clean(self.pending)
        self.pending = b""
        self.terminal_state = "text"
        return result


def read_secret_tables(tables):
    values = []
    for table in tables:
        path = Path(table["path"])
        column = table.get("key_column", 1)
        if type(column) is not int or not 0 <= column <= 8 or path.stat().st_size > 262144:
            raise ValueError("原认证表格式或大小不支持")
        found = False
        for line in path.read_text(encoding="utf-8").splitlines():
            if not line.strip() or line.lstrip().startswith("#"):
                continue
            fields = line.split()
            if table.get("label") and fields[0] != table["label"]:
                continue
            if len(fields) <= column:
                raise ValueError("原认证表缺少登记列")
            values.append(fields[column])
            found = True
        if not found:
            raise ValueError("原认证表缺少登记的测试身份")
    return values


def load_environment(plan):
    import pwd
    account = pwd.getpwnam(plan["user"])
    env = {"PATH": "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin", "HOME": account.pw_dir,
           "USER": account.pw_name, "LOGNAME": account.pw_name, "LANG": "C.UTF-8", "TERM": "xterm",
           "PYTHONUNBUFFERED": "1", "PYTHONDONTWRITEBYTECODE": "1"}
    secrets = read_secret_tables(plan.get("secret_tables", []))
    for path in plan.get("env_setup_files", []):
        # Only server-owned recipe paths reach this command; output stays in private process memory.
        command = ["/bin/bash", "-c", 'set -a; source "$1" >/dev/null; /usr/bin/env -0', "xas-env", path]
        result = subprocess.run(command, env=env, cwd=plan["cwd"], capture_output=True, timeout=10)
        if result.returncode:
            raise ValueError("原环境配置加载失败")
        env.update(dict(v.decode().split("=", 1) for v in result.stdout.split(b"\0") if b"=" in v))
        secrets.extend(v for k, v in env.items() if SENSITIVE_NAME.search(k) and v)
    for path in plan.get("credential_files", []):
        for line in Path(path).read_text().splitlines():
            if not line.strip() or line.lstrip().startswith("#"):
                continue
            match = re.match(r"^(?:export\s+)?([A-Za-z_][A-Za-z0-9_]*)=(.*)$", line.strip())
            if not match:
                raise ValueError("原认证配置格式不支持")
            parts = shlex.split(match[2], comments=True)
            if len(parts) > 1:
                raise ValueError("原认证配置格式不支持")
            env[match[1]] = parts[0] if parts else ""
            if SENSITIVE_NAME.search(match[1]) and env[match[1]]:secrets.append(env[match[1]])
    env.update(plan.get("environment", {}))
    secrets.extend(v for k, v in env.items() if SENSITIVE_NAME.search(k) and v)
    return env, secrets


def request(settings, payload):
    if not hasattr(socket, "AF_UNIX"):
        raise ValueError("原环境执行器仅在 Linux 服务器可用")
    try:
        with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
            connection.settimeout(15)
            connection.connect(str(settings.native_socket))
            connection.sendall(json.dumps(payload).encode() + b"\n")
            stream = connection.makefile("rb")
            data = stream.readline(131073)
            if len(data) > 131072:
                raise ValueError("原环境执行器响应过大")
            result = json.loads(data)
    except (OSError, json.JSONDecodeError) as exc:
        raise ValueError("原环境执行器暂不可用，请稍后重试") from exc
    if "error" in result:
        raise ValueError(result["error"])
    return result