SK-3156-addressed-review-comments

Author: skyflow-himanshupal

.githooks/pre-commit

@@ -3,7 +3,7 @@
 #   Tier 1: auto-redact gitleaks findings inside the Fern-generated code
 #           trees (driven live by .gitleaks.toml via
 #           scripts/patch_generated_secrets.py, not a hand-maintained list)
-#           and re-stage those directories. Never blocks the commit by
+#           and re-stage the files it touched. Never blocks the commit by
 #           itself - it either fixes generated code or leaves it untouched
 #           for tier 2 to catch.
 #   Tier 2: run the real gitleaks scan against the staged diff and block on
@@ -21,6 +21,20 @@ GENERATED_DIRS=(
   "flowvault/skyflow/generated"
 )
 
+# Snapshot what's already staged in the generated dirs before tier 1 runs,
+# so an intentionally-unstaged, in-progress change there (e.g. mid-way
+# through testing a Fern regen) isn't unconditionally swept into this
+# commit. Tier 1 only edits working-tree files, never the index, so this
+# snapshot stays accurate regardless of what it does next.
+already_staged_in_generated="$(git diff --cached --name-only -- "${GENERATED_DIRS[@]}" || true)"
+
+git_dir="$(git rev-parse --git-dir)"
+touched_file_list="$git_dir/leak-guard-touched-files.txt"
+# A list from a previous commit's tier 1 run must never be reused here - if
+# tier 1 doesn't run this time (python3 missing, below), stale entries from
+# that earlier run would otherwise get staged again.
+rm -f "$touched_file_list"
+
 # git invokes hooks with a leaner PATH than your interactive shell, so a
 # Python install managed by pyenv/conda/asdf (rather than a system package)
 # is often invisible here even though `python3` works fine in your
@@ -46,12 +60,24 @@ else
   fi
 fi
 
-# Stages the whole generated-code directories rather than just the files
-# tier 1 touched, since we don't get that list back from the script. Safe
-# in practice: these directories are machine-owned (Fern-generated, never
-# hand edited), so there's no legitimate "leave part of it unstaged" case to
-# worry about disturbing.
-git add -- "${GENERATED_DIRS[@]}"
+# Stages only what tier 1 actually touched this run (recorded to
+# $touched_file_list - see record_touched_files in the script) plus
+# whatever was already staged above. A file sitting in one of
+# GENERATED_DIRS that tier 1 didn't touch and the caller hadn't staged is
+# left alone rather than force-added.
+to_stage=()
+if [ -f "$touched_file_list" ]; then
+  while IFS= read -r f; do
+    [ -n "$f" ] && to_stage+=("$f")
+  done < "$touched_file_list"
+fi
+while IFS= read -r f; do
+  [ -n "$f" ] && to_stage+=("$f")
+done <<< "$already_staged_in_generated"
+
+if [ "${#to_stage[@]}" -gt 0 ]; then
+  git add -- "${to_stage[@]}"
+fi
 
 if ! command -v gitleaks >/dev/null 2>&1; then
   echo "[leak-guard] tier 2 skipped: 'gitleaks' binary not found locally."

scripts/patch_generated_secrets.py

@@ -163,6 +163,49 @@ def placeholder_for(rule_id: str) -> str:
     return f"<REDACTED_{re.sub(r'[^A-Z0-9]+', '_', rule_id.upper())}>"
 
 
