scripts/broker_sdk_environment.py
Rama main · réplica de lectura del repositorio original.
"""Repair the two source-pinned broker SDKs before runtime attestation.
Standalone/stdlib-only: Docker and pip-only setup run this before workspace
packages are importable. The compatibility Kotak entry point stays available.
"""
from __future__ import annotations
import json
import re
import shutil
import subprocess
import sys
import tomllib
from collections.abc import Callable
from pathlib import Path, PurePosixPath
from typing import Any
REPO = Path(__file__).resolve().parents[1]
BUILD_REQUIREMENTS_LOCK = REPO / "broker-sdk-build.lock"
KOTAK_REPO = "https://github.com/Kotak-Neo/kotak-neo-python.git"
_KOTAK_HOMEPAGE = KOTAK_REPO.removesuffix(".git")
_FULL_GIT_SHA = re.compile(r"[0-9a-f]{40}")
_IDENTITIES = {
"kotakneoapi": ("neo_api_client", KOTAK_REPO),
"upstox-python-sdk": ("upstox_client", "https://github.com/upstox/upstox-python.git"),
}
# Run inside the target interpreter. Never import an unverified SDK to decide
# whether its bytes are healthy. RECORD supplies ownership, not file authority.
_PROBE = r"""
import contextlib
import hashlib
import importlib
import importlib.metadata as md
import importlib.util
import io
import json
import os
from pathlib import Path
import re
import sys
pins = json.load(sys.stdin)
identities = {"kotakneoapi": "neo_api_client", "upstox-python-sdk": "upstox_client"}
result = {"owners": {module: [] for module in identities.values()}}
dists = list(md.distributions())
for dist in dists:
name = re.sub(r"[-_.]+", "-", str(dist.metadata.get("Name", "")).lower())
paths = [str(file).replace("\\", "/") for file in (dist.files or ())]
for module in identities.values():
if module in (dist.read_text("top_level.txt") or "").splitlines() or any(p.split("/")[0] == module for p in paths):
result["owners"][module].append(name)
if name not in {*identities, "neo-api-client"}:
continue
if name in result:
result[name]["duplicate"] = True
continue
try:
direct = json.loads(dist.read_text("direct_url.json") or "null")
except ValueError:
direct = None
root = Path(str(dist.locate_file(""))).resolve()
safe_remove = bool(paths) and all(not Path(p).is_absolute() and ".." not in Path(p).parts
and (root/p).resolve().is_relative_to(root) for p in paths)
row = {"version": dist.version, "direct_url": direct, "safe_remove": safe_remove,
"physical_ok": False, "import_ok": False, "root": str(root)}
result[name] = row
pin = pins.get(name)
if not pin:
continue
module = identities[name]
package = root/module
files = pin["build_files"]
try:
expected = {p for p in files if p.startswith(module + "/")}
actual = {p.relative_to(root).as_posix() for p in package.rglob("*")
if p.is_file() and "__pycache__" not in p.parts and p.suffix != ".pyc"}
links = [p for p in package.rglob("*") if p.is_symlink()]
hashes = {p: hashlib.sha256((root/p).read_bytes()).hexdigest() for p in files
if (root/p).is_file() and not (root/p).is_symlink()
and (root/p).resolve().is_relative_to(root)}
row["physical_ok"] = not package.is_symlink() and not links and actual == expected and hashes == files
spec = importlib.util.find_spec(module)
row["origin"] = spec.origin if spec else None
if row["physical_ok"] and spec and spec.origin and Path(spec.origin).resolve() == package/"__init__.py":
# Verified official modules can print banners; never let that turn
# the JSON protocol into a successful-looking partial observation.
# Health probes must not create the SDK's default rotating log.
# Only this disposable child environment changes, never the caller.
os.environ["NEO_LOG_FILE_ENABLED"] = "false"
with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()):
imported = importlib.import_module(module)
row["import_ok"] = Path(imported.__file__).resolve() == package/"__init__.py"
except (OSError, ValueError, TypeError, ImportError, AttributeError) as exc:
row["probe_error"] = type(exc).__name__
result["namespace_owners"] = result["owners"]["neo_api_client"]
print(json.dumps(result))
"""
def _pin() -> tuple[str, str]:
"""Compatibility metadata accessor; repair uses the complete source pin."""
data = tomllib.loads((REPO / "brokers.lock").read_text(encoding="utf-8"))
for entry in data.get("broker", []):
if entry.get("name") == "kotakneoapi":
version = str(entry.get("version", ""))
commit = str(entry.get("source_commit", ""))
if not version or entry.get("homepage") != _KOTAK_HOMEPAGE:
raise RuntimeError("Kotak Neo brokers.lock entry does not identify the official repository")
if _FULL_GIT_SHA.fullmatch(commit) is None:
raise RuntimeError("Kotak Neo source_commit must be a full lowercase Git commit")
return version, commit
raise RuntimeError("Kotak Neo is missing from brokers.lock")
def _pins(names: tuple[str, ...]) -> dict[str, dict[str, Any]]:
entries = tomllib.loads((REPO / "brokers.lock").read_text(encoding="utf-8")).get("broker", [])
pins = {}
for name in names:
selected = [p for p in entries if p.get("name") == name]
if len(selected) != 1:
raise RuntimeError(f"Expected one source pin for {name}")
pin = selected[0]
module, url = _IDENTITIES[name]
if (
pin.get("source_kind") != "git"
or pin.get("source_url") != url
or pin.get("homepage") != url.removesuffix(".git")
or pin.get("source_module") != module
or not pin.get("version")
):
raise RuntimeError(f"{name} does not identify the official repository/source module")
if any(_FULL_GIT_SHA.fullmatch(str(pin.get(k, ""))) is None for k in ("source_commit", "source_tree")):
raise RuntimeError(f"{name} requires a full lowercase Git commit and tree")
files = pin.get("build_files")
if (
not re.fullmatch(r"[0-9a-f]{64}", str(pin.get("build_wheel_sha256", "")))
or type(files) is not dict
or f"{module}/__init__.py" not in files
):
raise RuntimeError(f"{name} has no locked physical build identity")
for relative, digest in files.items():
if (
type(relative) is not str
or not relative
or "\\" in relative
or PurePosixPath(relative).is_absolute()
or ".." in PurePosixPath(relative).parts
or PurePosixPath(relative).as_posix() != relative
or not re.fullmatch(r"[0-9a-f]{64}", str(digest))
):
raise RuntimeError(f"{name} has an invalid locked build resource")
pins[name] = pin
# Consume the exact resolver export, not an independently invented VCS line.
requirements = [
line
for line in (REPO / "broker-sdk-git-requirements.txt").read_text().splitlines()
if line and not line.startswith("#")
]
all_pins = {p.get("name"): p for p in entries if p.get("name") in _IDENTITIES}
expected = [
f"{name} @ git+{url}@{all_pins.get(name, {}).get('source_commit', '')}"
for name, (_, url) in _IDENTITIES.items()
]
if sorted(requirements) != sorted(expected):
raise RuntimeError("Git requirements must contain exactly the two official full-commit pins")
return pins
def _invoke(
run: Callable[..., subprocess.CompletedProcess], args: list[str], *, input_text: str | None = None,
) -> subprocess.CompletedProcess:
options = {"input": input_text} if input_text is not None else {}
result = run(args, cwd=REPO, capture_output=True, text=True, check=False, **options)
if result.returncode != 0:
raise RuntimeError(
f"Broker SDK environment command failed: {args[0]} {args[1]} "
f"(exit {result.returncode}): {result.stderr.strip()}"
)
return result
def _probe(
python: Path, run: Callable[..., subprocess.CompletedProcess], pins: dict[str, dict[str, Any]] | None = None
) -> dict[str, Any]:
result = _invoke(run, [str(python), "-I", "-B", "-c", _PROBE], input_text=json.dumps(pins or {}))
try:
state = json.loads(result.stdout)
except (TypeError, ValueError) as exc:
raise RuntimeError("Broker SDK distribution probe returned invalid JSON") from exc
if not isinstance(state, dict):
raise RuntimeError("Broker SDK distribution probe returned invalid state")
return state
def _healthy_source(state: dict[str, Any], name: str, version: str, commit: str) -> bool:
module, url = _IDENTITIES[name]
owners = state.get("owners", {}).get(module)
if name == "kotakneoapi":
owners = state.get("namespace_owners", owners)
if "neo-api-client" in state:
return False
installed = state.get(name)
if (
owners != [name]
or not isinstance(installed, dict)
or installed.get("version") != version
or installed.get("duplicate")
or installed.get("physical_ok") is not True
or installed.get("import_ok") is not True
):
return False
direct = installed.get("direct_url")
if not isinstance(direct, dict) or direct.get("url") != url:
return False
vcs = direct.get("vcs_info")
return (
isinstance(vcs, dict)
and vcs.get("vcs") == "git"
and vcs.get("commit_id") == commit
and vcs.get("requested_revision") == commit
)
def _healthy(state: dict[str, Any], version: str, commit: str) -> bool:
return _healthy_source(state, "kotakneoapi", version, commit)
def _check_owners(state: dict[str, Any], names: tuple[str, ...]) -> None:
for name in names:
module = _IDENTITIES[name][0]
allowed = {name, "neo-api-client"} if name == "kotakneoapi" else {name}
owners = state.get("owners", {}).get(module, state.get("namespace_owners", []) if name == "kotakneoapi" else [])
if not isinstance(owners, list) or set(owners) - allowed:
raise RuntimeError(f"Refusing unrelated owner of {module}; no distributions changed")
for selected in allowed:
if selected in state and (
state[selected].get("duplicate") or state[selected].get("safe_remove") is not True
):
raise RuntimeError(f"Refusing ambiguous/unsafe removal of {selected}; no distributions changed")
def _commands(python: Path) -> tuple[list[str], list[str], str]:
uv = shutil.which("uv")
if uv:
return (
[uv, "pip", "install", "--python", str(python), "--break-system-packages"],
[uv, "pip", "uninstall", "--python", str(python), "--break-system-packages"],
"--reinstall",
)
return (
[str(python), "-m", "pip", "install", "--break-system-packages"],
[str(python), "-m", "pip", "uninstall", "-y", "--break-system-packages"],
"--force-reinstall",
)
def remove_kotak_distributions(
python: Path, *, run: Callable[..., subprocess.CompletedProcess] = subprocess.run
) -> None:
"""Legacy explicit removal; never uninstall an unrelated namespace owner."""
state = _probe(python, run)
_check_owners(state, ("kotakneoapi",))
present = [name for name in ("kotakneoapi", "neo-api-client") if name in state]
if present:
_invoke(run, [*_commands(python)[1], *present])
def _repair(python: Path, names: tuple[str, ...], run: Callable[..., subprocess.CompletedProcess]) -> None:
pins = _pins(names)
state = _probe(python, run, pins)
stale = [
name for name, pin in pins.items() if not _healthy_source(state, name, pin["version"], pin["source_commit"])
]
if not stale:
return # Healthy repeat is strictly probe-only.
_check_owners(state, names) # Preflight the whole requested union before writes.
install, uninstall, reinstall = _commands(python)
_invoke(run, [*install, "--require-hashes", "--only-binary=:all:", "--no-deps", "-r", str(BUILD_REQUIREMENTS_LOCK)])
for name in stale:
removable = (name, "neo-api-client") if name == "kotakneoapi" else (name,)
present = [n for n in removable if n in state]
if present:
_invoke(run, [*uninstall, *present])
pin = pins[name]
requirement = f"{name} @ git+{pin['source_url']}@{pin['source_commit']}"
_invoke(run, [*install, "--no-build-isolation", "--no-deps", reinstall, requirement])
final = _probe(python, run, pins)
if any(not _healthy_source(final, name, pin["version"], pin["source_commit"]) for name, pin in pins.items()):
raise RuntimeError(
"Broker SDK repair did not produce the pinned Git distribution with exclusive namespace and physical/import proof"
)
def repair_broker_sdk_environment(
python: Path, *, run: Callable[..., subprocess.CompletedProcess] = subprocess.run
) -> None:
"""Repair BOTH exact Git SDKs with hashed build tools and no dependency solve."""
_repair(python, tuple(_IDENTITIES), run)
def repair_kotakneo_environment(
python: Path, *, run: Callable[..., subprocess.CompletedProcess] = subprocess.run
) -> None:
"""Backward-compatible Kotak-only API, with the same physical source checks."""
_repair(python, ("kotakneoapi",), run)
def main(argv: list[str] | None = None) -> int:
args = argv if argv is not None else sys.argv[1:]
if args == ["repair"]:
repair_broker_sdk_environment(Path(sys.executable))
elif args == ["remove"]:
remove_kotak_distributions(Path(sys.executable))
else:
raise SystemExit("usage: broker_sdk_environment.py repair|remove")
return 0
if __name__ == "__main__":
raise SystemExit(main())