claude-plugin/plugins/reviews/skills/audit-terraform/scripts/reference_set.py

"""Compute reference sets for the consistency reviewer."""
from __future__ import annotations

import re
from collections import Counter, defaultdict
from pathlib import Path

from scripts.manifest import ModuleGraphEntry


def _peer_dirs_at_depth(repo_root: Path, dir_path: str) -> list[str]:
    parts = dir_path.split("/")
    depth = len(parts)
    parent = repo_root.joinpath(*parts[:-1]) if depth > 1 else repo_root
    if not parent.is_dir():
        return []
    out: list[str] = []
    for p in sorted(parent.iterdir()):
        if not p.is_dir():
            continue
        rel = p.relative_to(repo_root).as_posix()
        if rel != dir_path:
            out.append(rel)
    return out


def _terragrunt_refs(repo_root: Path, dir_path: str) -> list[str]:
    """Region peers + same-component cross-env (layout: live/<env>/<region>/<component>)."""
    parts = dir_path.split("/")
    if len(parts) != 4:
        return _peer_dirs_at_depth(repo_root, dir_path)
    live, env, region, component = parts
    refs: set[str] = set()
    region_dir = repo_root / live / env / region
    if region_dir.is_dir():
        for p in region_dir.iterdir():
            if p.is_dir() and (p / "terragrunt.hcl").is_file():
                rel = p.relative_to(repo_root).as_posix()
                if rel != dir_path:
                    refs.add(rel)
    live_dir = repo_root / live
    if live_dir.is_dir():
        for env_dir in live_dir.iterdir():
            candidate = env_dir / region / component
            if candidate.is_dir() and (candidate / "terragrunt.hcl").is_file():
                rel = candidate.relative_to(repo_root).as_posix()
                if rel != dir_path:
                    refs.add(rel)
    env_live_dir = repo_root / live / env
    if env_live_dir.is_dir():
        for region_dir2 in env_live_dir.iterdir():
            candidate = region_dir2 / component
            if candidate.is_dir() and (candidate / "terragrunt.hcl").is_file():
                rel = candidate.relative_to(repo_root).as_posix()
                if rel != dir_path:
                    refs.add(rel)
    return sorted(refs)


def compute_reference_sets(
    repo_root: Path | str,
    changed_dirs: set[str],
    module_graph: dict[str, ModuleGraphEntry],
) -> dict[str, list[str]]:
    root = Path(repo_root)
    out: dict[str, list[str]] = {}
    for d in sorted(changed_dirs):
        if d.startswith("modules/"):
            entry = module_graph.get(d)
            out[d] = list(entry.sibling_modules_at_callsites) if entry else []
            continue
        full = root / d
        if (full / "terragrunt.hcl").is_file():
            out[d] = _terragrunt_refs(root, d)
            continue
        out[d] = _peer_dirs_at_depth(root, d)
    return out


_ATTR_RE = re.compile(r"^\s*(?P<key>[A-Za-z0-9_]+)\s*=")
_MODULE_RE = re.compile(r'^\s*module\s+"(?P<name>[^"]+)"\s*\{')


def _iter_tf_lines(dir_path: Path) -> list[str]:
    lines: list[str] = []
    for tf in sorted(dir_path.glob("*.tf")):
        lines.extend(tf.read_text().splitlines())
    return lines


def _dir_signature(dir_path: Path) -> tuple[set[str], set[str]]:
    attrs: set[str] = set()
    modules: set[str] = set()
    for line in _iter_tf_lines(dir_path):
        attr_match = _ATTR_RE.match(line)
        if attr_match:
            attrs.add(attr_match.group("key"))
        module_match = _MODULE_RE.match(line)
        if module_match:
            modules.add(module_match.group("name"))
    return attrs, modules


def compute_consistency_norms(
    repo_root: Path | str,
    reference_sets: dict[str, list[str]],
) -> dict[str, dict[str, list[dict[str, object]]]]:
    root = Path(repo_root)
    out: dict[str, dict[str, list[dict[str, object]]]] = {}
    for changed_dir, refs in sorted(reference_sets.items()):
        attr_support: dict[str, list[str]] = defaultdict(list)
        module_support: dict[str, list[str]] = defaultdict(list)
        naming_tokens: Counter[str] = Counter()
        usable_refs: list[str] = []
        for ref in refs:
            ref_path = root / ref
            if not ref_path.is_dir():
                continue
            attrs, modules = _dir_signature(ref_path)
            if not attrs and not modules:
                continue
            usable_refs.append(ref)
            for attr in attrs:
                attr_support[attr].append(ref)
            for module in modules:
                module_support[module].append(ref)
            naming_tokens.update(Path(ref).parts[-1:])

        out[changed_dir] = {
            "attribute_norms": [
                {"attribute": attr, "peer_dirs": sorted(peer_dirs)}
                for attr, peer_dirs in sorted(attr_support.items())
                if len(peer_dirs) >= 2
            ],
            "module_wrapper_norms": [
                {"module": module, "peer_dirs": sorted(peer_dirs)}
                for module, peer_dirs in sorted(module_support.items())
                if len(peer_dirs) >= 2
            ],
            "naming_norms": [
                {"token": token, "peer_dirs": sorted(usable_refs)}
                for token, count in sorted(naming_tokens.items())
                if count >= 2
            ],
        }
    return out