+def _git_dir() -> Path:
+    output = subprocess.run(
+        ["git", "rev-parse", "--git-dir"],
+        cwd=REPO_ROOT,
+        capture_output=True,
+        text=True,
+        check=True,
+    ).stdout.strip()
+    return (REPO_ROOT / output).resolve()
+
+
+def record_touched_files(relative_files) -> None:
+    """Records exactly which files this run modified (an empty list if none).
+
+    Lets the pre-commit hook stage only those files - plus whatever was
+    already staged - instead of the entire generated-code directories,
+    which would otherwise sweep in unrelated or intentionally-unstaged
+    in-progress changes sitting in those same directories. Called at every
+    exit point so the hook never reads a stale list left over from a prior
+    run. Best-effort: if this can't be written, the hook simply finds no
+    list and stages nothing beyond what was already staged.
+    """
+    try:
+        list_path = _git_dir() / "leak-guard-touched-files.txt"
+        list_path.write_text(
+            "".join(f"{f}\n" for f in relative_files), encoding="utf-8"
+        )
+    except (OSError, subprocess.CalledProcessError):
+        pass
+
+
+def _rollback(original_contents: dict) -> None:
+    """Restores every file this run has written so far back to what it read
+    at the start. Called from every failure path after files start getting
+    written, so a build break, a final-scan tool failure, or leftover
+    findings after redaction never leaves a partially-redacted, unverified
+    file sitting in the working tree.
+    """
+    for file_path, content in original_contents.items():
+        file_path.write_text(content, encoding="utf-8")
+    record_touched_files([])
+
+
 def _line_start_offsets(text: str) -> list:
     """Absolute char offset where each 1-indexed line starts in `text`.
 
@@ -213,6 +256,7 @@ def main() -> int:
             "generated code. CI will still scan for this.",
             file=sys.stderr,
         )
+        record_touched_files([])
         return 0
 
     if not self_test_allowlist_support():
@@ -225,11 +269,13 @@ def main() -> int:
             "zricethezav/gitleaks:latest) and try again.",
             file=sys.stderr,
         )
+        record_touched_files([])
         return 1
 
     findings = scan_generated_dirs()
     if not findings:
         print("No gitleaks findings in generated code.")
+        record_touched_files([])
         return 0
 
     allowed_prefixes = tuple(GENERATED_DIRS)
@@ -243,6 +289,7 @@ def main() -> int:
             + "\n".join(f"  - {f['File']}" for f in out_of_scope),
             file=sys.stderr,
         )
+        record_touched_files([])
         return 1
 
     by_file = {}
@@ -284,11 +331,21 @@ def main() -> int:
                 continue
             spans.append((span[0], span[1], finding["RuleID"]))
 
-        # Deduplicate by (start, end): if two rules flag the exact same
-        # span, it must only be redacted once. Applied back-to-front
-        # (highest offset first) so replacing a later span never shifts
-        # the offsets of an earlier one still waiting to be processed.
-        unique_spans_by_file[file_path] = sorted(set(spans), reverse=True)
+        # Deduplicate by (start, end) only - not the full (start, end,
+        # rule_id) triple, which wouldn't collapse two different rules
+        # flagging the exact same span; a stale second replacement at an
+        # already-redacted offset would then corrupt the file or eat
+        # adjacent text. Keeps whichever rule_id was seen first for that
+        # span. Applied back-to-front (highest offset first) so replacing
+        # a later span never shifts the offsets of an earlier one still
+        # waiting to be processed.
+        spans_by_range = {}
+        for start, end, rule_id in spans:
+            spans_by_range.setdefault((start, end), rule_id)
+        unique_spans_by_file[file_path] = sorted(
+            ((start, end, rule_id) for (start, end), rule_id in spans_by_range.items()),
+            reverse=True,
+        )
 
     if unresolved:
         print(
@@ -300,9 +357,11 @@ def main() -> int:
             ),
             file=sys.stderr,
         )
+        record_touched_files([])
         return 1
 
     original_contents = dict(file_contents)
+    touched_files = set()
     redacted_count = 0
     for file_path, content in file_contents.items():
         by_rule = {}
@@ -317,6 +376,8 @@ def main() -> int:
             print(f"[{rule_id}] redacted in {file_path.relative_to(REPO_ROOT)} -> {placeholder}")
 
         file_path.write_text(content, encoding="utf-8")
+        if content != original_contents[file_path]:
+            touched_files.add(str(file_path.relative_to(REPO_ROOT)))
 
     build_ok = True
     for file_path in original_contents:
@@ -330,26 +391,43 @@ def main() -> int:
             )
 
     if not build_ok:
-        for file_path, content in original_contents.items():
-            file_path.write_text(content, encoding="utf-8")
+        _rollback(original_contents)
         print(
             "Rolled back all changes from this run - redaction needs manual review.",
             file=sys.stderr,
         )
         return 1
 
-    remaining = scan_generated_dirs()
+    # The final re-scan can itself fail for reasons other than "still
+    # leaks" (e.g. the gitleaks binary crashing mid-run) -
+    # run_gitleaks_detect (via scan_generated_dirs) raises in that case
+    # rather than returning a findings list. Either way, an unverified
+    # redaction must never be left on disk: roll back exactly as the
+    # build-failure path above does.
+    try:
+        remaining = scan_generated_dirs()
+    except RuntimeError as err:
+        _rollback(original_contents)
+        print(
+            f"Final gitleaks re-scan failed to run ({err}) - rolled back all "
+            "changes from this run. Redaction needs manual review.",
+            file=sys.stderr,
+        )
+        return 1
     if remaining:
+        _rollback(original_contents)
         print(
             f"Redacted {redacted_count} secret(s), but {len(remaining)} finding(s) "
-            "remain after re-scanning. Manual review needed:\n"
+            "remain after re-scanning. Rolled back all changes from this run. "
+            "Manual review needed:\n"
             + "\n".join(
                 f"  - [{f['RuleID']}] {f['File']}:{f['StartLine']}" for f in remaining
             ),
             file=sys.stderr,
         )
         return 1
 
+    record_touched_files(sorted(touched_files))
     print(
         f"\nDone. Redacted {redacted_count} secret(s) across {len(by_file)} file(s). "
         "Build and gitleaks re-scan both clean."