rmi-backend/scripts/fix_security_lints.py

#!/usr/bin/env python3
"""Bulk-fix S110, S324, and S311 ruff violations."""

from __future__ import annotations

import argparse
import json
import os
import py_compile
import re
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
DEFAULT_TARGETS = [ROOT / "app"]

parser = argparse.ArgumentParser(description="Fix S110/S324/S311 ruff violations")
parser.add_argument("targets", nargs="*", type=Path, default=DEFAULT_TARGETS)
ARGS = parser.parse_args()
TARGETS = [t if t.is_absolute() else ROOT / t for t in ARGS.targets]

S110_LOG_LINE = 'logging.getLogger(__name__).warning("swallowed exception", exc_info=True)'


def run_ruff(select: str, targets: list[Path] | None = None) -> list[dict]:
    targets = targets or TARGETS
    cmd = [
        sys.executable,
        "-m",
        "ruff",
        "check",
        "--select",
        select,
        "--output-format",
        "json",
        *[str(t) for t in targets],
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode not in (0, 1):
        print(f"ruff {select} failed: {result.stderr}", file=sys.stderr)
        return []
    return json.loads(result.stdout or "[]")


def ensure_import(text: str, module: str) -> str:
    """Add `import module` at the top-level import block if missing."""
    if re.search(rf"(?:^|\n)import\s+{re.escape(module)}(?:\s|$)", text):
        return text
    if re.search(rf"(?:^|\n)from\s+{re.escape(module)}\s+import", text):
        return text

    lines = text.splitlines(keepends=True)
    insert = 0
    paren = 0
    for i, ln in enumerate(lines):
        # track parentheses so we don't insert inside a multi-line parenthesized import
        paren += ln.count("(") - ln.count(")")
        if ln.startswith("#!/") or ln.startswith("# -*-") or ln.startswith("# coding"):
            insert = i + 1
            continue
        if ln.startswith("from __future__ import"):
            insert = i + 1
            continue
        if paren == 0 and (ln.startswith("import ") or ln.startswith("from ")):
            insert = i + 1
    if insert == len(lines) and lines and not lines[-1].endswith("\n"):
        lines.append("\n")
    lines.insert(insert, f"import {module}\n")
    return "".join(lines)


def remove_import(text: str, module: str) -> str:
    """Remove a top-level `import module` line if it exists."""
    lines = text.splitlines(keepends=True)
    out = []
    for ln in lines:
        if re.match(rf"^import\s+{re.escape(module)}\s*$", ln.strip("\n")):
            continue
        out.append(ln)
    return "".join(out)


# ---------------------------------------------------------------------------
# S110
# ---------------------------------------------------------------------------
SAME_LINE_PASS = re.compile(
    r"^(?P<indent>[ \t]*)except\b(?P<expr>[^:\n]*):(?P<space>\s*)pass\s*(?:#.*)?$",
    re.MULTILINE,
)


def fix_s110(targets: list[Path] | None = None) -> tuple[int, int]:
    errors = run_ruff("S110", targets)
    by_file: dict[str, list[int]] = {}
    for e in errors:
        by_file.setdefault(e["filename"], []).append(e["location"]["row"])

    changed = 0
    for path_str, rows in by_file.items():
        path = Path(path_str)
        text = path.read_text()
        orig = text

        lines = text.splitlines(keepends=True)
        # sort rows descending so line numbers stay stable
        for row in sorted(rows, reverse=True):
            if row < 1 or row > len(lines):
                continue
            idx = row - 1
            line = lines[idx]
            m = SAME_LINE_PASS.match(line)
            if m:
                indent = m.group("indent")
                expr = m.group("expr").strip()
                new_lines = [f"{indent}except {expr}:\n", f"{indent}    {S110_LOG_LINE}\n"]
                lines[idx : idx + 1] = new_lines
                continue

            # multiline: find the next non-comment/blank line that is a bare pass
            j = idx + 1
            while j < len(lines):
                content = lines[j].strip()
                if content == "" or content.startswith("#"):
                    j += 1
                    continue
                if content == "pass" or content.startswith("pass ") or content.startswith("pass\t"):
                    indent = len(lines[j]) - len(lines[j].lstrip())
                    lines[j] = " " * indent + S110_LOG_LINE + "\n"
                    break
                # next executable line is not a pass; give up
                break

        text = "".join(lines)
        # ensure logging is imported after all replacements so row indices remain valid
        text = ensure_import(text, "logging")
        if text != orig:
            path.write_text(text)
            changed += 1

    return len(errors), changed


# ---------------------------------------------------------------------------
# S324
# ---------------------------------------------------------------------------
MD5_RE = re.compile(r"hashlib\.md5\s*\(")
SHA1_RE = re.compile(r"hashlib\.sha1\s*\(")
NEW_MD5_RE = re.compile(r"hashlib\.new\s*\(\s*(['\"])md5\1\s*\)")
NEW_SHA1_RE = re.compile(r"hashlib\.new\s*\(\s*(['\"])sha1\1\s*\)")


def fix_s324(targets: list[Path] | None = None) -> tuple[int, int]:
    errors = run_ruff("S324", targets)
    by_file: dict[str, list[dict]] = {}
    for e in errors:
        by_file.setdefault(e["filename"], []).append(e)

    changed = 0
    for path_str in by_file:
        path = Path(path_str)
        text = path.read_text()
        orig = text
        text = MD5_RE.sub("hashlib.sha256(", text)
        text = SHA1_RE.sub("hashlib.sha256(", text)
        text = NEW_MD5_RE.sub(r"hashlib.new(\1sha256\1)", text)
        text = NEW_SHA1_RE.sub(r"hashlib.new(\1sha256\1)", text)
        if text != orig:
            path.write_text(text)
            changed += 1

    return len(errors), changed


# ---------------------------------------------------------------------------
# S311
# ---------------------------------------------------------------------------
RANDOM_CALL_RE = re.compile(r"random\.((?:random|randint|uniform|randrange))\s*\(")
RANDOM_CHOICE_RE = re.compile(r"random\.choice\s*\(")


def fix_s311(targets: list[Path] | None = None) -> tuple[int, int]:
    errors = run_ruff("S311", targets)
    by_file: dict[str, list[dict]] = {}
    for e in errors:
        by_file.setdefault(e["filename"], []).append(e)

    changed = 0
    for path_str in by_file:
        path = Path(path_str)
        text = path.read_text()
        orig = text

        text = RANDOM_CALL_RE.sub(r"secrets.SystemRandom().\1(", text)
        text = RANDOM_CHOICE_RE.sub("secrets.choice(", text)

        if text != orig:
            text = ensure_import(text, "secrets")
            # remove `import random` only if no other random.* calls remain
            if not re.search(r"\brandom\.\w+\s*\(", text):
                text = remove_import(text, "random")
            path.write_text(text)
            changed += 1

    return len(errors), changed


# ---------------------------------------------------------------------------
# verification
# ---------------------------------------------------------------------------
def verify_changed(files: list[Path]) -> list[str]:
    bad = []
    for p in files:
        try:
            py_compile.compile(str(p), doraise=True)
        except py_compile.PyCompileError as e:
            bad.append(f"{p}: {e}")
    return bad


def main() -> None:
    os.chdir(ROOT)
    targets = TARGETS
    s110_n, s110_files = fix_s110(targets)
    s324_n, s324_files = fix_s324(targets)
    s311_n, s311_files = fix_s311(targets)
    print(f"S110: {s110_n} errors -> touched {s110_files} files")
    print(f"S324: {s324_n} errors -> touched {s324_files} files")
    print(f"S311: {s311_n} errors -> touched {s311_files} files")

    # compile check every touched file
    changed = set()
    # collect all files that were touched via git status
    proc = subprocess.run(
        ["git", "status", "--short"] + [str(t) for t in targets], capture_output=True, text=True
    )
    for ln in proc.stdout.splitlines():
        if ln.startswith(" M ") or ln.startswith("M  "):
            changed.add(ROOT / ln[3:].strip())
    if changed:
        bad = verify_changed(list(changed))
        if bad:
            print("\nPY_COMPILE ERRORS:", file=sys.stderr)
            for b in bad:
                print(b, file=sys.stderr)
            sys.exit(1)
        print(f"py_compile OK for {len(changed)} changed files")


if __name__ == "__main__":
    main()