fef2's picture
Upload folder using huggingface_hub
0936c1f verified
Raw
History Blame Contribute Delete
38.1 kB
from __future__ import annotations
import argparse
import json
import math
import multiprocessing
import os
import re
import sys
import warnings
from collections import Counter
from dataclasses import dataclass
from pathlib import Path
from typing import Literal
from ._shim import LIVE_PREFIX_RE, PLACEHOLDER_RE, REFERENCE_RE, SLOT_RE, add_common_args, is_inert_value, nfc, read_jsonl, seeded_rng, stable_id, write_jsonl, write_report
MARKER_RE = re.compile(r"\[REDACTED:([A-Z_]+)\]")
KV_RE = re.compile(r"(?im)(?P<key>[A-Za-zА-Яа-яЁё][\w.\-/ ]{0,64}?)[\"']?\s*(?P<sep>[:=])\s*(?:(?P<quote>[\"'])(?P<quoted_value>[^\n\"']*)(?P=quote)|[«“](?P<smart_value>[^\n»”]*)[»”]|(?P<bare_value>[^\s,;&\]\}\n]+))")
EMAIL_RE = re.compile(r"\b[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}\b")
PHONE_RE = re.compile(r"(?<!\w)(?:\+7|8)[\s(\-]*\d{3}[\s)\-]*\d{3}[\s\-]*\d{2}[\s\-]*\d{2}(?!\w)")
INN_RE = re.compile(r"(?iu)\b(?:инн|inn)\D{0,12}(\d{10}(?:\d{2})?)")
SNILS_RE = re.compile(r"(?<!\d)\d{3}-\d{3}-\d{3}\s?\d{2}(?!\d)")
CARD_RE = re.compile(r"(?<!\d)(?:\d[ -]?){15}\d(?!\d)")
PASSPORT_RE = re.compile(r"(?iu)\b(?:паспорт|passport)\D{0,12}(\d{4}\s?\d{6})")
CONTRACT_RE = re.compile(r"(?iu)\b(?:договор|контракт|соглашени[ея])\s*(?:№|N[oо.]?)?\s*([A-ZА-Я0-9][A-ZА-Я0-9\-/]{7,})")
OPAQUE_RE = re.compile(r"(?<![A-Za-z0-9+/=_-])[A-Za-z0-9+/=_-]{16,}(?![A-Za-z0-9+/=_-])")
GitleaksRule = "tuple[re.Pattern[str], str | None, tuple[str, ...], int, float]"
_ONTOLOGY_PATH = Path(__file__).with_name("slot_ontology.yaml")
_ONTOLOGY = json.loads(_ONTOLOGY_PATH.read_text(encoding="utf-8"))
MARKER_LABELS = {"PASSWORD": "PASSWORD", "TOKEN": "AUTH_TOKEN", "AUTH_TOKEN": "AUTH_TOKEN", "API_KEY": "SECRET_KEY", "SECRET": "SECRET_KEY", "KEY": "SECRET_KEY", "LOGIN": "LOGIN", "USERNAME": "LOGIN", "USER": "LOGIN", "CONTRACT": "CONTRACT_NUMBER"}
_PROJECT_PII_RECOGNIZERS: list[object] | None = None
SCRUBBER_POLICY_VERSION = "cli-credentials-v2"
@dataclass(frozen=True)
class CredentialSite:
start: int
end: int
label: str
detector: str
disposition: str
@dataclass(frozen=True)
class CliCredentialFlag:
commands: frozenset[str] | None
mode: Literal["userinfo", "password", "token"]
attached: bool = False
def _entropy(value: str) -> float:
frequencies = Counter(value)
return -sum((count / len(value)) * math.log2(count / len(value)) for count in frequencies.values()) if value else 0.0
def _key_parts(key: str) -> tuple[str, list[str]]:
camel = re.sub(r"(?<=[a-zа-яё])(?=[A-ZА-ЯЁ])", "_", key)
camel = re.sub(r"(?<=[A-ZА-ЯЁ])(?=[A-ZА-ЯЁ][a-zа-яё])", "_", camel)
normalized = re.sub(r"[\s.\-/]+", "_", camel.lower()).strip("_")
return normalized, [part for part in normalized.split("_") if part]
_FILE_EXT_RE = re.compile(r"\.(?:md|mdx|go|js|mjs|cjs|ts|tsx|jsx|py|sh|rb|rs|java|kt|cs|cpp|cc|c|h|hpp|yml|yaml|json|toml|txt|cfg|conf|ini|env|html|css|scss|sql|log|xml|lock)(?:$|[-:]\d)", re.I)
def classify_key(key: str) -> str | None:
if "/" in key or _FILE_EXT_RE.search(key.strip()):
return None
normalized, parts = _key_parts(key)
if normalized in set(_ONTOLOGY["stop"]) or any(part in set(_ONTOLOGY["stop"]) for part in parts):
return None
for label, names in _ONTOLOGY["exact"].items():
if normalized in names:
return label
for label, segments in _ONTOLOGY["strong"].items():
if any(part in segments for part in parts):
return label
for label, weak in _ONTOLOGY["weak"].items():
for segment, qualifiers in weak.items():
if segment not in parts:
continue
if not qualifiers or any(part in qualifiers for part in parts):
return label
return None
def _is_live_opaque(value: str) -> bool:
compact = value.strip("\"'«»“”")
return len(compact) >= 16 and any(character.isdigit() for character in compact) and re.fullmatch(r"[A-Za-z0-9+/=_-]+", compact) is not None and _entropy(compact) >= 3.5
def _format_preserving(value: str, *, seed: int, record_key: str, offset: int) -> str:
rng = seeded_rng(seed, record_key, offset)
if value.isdigit():
return "".join(rng.choice("0123456789") for _ in value)
if re.fullmatch(r"[0-9A-Fa-f]+", value):
alphabet = "0123456789ABCDEF" if value.upper() == value else "0123456789abcdef"
return "".join(rng.choice(alphabet) for _ in value)
replacement = []
for character in value:
if character.isdigit():
replacement.append(rng.choice("0123456789"))
elif character.isupper():
replacement.append(rng.choice("ABCDEFGHIJKLMNOPQRSTUVWXYZ"))
elif character.islower():
replacement.append(rng.choice("abcdefghijklmnopqrstuvwxyz"))
else:
replacement.append(character)
return "".join(replacement)
_LITERAL_ESCAPE_RE = re.compile(r"\\+[nrt]")
_STRIP_CHARS = "\"'«»“”\\()"
def _kv_value(match: re.Match[str]) -> tuple[int, int, str, bool]:
for group, quoted in (("quoted_value", True), ("smart_value", True), ("bare_value", False)):
if match.group(group) is not None:
start, end = match.span(group)
value = match.group(group)
if not quoted:
cut = _LITERAL_ESCAPE_RE.search(value)
if cut and cut.start() > 0:
value = value[:cut.start()]; end = start + len(value)
text = match.string
opened = value.count("{") - value.count("}") + value.count("[") - value.count("]")
while opened > 0 and end < len(text) and text[end] in "}]":
value += text[end]; end += 1; opened -= 1
return start, end, value, quoted
raise AssertionError("K/V regex produced no value")
def _kv_core(start: int, end: int, value: str) -> tuple[int, int, str, bool]:
core = value.strip(_STRIP_CHARS)
lead = value.find(core) if core else 0
opening, closing = value[:lead], value[lead + len(core):]
quotes = "\"'«»“”"
wrapped = any(character in quotes for character in opening) and any(character in quotes for character in closing)
return start + lead, start + lead + len(core), core, wrapped
def _kv_matches(text: str):
tokens = [match.span() for match in SLOT_RE.finditer(text)]
position = 0
while True:
match = KV_RE.search(text, position)
if match is None:
return
key_start = match.start("key")
enclosing = next((span for span in tokens if span[0] <= key_start < span[1]), None)
if enclosing:
position = enclosing[1]
continue
yield match
position = match.end("key")
def _overlaps(start: int, end: int, replacements: list[tuple[int, int, str]]) -> bool:
return any(not (end <= left or start >= right) for left, right, _ in replacements)
def _replace(text: str, replacements: list[tuple[int, int, str]]) -> str:
for start, end, replacement in sorted(replacements, reverse=True):
text = text[:start] + replacement + text[end:]
return text
def _merge_ranges(ranges: list[list[int]]) -> list[list[int]]:
merged: list[list[int]] = []
for start, end in sorted(ranges):
if merged and start <= merged[-1][1]:
merged[-1][1] = max(merged[-1][1], end)
else:
merged.append([start, end])
return merged
def sanitize_synthetic_text(text: str, covered_ranges: list[list[int]], *, seed: int, record_key: str, rules: list[GitleaksRule], max_rounds: int = 64) -> tuple[str, list[list[int]], int, int]:
covered = _merge_ranges([list(item) for item in covered_ranges]); added: list[list[int]] = []; replaced_count = 0; rounds = 0
for _iteration in range(max_rounds):
pending: list[tuple[int, int, str]] = []
def add_uncovered(start: int, end: int) -> None:
pieces = [(start, end)]
for covered_start, covered_end in covered:
remainder = []
for piece_start, piece_end in pieces:
if piece_end <= covered_start or piece_start >= covered_end:
remainder.append((piece_start, piece_end))
else:
if piece_start < covered_start:
remainder.append((piece_start, covered_start))
if covered_end < piece_end:
remainder.append((covered_end, piece_end))
pieces = remainder
for piece_start, piece_end in pieces:
if piece_start >= piece_end or SLOT_RE.search(text[piece_start:piece_end]) or _overlaps(piece_start, piece_end, pending):
continue
pending.append((piece_start, piece_end, _format_preserving(text[piece_start:piece_end], seed=seed, record_key=record_key, offset=piece_start)))
for site in credential_sites(text, rules):
if site.disposition != "quarantine":
add_uncovered(site.start, site.end)
if not pending:
break
text = _replace(text, pending)
fresh = [[start, end] for start, end, _replacement in pending]
added.extend(fresh); covered = _merge_ranges([*covered, *fresh])
replaced_count += len(pending); rounds += 1
return text, _merge_ranges(added), replaced_count, rounds
_CREDENTIAL_ENTITIES = frozenset({"PASSWORD", "LOGIN", "JWT", "BEARER_TOKEN", "API_KEY", "PRIVATE_KEY", "DB_URL"})
def _pii_recognizers() -> list:
global _PROJECT_PII_RECOGNIZERS
if _PROJECT_PII_RECOGNIZERS is None:
try:
candidates = (Path(__file__).resolve().parents[2] / "presidio", Path("/app"))
for directory in candidates:
if (directory / "recognizers").is_dir() and str(directory) not in sys.path:
sys.path.insert(0, str(directory))
from recognizers import ALL_RECOGNIZERS
selected = []
for factory in ALL_RECOGNIZERS:
recognizer = factory()
supported = set(getattr(recognizer, "supported_entities", []) or [])
if supported and supported.isdisjoint(_CREDENTIAL_ENTITIES):
selected.append(recognizer)
_PROJECT_PII_RECOGNIZERS = selected
except Exception:
_PROJECT_PII_RECOGNIZERS = []
return _PROJECT_PII_RECOGNIZERS
def pii_recognizer_names() -> list[str]:
return sorted(type(recognizer).__name__ for recognizer in _pii_recognizers())
def _project_pii_matches(text: str) -> list[tuple[int, int]]:
matches = []
for recognizer in _pii_recognizers():
try:
matches.extend((result.start, result.end) for result in recognizer.analyze(text, entities=None, nlp_artifacts=None))
except Exception:
continue
return matches
def load_gitleaks_rules(path: Path | None) -> tuple[list[GitleaksRule], int]:
if path is None or not path.exists():
return [], 0
try:
import tomllib
raw_rules = tomllib.loads(path.read_text(encoding="utf-8")).get("rules", [])
except ImportError:
raw_rules = []
for block in path.read_text(encoding="utf-8").split("[[rules]]")[1:]:
identifier = re.search(r'^id\s*=\s*"([^"]+)"', block, re.M)
regex = re.search(r"^regex\s*=\s*'''(.*?)'''", block, re.M | re.S)
if identifier and regex:
keywords = re.search(r"^keywords\s*=\s*\[(.*?)\]", block, re.M | re.S)
secret_group = re.search(r"^secretGroup\s*=\s*(\d+)", block, re.M)
entropy = re.search(r"^entropy\s*=\s*([\d.]+)", block, re.M)
raw_rules.append({"id": identifier.group(1), "regex": regex.group(1), "keywords": re.findall(r'"([^"]+)"', keywords.group(1)) if keywords else [], "secretGroup": int(secret_group.group(1)) if secret_group else None, "entropy": float(entropy.group(1)) if entropy else None})
rules = []; skipped = 0
for rule in raw_rules:
identifier = str(rule.get("id", "")).lower(); regex = rule.get("regex")
if not isinstance(regex, str):
continue
label = "AUTH_TOKEN" if "token" in identifier else "SECRET_KEY" if "key" in identifier or "secret" in identifier else None
if label is None:
continue
try:
keywords = tuple(str(item).lower() for item in rule.get("keywords", []) if isinstance(item, str))
with warnings.catch_warnings():
warnings.simplefilter("error", DeprecationWarning)
warnings.simplefilter("error", FutureWarning)
pattern = re.compile(regex)
group = int(rule.get("secretGroup") or (1 if pattern.groups else 0))
rules.append((pattern, label, keywords, min(group, pattern.groups), float(rule.get("entropy") or 0.0)))
except (re.error, DeprecationWarning, FutureWarning):
skipped += 1
return rules, skipped
def _value_allowed(label: str, value: str, quoted: bool) -> bool:
value = value.strip()
if value.isdigit() and label != "PASSWORD":
return False
if label == "LOGIN":
return bool(re.fullmatch(r"[\w.@][\w.@\-]{2,63}", value))
return 4 <= len(value) <= 256 and (quoted or not any(character.isspace() for character in value))
def detector_core(text: str, start: int, end: int) -> tuple[int, int]:
value = text[start:end]
if "=" in value:
prefix, suffix = value.split("=", 1)
if suffix and re.fullmatch(r"[?&]?[A-Za-zА-Яа-яЁё_][\w.-]{0,64}", prefix):
start += len(prefix) + 1
while start < end and text[start] in "/?&":
start += 1
while end > start and text[end - 1] in ".,;)]}":
end -= 1
return start, end
def _kv_site(match: re.Match[str], text: str) -> tuple[int, int, str, str | None, str | None]:
start, end, value, quoted = _kv_value(match)
start, end, value, wrapped = _kv_core(start, end, value); quoted = quoted or wrapped
label = classify_key(match.group("key"))
if not label or value.startswith("[PII_REDACTED"):
return start, end, value, None, None
if PLACEHOLDER_RE.match(value):
return start, end, value, None, "placeholder"
if is_inert_value(value, quoted):
return start, end, value, None, "code_reference"
if not _value_allowed(label, value, quoted) or _looks_like_user_sentence(match, text, end, quoted):
return start, end, value, None, None
start, end = detector_core(text, start, end)
value = text[start:end]
if not value:
return start, end, value, None, "empty_core"
return start, end, value, label, None
def kv_credential_sites(text: str) -> list[tuple[int, int, str, str]]:
sites = []
for match in _kv_matches(text):
start, end, value, label, _reason = _kv_site(match, text)
if label:
sites.append((start, end, value, label))
return sites
_CLI_CREDENTIAL_FLAGS = {
"-u": CliCredentialFlag(frozenset({"curl"}), "userinfo", True), "--user": CliCredentialFlag(frozenset({"curl"}), "userinfo"),
"-U": CliCredentialFlag(frozenset({"curl"}), "userinfo", True), "--proxy-user": CliCredentialFlag(frozenset({"curl"}), "userinfo"),
"--oauth2-bearer": CliCredentialFlag(frozenset({"curl"}), "token"),
"--password": CliCredentialFlag(None, "password"), "--passwd": CliCredentialFlag(None, "password"), "--pass": CliCredentialFlag(None, "password"),
"--http-password": CliCredentialFlag(frozenset({"wget"}), "password"), "--proxy-password": CliCredentialFlag(frozenset({"wget"}), "password"), "--ftp-password": CliCredentialFlag(frozenset({"wget"}), "password"),
"--auth": CliCredentialFlag(frozenset({"http", "https"}), "userinfo"), "-a": CliCredentialFlag(frozenset({"http", "https"}), "userinfo", True),
"-p": CliCredentialFlag(frozenset({"mysql"}), "password", True),
}
_CLI_FLAG_RE = re.compile(r"(?<![\w=-])(--[A-Za-z][\w-]*|-[A-Za-z])")
_CLI_REFERENCE_RE = re.compile(r"^(?:\$\w+|\$\{[A-Za-z_]\w*\}|%\w+%|\$\([^)]*\)|`[^`]*`)$")
_CLI_WRAPPERS = frozenset({"sudo", "env", "command", "xargs"})
def _cli_read_value(text: str, position: int) -> tuple[int, int, str] | None:
if position >= len(text):
return None
if text[position] in "\"'":
closing = text.find(text[position], position + 1)
if closing == -1:
return None
return position + 1, closing, text[position + 1:closing]
match = _CLI_VALUE_RE.match(text, position)
if not match:
return None
start, end = match.start(), match.end()
while end > start:
char = text[end - 1]
if char in "\"',;":
end -= 1
elif char in ")]}" and {")": "(", "]": "[", "}": "{"}[char] not in text[start:end - 1]:
end -= 1
else:
break
return (start, end, text[start:end]) if end > start else None
def _cli_part_kind(part: str) -> str:
part = part.strip().strip("\"'")
if not part:
return "empty"
if re.fullmatch(r"\$\{[A-Za-z_]\w*:-\}", part):
return "empty"
if re.fullmatch(r"\$\{[A-Za-z_]\w*:-.+\}", part):
return "dynamic_with_literal_fallback"
if _CLI_REFERENCE_RE.match(part) or PLACEHOLDER_RE.match(part) or is_inert_value(part, quoted=True):
return "reference"
return "literal"
def _cli_read_shell_value(text: str, position: int) -> tuple[int, int, str, str] | None:
if position >= len(text) or text[position].isspace(): return None
start = position; quote = None; escaped = backtick = saw_escape = False; parens = braces = 0
while position < len(text):
char = text[position]
if escaped: escaped = False; position += 1; continue
if char == "\\": escaped = saw_escape = True; position += 1; continue
if quote:
if char == quote: quote = None
position += 1; continue
if backtick:
if char == "`": backtick = False
position += 1; continue
if char in "\"'": quote = char; position += 1; continue
if char == "`": backtick = True; position += 1; continue
if text.startswith("$(", position): parens += 1; position += 2; continue
if text.startswith("${", position): braces += 1; position += 2; continue
if char == ")" and parens: parens -= 1; position += 1; continue
if char == "}" and braces: braces -= 1; position += 1; continue
if char.isspace() and not parens and not braces: break
position += 1
end = position
malformed = any((quote, escaped, backtick, parens, braces, saw_escape))
wrapped = end - start >= 2 and text[start] in "\"'" and text[end - 1] == text[start]
if wrapped:
start += 1; end -= 1
elif not malformed:
while end > start and text[end - 1] in "\"',;": end -= 1
value = text[start:end]
if not value: return None
kind = "ambiguous" if malformed else _cli_part_kind(value)
if not malformed and (value.startswith("$(") or value.startswith("${") or value.startswith("`")):
kind = "reference"
return start, end, value, kind
def _cli_command(text: str, position: int) -> str | None:
segment = re.split(r"(?:;|&&|\|\| |\|)", text[:position])[-1]
embedded = re.findall(r"(?:^|[\s\"'])((?:curl|wget|mysql|http|https))(?:\.exe)?(?=$|[\s\"'])", segment, re.I)
if embedded:
return embedded[-1].lower()
for word in re.findall(r"(?:[^\s'\"]+|'[^']*'|\"[^\"]*\")+", segment):
token = re.sub(r"\.exe$", "", word.strip("'\"").rsplit("/", 1)[-1].lower())
if token in _CLI_WRAPPERS or token.startswith("-") or re.fullmatch(r"[A-Za-z_]\w*=.*", token): continue
return token
return None
def _cli_credential_parts(text: str):
for flag_match in _CLI_FLAG_RE.finditer(text):
flag = flag_match.group(1)
config = _CLI_CREDENTIAL_FLAGS.get(flag)
if config is None:
continue
command = _cli_command(text, flag_match.start())
if config.commands is not None and command not in config.commands:
continue
after = flag_match.end()
if after < len(text) and text[after] == "=":
read = _cli_read_shell_value(text, after + 1)
elif config.attached and after < len(text) and not text[after].isspace():
read = _cli_read_shell_value(text, after)
else:
cursor = after
while cursor < len(text) and text[cursor] in " \t":
cursor += 1
read = _cli_read_shell_value(text, cursor) if cursor > after or (after < len(text) and text[after].isspace()) else None
if read is None:
continue
value_start, value_end, value, value_kind = read
part = {"password": None, "username": None}
if config.mode == "userinfo":
colon = value.find(":")
if colon == -1:
continue
username = value[:colon]
part["username"] = (value_start, value_start + colon, username, "LOGIN", _cli_part_kind(username))
password_kind = value_kind if value_kind in {"ambiguous", "dynamic_with_literal_fallback"} else _cli_part_kind(value[colon + 1:])
part["password"] = (value_start + colon + 1, value_end, value[colon + 1:], "PASSWORD", password_kind)
else:
part["password"] = (value_start, value_end, value, "AUTH_TOKEN" if config.mode == "token" else "PASSWORD", value_kind)
yield part
def cli_credential_sites(text: str) -> list[tuple[int, int, str, str, str]]:
sites: list[tuple[int, int, str, str, str]] = []
for part in _cli_credential_parts(text):
password = part["password"]
if password and password[4] != "empty":
sites.append(password)
return sites
def cli_credential_slots(text: str) -> list[tuple[int, int, str, str]]:
slots: list[tuple[int, int, str, str]] = []
for part in _cli_credential_parts(text):
password = part["password"]
if not password or password[4] != "literal":
continue
username = part["username"]
if username and username[4] == "literal":
slots.append((username[0], username[1], "LOGIN", "cli"))
slots.append((password[0], password[1], password[3], "cli"))
return slots
def credential_sites(text: str, rules: list[GitleaksRule]) -> list[CredentialSite]:
sites = [CredentialSite(start, end, label, "kv", "slot") for start, end, _value, label in kv_credential_sites(text)]
sites += [CredentialSite(start, end, label, "cli", "slot" if kind == "literal" else "quarantine") for start, end, _value, label, kind in cli_credential_sites(text) if kind in {"literal", "ambiguous", "dynamic_with_literal_fallback"}]
sites += [CredentialSite(start, end, label or "SECRET", "opaque", "replace") for start, end, label, _match_start, _match_end in _find_live(text, rules)]
return sorted(set(sites), key=lambda item: (item.start, item.end, item.label, item.detector))
def _looks_like_user_sentence(match: re.Match[str], text: str, value_end: int, quoted: bool) -> bool:
key = match.group("key")
_, parts = _key_parts(key)
if quoted:
return False
remainder = text[value_end:text.find("\n", value_end) if text.find("\n", value_end) != -1 else len(text)]
continues_as_sentence = bool(re.match(r"\s+[\wА-Яа-яЁё]{2,}", remainder))
return continues_as_sentence and (
"user" in parts or
(match.group("sep") == ":" and " " in key and len(parts) >= 3)
)
def _find_live(text: str, rules: list[GitleaksRule]) -> list[tuple[int, int, str | None, int, int]]:
found = []
for match in OPAQUE_RE.finditer(text):
if not _is_live_opaque(match.group()):
continue
start, end, value, label = match.start(), match.end(), match.group(), None
if "=" in value:
prefix, suffix = value.split("=", 1)
classified = classify_key(prefix)
if classified and suffix:
start += len(prefix) + 1; value = suffix; label = classified
found.append((start, end, label, match.start(), match.end()))
lowered = text.lower()
for pattern, label, keywords, group, entropy in rules:
if keywords and not any(keyword in lowered for keyword in keywords):
continue
for match in pattern.finditer(text):
start, end = match.span(group)
if start < 0 or end <= start:
start, end = match.span()
if entropy and _entropy(text[start:end]) < entropy:
continue
found.append((start, end, label, match.start(), match.end()))
normalized = []
for start, end, label, match_start, match_end in found:
start, end = detector_core(text, start, end)
if start < end:
normalized.append((start, end, label, match_start, match_end))
return sorted(set(normalized), key=lambda item: (item[0], item[1], item[2] or ""))
def scrub_row(raw: dict, slot_counter: int, *, seed: int = 42, gitleaks_rules: list[GitleaksRule] | None = None) -> tuple[dict | None, list[dict], dict[str, int], int]:
text = nfc(str(raw.get("content", raw.get("text", ""))))
record_key = str(raw.get("group_id") or raw.get("conversation_id") or raw.get("session_id") or stable_id(text))
counts: Counter[str] = Counter(); slots: list[dict] = []; flags = list(raw.get("flags", [])); negative_reason = raw.get("negative_reason")
pii_replacements: list[tuple[int, int, str]] = []
for pattern in (EMAIL_RE, PHONE_RE, SNILS_RE, CARD_RE, PASSPORT_RE, INN_RE):
for match in pattern.finditer(text):
start, end = match.span(1) if pattern in {INN_RE, PASSPORT_RE} else match.span()
if not _overlaps(start, end, pii_replacements):
pii_replacements.append((start, end, "[PII_REDACTED]")); counts["pii_replaced"] += 1
for start, end in _project_pii_matches(text):
if not _overlaps(start, end, pii_replacements):
pii_replacements.append((start, end, "[PII_REDACTED]")); counts["pii_replaced"] += 1
text = _replace(text, pii_replacements)
if pii_replacements:
flags.append("pii_replaced")
replacements: list[tuple[int, int, str]] = []; entropy_replacements: list[tuple[int, int, str]] = []
def add_slot(start: int, end: int, label: str, slot_type: str) -> None:
nonlocal slot_counter
if _overlaps(start, end, replacements):
return
token = f"⟦SLOT:{label}:{slot_counter}⟧"; slot_counter += 1
replacements.append((start, end, token))
slots.append({"slot_id": token, "record_key": record_key, "klass": label, "injection_slot_type": slot_type})
counts["slots"] += 1
for marker in MARKER_RE.finditer(text):
kind = marker.group(1); label = MARKER_LABELS.get(kind) or classify_key(kind)
if kind == "ENV":
context = text[max(0, marker.start() - 120):marker.end() + 120]
key_match = re.search(r"([\w.-]+)\s*[:=]\s*\[REDACTED:ENV\]", context, re.I)
label = classify_key(key_match.group(1)) if key_match else None
if label:
add_slot(marker.start(), marker.end(), label, "redacted_marker")
cli_quarantine_ranges = [(site.start, site.end) for site in credential_sites(text, gitleaks_rules or []) if site.detector == "cli" and site.disposition == "quarantine"]
for match in _kv_matches(text):
start, end, _value, label, reason = _kv_site(match, text)
if any(not (end <= left or start >= right) for left, right in cli_quarantine_ranges):
continue
if reason:
negative_reason = negative_reason or reason; counts[reason] += 1; continue
if label:
add_slot(start, end, label, "dsn" if "://" in text[max(0, match.start() - 50):match.end() + 50] else "kv_pair")
for match in CONTRACT_RE.finditer(text):
if not _overlaps(match.start(1), match.end(1), replacements):
add_slot(match.start(1), match.end(1), "CONTRACT_NUMBER", "kv_pair")
for start, end, label, slot_type in cli_credential_slots(text):
add_slot(start, end, label, slot_type)
for start, end, _label, _match_start, _match_end in _find_live(text, gitleaks_rules or []):
if _overlaps(start, end, replacements):
overlapping = sorted((left, right) for left, right, _ in replacements if left < end and right > start)
leftovers = []; cursor = start
for left, right in overlapping:
if left > cursor:
leftovers.append((cursor, left))
cursor = max(cursor, right)
if cursor < end:
leftovers.append((cursor, end))
if any(_is_live_opaque(text[piece_start:piece_end]) for piece_start, piece_end in leftovers):
counts["quarantined_conflict"] += 1
return None, [], dict(counts), slot_counter
continue
replacement = _format_preserving(text[start:end], seed=seed, record_key=record_key, offset=start)
item = (start, end, replacement); replacements.append(item); entropy_replacements.append(item); counts["entropy_replaced_values"] += 1
safe_text = _replace(text, replacements)
final_ranges = []; delta = 0
entropy_starts = {(start, end, replacement) for start, end, replacement in entropy_replacements}
for start, end, replacement in sorted(replacements):
final_start = start + delta; final_end = final_start + len(replacement)
if (start, end, replacement) in entropy_starts:
final_ranges.append([final_start, final_end])
delta += len(replacement) - (end - start)
final_ranges = _merge_ranges(final_ranges)
if entropy_replacements:
flags.append("entropy_replaced")
for _iteration in range(64):
pending: list[tuple[int, int, str]] = []
def add_post_replacement(start: int, end: int) -> None:
pieces = [(start, end)]
for covered_start, covered_end in final_ranges:
remainder = []
for piece_start, piece_end in pieces:
if piece_end <= covered_start or piece_start >= covered_end:
remainder.append((piece_start, piece_end))
else:
if piece_start < covered_start:
remainder.append((piece_start, covered_start))
if covered_end < piece_end:
remainder.append((covered_end, piece_end))
pieces = remainder
for piece_start, piece_end in pieces:
if piece_start >= piece_end or SLOT_RE.search(safe_text[piece_start:piece_end]) or _overlaps(piece_start, piece_end, pending):
continue
pending.append((piece_start, piece_end, _format_preserving(safe_text[piece_start:piece_end], seed=seed, record_key=record_key, offset=piece_start)))
for site in credential_sites(safe_text, gitleaks_rules or []):
if site.disposition != "quarantine":
add_post_replacement(site.start, site.end)
if not pending:
break
safe_text = _replace(safe_text, pending)
final_ranges = _merge_ranges([*final_ranges, *[[start, end] for start, end, _replacement in pending]])
counts["postscan_replaced_values"] += len(pending)
counts["postscan_rounds"] += 1
if "entropy_replaced" not in flags:
flags.append("entropy_replaced")
slot_spans = [list(match.span()) for match in SLOT_RE.finditer(safe_text)]
covered = _merge_ranges([*final_ranges, *slot_spans])
for site in credential_sites(safe_text, gitleaks_rules or []):
if site.detector == "cli" and site.disposition == "quarantine" and not any(left <= site.start and site.end <= right for left, right in covered):
counts["cli_credential_quarantine"] += 1
return None, [], dict(counts), slot_counter
row = dict(raw); row.pop("content", None); row["text"] = safe_text; row["flags"] = sorted(set(flags)); row["negative_reason"] = negative_reason; row["entropy_replaced_ranges"] = final_ranges
return row, slots, dict(counts), slot_counter
def unresolved_live_values(row: dict, rules: list[GitleaksRule]) -> int:
ranges = _merge_ranges([list(item) for item in row.get("entropy_replaced_ranges", [])])
slots = [match.span() for match in SLOT_RE.finditer(row["text"])]
def covered(start: int, end: int) -> bool:
return any(left <= start and end <= right for left, right in [*ranges, *slots])
return sum(not covered(site.start, site.end) for site in credential_sites(row["text"], rules))
_SLOT_BASE_STRIDE = 65536
_WORKER_RULES: list[GitleaksRule] = []
def _init_worker(config: str) -> None:
global _WORKER_RULES
_WORKER_RULES = load_gitleaks_rules(Path(config))[0] if config else []
def _scrub_chunk(payload: tuple[int, int, list[str]]) -> tuple[list[str], list[dict], dict[str, int]]:
seed, start_index, lines = payload
out_rows: list[str] = []; out_slots: list[dict] = []; counts: Counter[str] = Counter()
for offset, line in enumerate(lines):
raw = json.loads(line); counts["input"] += 1
safe, found, row_counts, _ = scrub_row(raw, (start_index + offset) * _SLOT_BASE_STRIDE, seed=seed, gitleaks_rules=_WORKER_RULES)
counts.update(row_counts)
if safe is None:
counts["quarantined"] += 1
else:
counts["safe"] += 1; out_rows.append(json.dumps(safe, ensure_ascii=False, sort_keys=True, separators=(",", ":"))); out_slots.extend(found)
return out_rows, out_slots, dict(counts)
def _unresolved_chunk(lines: list[str]) -> int:
return sum(unresolved_live_values(json.loads(line), _WORKER_RULES) for line in lines)
def _line_chunks(paths: list[Path], size: int):
index = 0; buffer: list[str] = []
for path in paths:
with path.open(encoding="utf-8") as stream:
for line in stream:
if not line.strip():
continue
buffer.append(line)
if len(buffer) >= size:
yield index, buffer; index += len(buffer); buffer = []
if buffer:
yield index, buffer
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__); parser.add_argument("--input", type=Path, action="append", help="turns JSONL; defaults to both extract outputs"); parser.add_argument("--gitleaks-config", type=Path, default=Path("../dataset/gitleaks/config/gitleaks.toml")); parser.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 2) - 2)); parser.add_argument("--verify-only", action="store_true", help="re-scan the existing scrubbed_turns.jsonl without rewriting it"); parser.add_argument("--allow-missing-pii-recognizers", action="store_true", help="proceed without personal-data recognizers; the output is then NOT safe to export"); add_common_args(parser); args = parser.parse_args()
inputs = args.input or [args.staging / "codechat_turns.jsonl", args.staging / "swechat_turns.jsonl"]
recognizers = pii_recognizer_names()
if not recognizers and not args.allow_missing_pii_recognizers:
raise SystemExit("PII recognizers unavailable (is presidio_analyzer installed?); "
"re-run inside the pipeline image or pass --allow-missing-pii-recognizers")
rules, skipped_rules = load_gitleaks_rules(args.gitleaks_config); slots: list[dict] = []; counts: Counter[str] = Counter()
config = str(args.gitleaks_config) if args.gitleaks_config and Path(args.gitleaks_config).exists() else ""
pool = multiprocessing.get_context("spawn").Pool(args.workers, initializer=_init_worker, initargs=(config,))
scrubbed_path = args.staging / "scrubbed_turns.jsonl"; scrubbed_path.parent.mkdir(parents=True, exist_ok=True)
if args.verify_only:
report_path = args.staging / "quarantine_report.json"
report = json.loads(report_path.read_text(encoding="utf-8")) if report_path.is_file() else {}
if report.get("scrubber_policy_version") != SCRUBBER_POLICY_VERSION:
raise SystemExit("staging was not produced by the current scrubber policy")
with pool:
unresolved = sum(pool.imap_unordered(_unresolved_chunk, (lines for _, lines in _line_chunks([scrubbed_path], 2000))))
print(json.dumps({"unresolved_live_values": unresolved, "workers": args.workers}, sort_keys=True))
raise SystemExit(1 if unresolved else 0)
with pool, scrubbed_path.open("w", encoding="utf-8", newline="\n") as out:
payloads = ((args.seed, index, lines) for index, lines in _line_chunks(inputs, 2000))
for out_rows, out_slots, chunk_counts in pool.imap(_scrub_chunk, payloads):
counts.update(chunk_counts); slots.extend(out_slots)
for line in out_rows:
out.write(line + "\n")
write_jsonl(args.staging / "slots.jsonl", slots)
out.flush()
unresolved = sum(pool.imap_unordered(_unresolved_chunk, (lines for _, lines in _line_chunks([scrubbed_path], 2000))))
report = {"scrubber_policy_version": SCRUBBER_POLICY_VERSION, "input": counts["input"], "safe": counts["safe"], "quarantined": counts["quarantined"], "unresolved_live_values": unresolved, "gitleaks_rules": len(rules), "gitleaks_rules_skipped": skipped_rules, "workers": args.workers, "pii_recognizers": recognizers, "by_type": dict(sorted(counts.items()))}
(args.staging / "quarantine_report.json").write_text(json.dumps(report, ensure_ascii=False, sort_keys=True, indent=2) + "\n", encoding="utf-8"); write_report(args.staging, "scrubber", **report); print(json.dumps(report, ensure_ascii=False, sort_keys=True))
if unresolved:
raise SystemExit(1)
if __name__ == "__main__":
main()