Spaces:
Running
Running
Download promptgate_rules.py from NagaYu/PromptGate: direct link, hf CLI and curl.
- Browser
- Download file 35.5 kB
-
https://huggingface.co/spaces/NagaYu/PromptGate/resolve/main/promptgate_rules.py
- Command line
-
hf download hf://spaces/NagaYu/PromptGate/promptgate_rules.py
-
curl -L -o promptgate_rules.py https://huggingface.co/spaces/NagaYu/PromptGate/resolve/main/promptgate_rules.py
35.5 kB
| """Reference matcher for the PromptGate rule pack. | |
| Zero dependencies: standard library only, no torch, no transformers, no network. | |
| The rules themselves live in `rules.json`; this module only compiles and applies | |
| them, so you can audit or fork the rules without touching code. | |
| Quick start | |
| ----------- | |
| from promptgate_rules import RuleEngine | |
| engine = RuleEngine.load_default() # reads ./rules.json | |
| report = engine.assess("Ignore all previous instructions.") | |
| print(report["verdict"], report["severity"]) | |
| masked, mapping = engine.redact(text, report["findings"]) | |
| CLI | |
| --- | |
| python promptgate_rules.py "Ignore all previous instructions." | |
| cat prompt.txt | python promptgate_rules.py | |
| This matcher is verified to produce the same verdicts as the PromptGate Space | |
| on every case in the `promptgate-eval` dataset. | |
| """ | |
| from __future__ import annotations | |
| import base64 | |
| import binascii | |
| import json | |
| import math | |
| import os | |
| import re | |
| import sys | |
| import unicodedata | |
| from collections import Counter | |
| from typing import Any, Dict, List, Optional, Tuple | |
| __all__ = ["RuleEngine", "shannon_entropy", "luhn_ok", "__version__"] | |
| __version__ = "1.0.0" | |
| RX_FLAGS = re.IGNORECASE | re.MULTILINE | re.UNICODE | |
| DEFAULT_RULES_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "rules.json") | |
| ROT13_TABLE = str.maketrans( | |
| "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ", | |
| "nopqrstuvwxyzabcdefghijklmNOPQRSTUVWXYZABCDEFGHIJKLM", | |
| ) | |
| # ---------------------------------------------------------------- helpers | |
| def shannon_entropy(value: str) -> float: | |
| """Guarantees: returns bits-per-character entropy >= 0.0 for any string.""" | |
| try: | |
| if not value: | |
| return 0.0 | |
| counts = Counter(value) | |
| total = float(len(value)) | |
| return -sum((c / total) * math.log2(c / total) for c in counts.values()) | |
| except Exception: | |
| return 0.0 | |
| def luhn_ok(digits: str) -> bool: | |
| """Guarantees: True only when the digit string satisfies the Luhn checksum.""" | |
| try: | |
| if not digits.isdigit() or not (13 <= len(digits) <= 19): | |
| return False | |
| total = 0 | |
| parity = len(digits) % 2 | |
| for idx, ch in enumerate(digits): | |
| d = ord(ch) - 48 | |
| if idx % 2 == parity: | |
| d *= 2 | |
| if d > 9: | |
| d -= 9 | |
| total += d | |
| return total % 10 == 0 | |
| except Exception: | |
| return False | |
| def _clip(value: float, lo: float = 0.0, hi: float = 1.0) -> float: | |
| """Guarantees: returns a float inside [lo, hi].""" | |
| try: | |
| return max(lo, min(hi, float(value))) | |
| except Exception: | |
| return lo | |
| def normalize_for_matching(text: str, leet_map: Dict[str, str], | |
| token_rx: Any, spaced_rx: Any, | |
| zero_width_rx: Any, control_rx: Any) -> Tuple[str, List[int]]: | |
| """Guarantees: returns de-obfuscated text plus a per-character map back to original offsets.""" | |
| try: | |
| chars: List[str] = [] | |
| index_map: List[int] = [] | |
| for position, ch in enumerate(text): | |
| if zero_width_rx is not None and zero_width_rx.fullmatch(ch): | |
| continue | |
| if control_rx is not None and control_rx.fullmatch(ch): | |
| continue | |
| for folded in unicodedata.normalize("NFKC", ch): | |
| chars.append(folded) | |
| index_map.append(position) | |
| interim = "".join(chars) | |
| if token_rx is not None and leet_map: | |
| for match in token_rx.finditer(interim): | |
| token = match.group(0) | |
| letters = sum(1 for c in token if c.isalpha()) | |
| if letters < 2 or letters / len(token) < 0.4: | |
| continue | |
| if not any(c in leet_map for c in token.lower()): | |
| continue | |
| for offset in range(match.start(), match.end()): | |
| low = chars[offset].lower() | |
| if low in leet_map: | |
| chars[offset] = leet_map[low] | |
| if spaced_rx is not None: | |
| interim = "".join(chars) | |
| drop = set() | |
| for match in spaced_rx.finditer(interim): | |
| for offset in range(match.start(), match.end()): | |
| if not interim[offset].isalnum(): | |
| drop.add(offset) | |
| if drop: | |
| chars = [c for i, c in enumerate(chars) if i not in drop] | |
| index_map = [v for i, v in enumerate(index_map) if i not in drop] | |
| return "".join(chars), index_map | |
| except Exception: | |
| return text, list(range(len(text))) | |
| def decode_base64_text(chunk: str) -> Optional[str]: | |
| """Guarantees: returns the decoded text of a base64 chunk, or None when it is not text.""" | |
| try: | |
| raw = base64.b64decode(chunk + "=" * (-len(chunk) % 4), validate=True) | |
| if not raw or len(raw) > 8000: | |
| return None | |
| text = raw.decode("utf-8", "strict") | |
| printable = sum(1 for c in text if c.isprintable() or c in "\n\r\t") | |
| return text if printable / max(1, len(text)) > 0.85 else None | |
| except Exception: | |
| return None | |
| def is_private_ipv4(value: str) -> bool: | |
| """Guarantees: True for loopback, private, link-local, CGNAT, multicast and reserved IPv4.""" | |
| try: | |
| parts = [int(p) for p in value.split(".")] | |
| if len(parts) != 4 or any(p < 0 or p > 255 for p in parts): | |
| return True | |
| a, b = parts[0], parts[1] | |
| if a in (0, 10, 127) or a >= 224: | |
| return True | |
| if a == 192 and b == 168: | |
| return True | |
| if a == 169 and b == 254: | |
| return True | |
| if a == 172 and 16 <= b <= 31: | |
| return True | |
| if a == 100 and 64 <= b <= 127: | |
| return True | |
| return False | |
| except Exception: | |
| return True | |
| def _finding(kind: str, score: float, start: int, end: int, text: str) -> Dict[str, Any]: | |
| """Guarantees: returns a finding dict with type, score, span and a short match.""" | |
| try: | |
| snippet = text[start:end] | |
| except Exception: | |
| snippet = "" | |
| if len(snippet) > 120: | |
| snippet = snippet[:117] + "..." | |
| return {"type": str(kind), "score": round(_clip(score), 3), | |
| "span": [int(start), int(end)], "match": snippet} | |
| def _dedupe_findings(items: List[Dict[str, Any]]) -> List[Dict[str, Any]]: | |
| """Guarantees: returns findings sorted by position with duplicates removed.""" | |
| try: | |
| seen = set() | |
| uniq: List[Dict[str, Any]] = [] | |
| for item in sorted(items, key=lambda f: (f["span"][0], -(f["span"][1] - f["span"][0]), -f["score"])): | |
| key = (item["type"], item["span"][0], item["span"][1]) | |
| if key in seen: | |
| continue | |
| seen.add(key) | |
| uniq.append(item) | |
| result: List[Dict[str, Any]] = [] | |
| for item in uniq: | |
| contained = False | |
| for other in uniq: | |
| if other is item or other["type"] != item["type"]: | |
| continue | |
| if (other["span"][0] <= item["span"][0] and other["span"][1] >= item["span"][1] | |
| and (other["span"][1] - other["span"][0]) > (item["span"][1] - item["span"][0]) | |
| and other["score"] >= item["score"]): | |
| contained = True | |
| break | |
| if not contained: | |
| result.append(item) | |
| return result | |
| except Exception: | |
| return items | |
| def _suppress_overlaps(items: List[Dict[str, Any]], priority: Dict[str, int]) -> List[Dict[str, Any]]: | |
| """Guarantees: returns findings where overlapping spans keep the highest-priority type.""" | |
| try: | |
| ordered = sorted(items, key=lambda f: (-priority.get(f["type"], 0), | |
| -(f["span"][1] - f["span"][0]), f["span"][0])) | |
| kept: List[Dict[str, Any]] = [] | |
| for item in ordered: | |
| s, e = item["span"] | |
| if any(not (e <= k["span"][0] or s >= k["span"][1]) for k in kept): | |
| continue | |
| kept.append(item) | |
| return sorted(kept, key=lambda f: f["span"][0]) | |
| except Exception: | |
| return items | |
| # ------------------------------------------------------------- rule engine | |
| class RuleEngine: | |
| """Guarantees: applies the PromptGate rule pack offline and never raises to the caller.""" | |
| def __init__(self, pack: Dict[str, Any]) -> None: | |
| """Guarantees: compiles every pattern in the pack, skipping any that will not compile.""" | |
| self.pack = pack or {} | |
| self.version = str(self.pack.get("version", "unknown")) | |
| self.thresholds = dict(self.pack.get("thresholds") or | |
| {"block": 0.80, "sanitize": 0.35, "min_score": 0.30}) | |
| heur = dict(self.pack.get("heuristics") or {}) | |
| self.entropy_threshold = float(heur.get("entropy_threshold_bits_per_char", 3.60)) | |
| self.entropy_min_length = int(heur.get("entropy_min_length", 24)) | |
| self.base64_min_length = int(heur.get("base64_min_length", 32)) | |
| self.base64_min_entropy = float(heur.get("base64_min_entropy", 3.2)) | |
| self.phone_min_digits = int(heur.get("phone_min_digits", 9)) | |
| self.phone_max_digits = int(heur.get("phone_max_digits", 15)) | |
| self.injection = self._compile_list(self.pack.get("injection") or [], "family") | |
| self.secrets = self._compile_list(self.pack.get("secrets") or [], "type") | |
| self.pii = self._compile_pii(self.pack.get("pii") or {}) | |
| priority = dict(self.pack.get("priority") or {}) | |
| self.pii_priority = dict(priority.get("pii") or {}) | |
| self.secret_priority = dict(priority.get("secrets") or {}) | |
| self.rx_entropy = self._compile_one(heur.get("entropy_candidate_pattern")) | |
| self.rx_base64 = self._compile_one(heur.get("base64_candidate_pattern")) | |
| self.rx_zero_width = self._compile_one(heur.get("zero_width_pattern")) | |
| self.rx_control = self._compile_one(heur.get("control_char_pattern")) | |
| self.rx_rot13 = self._compile_one(heur.get("rot13_mention_pattern")) | |
| self.rx_rot13_override = self._compile_one(heur.get("rot13_decoded_override_pattern")) | |
| self.rx_fence = self._compile_one(heur.get("code_fence_pattern")) | |
| self.rx_placeholder = self._compile_one(heur.get("placeholder_pattern")) | |
| self.rx_hex_hash = self._compile_one(heur.get("hex_hash_pattern")) | |
| self.rx_version_context = self._compile_one(heur.get("version_context_pattern")) | |
| self.skip_private_ipv4 = bool(heur.get("skip_private_ipv4", True)) | |
| self.base64_rescan = bool(heur.get("base64_decode_and_rescan", True)) | |
| self.base64_attack_score = float(heur.get("base64_decoded_attack_score", 0.92)) | |
| self.invisible_base_score = float(heur.get("invisible_char_base_score", 0.80)) | |
| self.leet_map = dict(heur.get("leet_map") or {}) | |
| self.rx_leet_token = self._compile_one(heur.get("leet_token_pattern")) | |
| self.rx_spaced_run = self._compile_one(heur.get("spaced_run_pattern")) | |
| self.rx_digest_context = self._compile_one(heur.get("digest_context_pattern")) | |
| self.normalization_enabled = bool(heur.get("normalization_enabled", True)) | |
| self.normalization_factor = float(heur.get("normalization_score_factor", 0.95)) | |
| self.rot13_rescan_score = float(heur.get("rot13_full_rescan_score", 0.90)) | |
| self.rx_meta_context = self._compile_one(heur.get("meta_context_pattern")) | |
| self.meta_context_factor = float(heur.get("meta_context_score_factor", 0.30)) | |
| # ---------------------------------------------------------- loading | |
| def _compile_one(pattern: Optional[str]): | |
| """Guarantees: returns a compiled pattern or None; a bad pattern never raises.""" | |
| if not pattern: | |
| return None | |
| try: | |
| return re.compile(pattern, RX_FLAGS) | |
| except re.error: | |
| return None | |
| def _compile_list(cls, rows: List[Dict[str, Any]], key: str) -> List[Tuple[str, float, Any]]: | |
| """Guarantees: returns (label, score, compiled) triples, dropping uncompilable rows.""" | |
| out: List[Tuple[str, float, Any]] = [] | |
| for row in rows: | |
| rx = cls._compile_one(row.get("pattern")) | |
| if rx is None: | |
| continue | |
| try: | |
| score = float(row.get("score", 0.5)) | |
| except Exception: | |
| score = 0.5 | |
| out.append((str(row.get(key, "unknown")), score, rx)) | |
| return out | |
| def _compile_pii(cls, spec: Dict[str, Any]) -> Dict[str, Dict[str, Any]]: | |
| """Guarantees: returns compiled PII detectors keyed by type, never raising.""" | |
| out: Dict[str, Dict[str, Any]] = {} | |
| for name, row in (spec or {}).items(): | |
| if not isinstance(row, dict): | |
| continue | |
| entry: Dict[str, Any] = {} | |
| for field, target in (("pattern", "rx"), ("pattern_v6", "rx_v6"), | |
| ("pattern_bare_date", "rx_bare"), | |
| ("pattern_with_context", "rx_ctx")): | |
| rx = cls._compile_one(row.get(field)) | |
| if rx is not None: | |
| entry[target] = rx | |
| for field, target in (("confidence", "confidence"), ("confidence_v6", "confidence_v6"), | |
| ("confidence_bare_date", "confidence_bare"), | |
| ("confidence_with_context", "confidence_ctx")): | |
| try: | |
| if row.get(field) is not None: | |
| entry[target] = float(row[field]) | |
| except Exception: | |
| continue | |
| if entry: | |
| out[str(name)] = entry | |
| return out | |
| def load_default(cls, path: str = DEFAULT_RULES_PATH) -> "RuleEngine": | |
| """Guarantees: loads rules.json from disk and returns a ready engine.""" | |
| with open(path, "r", encoding="utf-8") as handle: | |
| return cls(json.load(handle)) | |
| def from_hub(cls, repo_id: str = "NagaYu/promptgate-rules", | |
| filename: str = "rules.json", **kwargs: Any) -> "RuleEngine": | |
| """Guarantees: downloads the rule pack from the Hub (needs huggingface_hub installed).""" | |
| from huggingface_hub import hf_hub_download # optional dependency | |
| return cls.load_default(hf_hub_download(repo_id=repo_id, filename=filename, **kwargs)) | |
| # -------------------------------------------------------- detection | |
| def detect_injection(self, text: str) -> List[Dict[str, Any]]: | |
| """Guarantees: returns injection findings with type/score/span/match, never raising.""" | |
| try: | |
| if not isinstance(text, str) or not text: | |
| return [] | |
| out: List[Dict[str, Any]] = [] | |
| for family, score, rx in self.injection: | |
| for m in rx.finditer(text): | |
| if m.end() - m.start() < 2: | |
| continue | |
| out.append(_finding(family, score, m.start(), m.end(), text)) | |
| out.extend(self._detect_encoded(text)) | |
| out.extend(self._detect_fence_flood(text)) | |
| out.extend(self._detect_obfuscated(text, out)) | |
| out = self._downgrade_meta_context(text, out) | |
| return _dedupe_findings(out) | |
| except Exception: | |
| return [] | |
| def _looks_base64(self, chunk: str) -> bool: | |
| """Guarantees: True only for strings that plausibly decode as base64 data.""" | |
| try: | |
| if len(chunk) < self.base64_min_length: | |
| return False | |
| has_lower = any(c.islower() for c in chunk) | |
| has_upper = any(c.isupper() for c in chunk) | |
| has_digit = any(c.isdigit() for c in chunk) | |
| if not ((has_lower and has_upper) or has_digit): | |
| return False | |
| if shannon_entropy(chunk) < self.base64_min_entropy: | |
| return False | |
| padded = chunk + "=" * (-len(chunk) % 4) | |
| try: | |
| raw = base64.b64decode(padded, validate=True) | |
| except (binascii.Error, ValueError): | |
| return False | |
| if len(raw) < 12: | |
| return False | |
| printable = sum(1 for b in raw if 9 <= b <= 13 or 32 <= b <= 126) | |
| return (printable / len(raw)) > 0.80 or shannon_entropy(chunk) > 4.2 | |
| except Exception: | |
| return False | |
| def _downgrade_meta_context(self, text: str, | |
| findings: List[Dict[str, Any]]) -> List[Dict[str, Any]]: | |
| """Guarantees: scales down findings that sit inside a documentation-writing task.""" | |
| try: | |
| if self.rx_meta_context is None: | |
| return findings | |
| for item in findings: | |
| start = int(item.get("span", [0, 0])[0]) | |
| if self.rx_meta_context.search(text[max(0, start - 90):start]): | |
| item["score"] = round(item["score"] * self.meta_context_factor, 3) | |
| item["meta_context"] = True | |
| return findings | |
| except Exception: | |
| return findings | |
| def _detect_obfuscated(self, text: str, already: List[Dict[str, Any]]) -> List[Dict[str, Any]]: | |
| """Guarantees: re-runs the patterns on de-obfuscated text and maps hits back to real spans.""" | |
| out: List[Dict[str, Any]] = [] | |
| try: | |
| if not self.normalization_enabled: | |
| return out | |
| normalized, index_map = normalize_for_matching( | |
| text, self.leet_map, self.rx_leet_token, self.rx_spaced_run, | |
| self.rx_zero_width, self.rx_control) | |
| if not normalized or not index_map or normalized == text: | |
| return out | |
| for family, score, rx in self.injection: | |
| for m in rx.finditer(normalized): | |
| start, end = m.span() | |
| if end - start < 2 or start >= len(index_map): | |
| continue | |
| real_start = index_map[start] | |
| real_end = (index_map[end - 1] + 1) if (end - 1) < len(index_map) else len(text) | |
| if any(f["type"] == family | |
| and not (real_end <= f["span"][0] or real_start >= f["span"][1]) | |
| for f in already): | |
| continue | |
| item = _finding(family, min(1.0, score * self.normalization_factor), | |
| real_start, real_end, text) | |
| item["normalized"] = True | |
| out.append(item) | |
| return out | |
| except Exception: | |
| return out | |
| def _has_injection_signal(self, text: str, min_score: float = 0.60) -> bool: | |
| """Guarantees: True when decoded text matches a strong injection pattern.""" | |
| try: | |
| for _family, score, rx in self.injection: | |
| if score >= min_score and rx.search(text): | |
| return True | |
| return False | |
| except Exception: | |
| return False | |
| def _detect_encoded(self, text: str) -> List[Dict[str, Any]]: | |
| """Guarantees: returns encoded/obfuscated payload findings, never raising.""" | |
| out: List[Dict[str, Any]] = [] | |
| try: | |
| if self.rx_base64 is not None: | |
| for m in self.rx_base64.finditer(text): | |
| chunk = m.group(0) | |
| if (self.rx_digest_context is not None | |
| and self.rx_digest_context.search(text[max(0, m.start() - 24):m.start()])): | |
| continue # a hash or fingerprint, not a smuggled payload | |
| if self._looks_base64(chunk): | |
| score = 0.55 + min(0.30, len(chunk) / 1200.0) | |
| if self.base64_rescan: | |
| decoded = decode_base64_text(chunk) | |
| if decoded and self._has_injection_signal(decoded): | |
| score = self.base64_attack_score | |
| out.append(_finding("encoded_payload", score, m.start(), m.end(), text)) | |
| if self.rx_zero_width is not None: | |
| for m in self.rx_zero_width.finditer(text): | |
| score = self.invisible_base_score + min(0.15, (m.end() - m.start()) / 20.0) | |
| out.append(_finding("encoded_payload", score, m.start(), m.end(), text)) | |
| if self.rx_control is not None: | |
| for m in self.rx_control.finditer(text): | |
| out.append(_finding("encoded_payload", 0.55, m.start(), m.end(), text)) | |
| if self.rx_rot13 is not None: | |
| for m in self.rx_rot13.finditer(text): | |
| out.append(_finding("encoded_payload", 0.45, m.start(), m.end(), text)) | |
| if self.rx_rot13_override is not None and len(text) < 20000: | |
| decoded = text.translate(ROT13_TABLE) | |
| for m in self.rx_rot13_override.finditer(decoded): | |
| out.append(_finding("encoded_payload", 0.80, m.start(), m.end(), text)) | |
| # Only report this when the attack is visible *after* decoding: | |
| # rot13 is a no-op on non-ASCII text, so a plain-text attack | |
| # would otherwise be relabelled as an encoded payload. | |
| if (decoded != text and self._has_injection_signal(decoded, 0.80) | |
| and not self._has_injection_signal(text, 0.80)): | |
| out.append(_finding("encoded_payload", self.rot13_rescan_score, | |
| 0, min(len(text), 200), text)) | |
| except Exception: | |
| return out | |
| return out | |
| def _detect_fence_flood(self, text: str) -> List[Dict[str, Any]]: | |
| """Guarantees: flags injected code-fence delimiters, never raising.""" | |
| out: List[Dict[str, Any]] = [] | |
| try: | |
| if self.rx_fence is None: | |
| return out | |
| fences = list(self.rx_fence.finditer(text)) | |
| if len(fences) >= 3 and len(fences) % 2 == 1: | |
| m = fences[-1] | |
| out.append(_finding("delimiter_escape", 0.70, m.start(), m.end(), text)) | |
| elif len(fences) >= 6: | |
| m = fences[0] | |
| out.append(_finding("delimiter_escape", 0.45, m.start(), m.end(), text)) | |
| except Exception: | |
| return out | |
| return out | |
| def detect_pii(self, text: str) -> List[Dict[str, Any]]: | |
| """Guarantees: returns PII findings; credit cards are only reported when Luhn passes.""" | |
| try: | |
| if not isinstance(text, str) or not text: | |
| return [] | |
| out: List[Dict[str, Any]] = [] | |
| email = self.pii.get("EMAIL", {}) | |
| if email.get("rx") is not None: | |
| for m in email["rx"].finditer(text): | |
| out.append(_finding("EMAIL", email.get("confidence", 0.95), m.start(), m.end(), text)) | |
| card = self.pii.get("CREDIT_CARD", {}) | |
| if card.get("rx") is not None: | |
| for m in card["rx"].finditer(text): | |
| digits = re.sub(r"\D", "", m.group(0)) | |
| if luhn_ok(digits): | |
| item = _finding("CREDIT_CARD", card.get("confidence", 0.96), | |
| m.start(), m.end(), text) | |
| item["match"] = "*" * max(0, len(digits) - 4) + digits[-4:] | |
| out.append(item) | |
| phone = self.pii.get("PHONE", {}) | |
| if phone.get("rx") is not None: | |
| for m in phone["rx"].finditer(text): | |
| raw = m.group(0) | |
| digits = re.sub(r"\D", "", raw) | |
| if not (self.phone_min_digits <= len(digits) <= self.phone_max_digits): | |
| continue | |
| stripped = raw.strip() | |
| ip_rx = self.pii.get("IP_ADDRESS", {}).get("rx") | |
| if ip_rx is not None and ip_rx.fullmatch(stripped): | |
| continue | |
| parts = stripped.split(".") | |
| if len(parts) == 4 and all(p.isdigit() for p in parts): | |
| continue | |
| has_sep = bool(re.search(r"[\-.() \u3000]", raw)) | |
| plus = raw.lstrip().startswith("+") | |
| if not (has_sep or plus or len(digits) in (10, 11)): | |
| continue | |
| if luhn_ok(digits) and len(digits) >= 13: | |
| continue | |
| conf = 0.88 if plus else (0.74 if has_sep else 0.55) | |
| out.append(_finding("PHONE", conf, m.start(), m.end(), text)) | |
| ip = self.pii.get("IP_ADDRESS", {}) | |
| if ip.get("rx") is not None: | |
| for m in ip["rx"].finditer(text): | |
| if self.skip_private_ipv4 and is_private_ipv4(m.group(0)): | |
| continue | |
| if (self.rx_version_context is not None | |
| and self.rx_version_context.search(text[max(0, m.start() - 14):m.start()])): | |
| continue | |
| out.append(_finding("IP_ADDRESS", ip.get("confidence", 0.72), | |
| m.start(), m.end(), text)) | |
| if ip.get("rx_v6") is not None: | |
| for m in ip["rx_v6"].finditer(text): | |
| if m.group(0).count(":") >= 4: | |
| out.append(_finding("IP_ADDRESS", ip.get("confidence_v6", 0.55), | |
| m.start(), m.end(), text)) | |
| dob = self.pii.get("DATE_OF_BIRTH", {}) | |
| if dob.get("rx") is not None: | |
| for m in dob["rx"].finditer(text): | |
| out.append(_finding("DATE_OF_BIRTH", dob.get("confidence", 0.85), | |
| m.start(), m.end(), text)) | |
| if dob.get("rx_bare") is not None: | |
| for m in dob["rx_bare"].finditer(text): | |
| out.append(_finding("DATE_OF_BIRTH", dob.get("confidence_bare", 0.50), | |
| m.start(), m.end(), text)) | |
| postal = self.pii.get("POSTAL_CODE", {}) | |
| if postal.get("rx") is not None: | |
| for m in postal["rx"].finditer(text): | |
| out.append(_finding("POSTAL_CODE", postal.get("confidence", 0.58), | |
| m.start(), m.end(), text)) | |
| if postal.get("rx_ctx") is not None: | |
| for m in postal["rx_ctx"].finditer(text): | |
| out.append(_finding("POSTAL_CODE", postal.get("confidence_ctx", 0.78), | |
| m.start(), m.end(), text)) | |
| return _suppress_overlaps(_dedupe_findings(out), self.pii_priority) | |
| except Exception: | |
| return [] | |
| def detect_secrets(self, text: str) -> List[Dict[str, Any]]: | |
| """Guarantees: returns secret findings from prefixes, JWTs and high-entropy strings.""" | |
| try: | |
| if not isinstance(text, str) or not text: | |
| return [] | |
| out: List[Dict[str, Any]] = [] | |
| for kind, score, rx in self.secrets: | |
| for m in rx.finditer(text): | |
| raw = m.group(0) | |
| if self.rx_placeholder is not None and self.rx_placeholder.search(raw): | |
| continue # documentation placeholder, not a live credential | |
| item = _finding(kind, score, m.start(), m.end(), text) | |
| item["match"] = (raw[:6] + "***" + raw[-4:]) if len(raw) > 14 else "***" | |
| out.append(item) | |
| if self.rx_entropy is not None: | |
| for m in self.rx_entropy.finditer(text): | |
| chunk = m.group(0) | |
| if not any(c.isdigit() for c in chunk) or not any(c.isalpha() for c in chunk): | |
| continue | |
| if self.rx_hex_hash is not None and self.rx_hex_hash.match(chunk): | |
| continue # digest, not a credential | |
| if self.rx_placeholder is not None and self.rx_placeholder.search(chunk): | |
| continue | |
| if (self.rx_digest_context is not None | |
| and self.rx_digest_context.search(text[max(0, m.start() - 24):m.start()])): | |
| continue # "SHA256:..." is a fingerprint, not a secret | |
| entropy = shannon_entropy(chunk) | |
| if entropy < self.entropy_threshold: | |
| continue | |
| conf = _clip(0.40 + (entropy - self.entropy_threshold) * 0.35, 0.40, 0.85) | |
| item = _finding("HIGH_ENTROPY_STRING", conf, m.start(), m.end(), text) | |
| item["entropy"] = round(entropy, 3) | |
| item["match"] = chunk[:6] + "***" + chunk[-4:] | |
| out.append(item) | |
| return _suppress_overlaps(_dedupe_findings(out), self.secret_priority) | |
| except Exception: | |
| return [] | |
| # ------------------------------------------------------- redaction | |
| def redact(self, text: str, findings: Any) -> Tuple[str, Dict[str, str]]: | |
| """Guarantees: returns (masked_text, mapping); the mapping is returned, never logged.""" | |
| try: | |
| if not isinstance(text, str) or not text: | |
| return "", {} | |
| flat: List[Tuple[str, Dict[str, Any]]] = [] | |
| if isinstance(findings, dict): | |
| for category, items in findings.items(): | |
| if isinstance(items, list): | |
| for item in items: | |
| if isinstance(item, dict) and "span" in item: | |
| flat.append((str(category), item)) | |
| elif isinstance(findings, list): | |
| for item in findings: | |
| if isinstance(item, dict) and "span" in item: | |
| flat.append((str(item.get("category", "risk")), item)) | |
| if not flat: | |
| return text, {} | |
| prefixes = {"pii": "PII", "secrets": "SECRET", "injection": "RISK"} | |
| flat.sort(key=lambda pair: (pair[1]["span"][0], -(pair[1]["span"][1] - pair[1]["span"][0]))) | |
| claimed: List[Tuple[int, int]] = [] | |
| planned: List[Tuple[int, int, str]] = [] | |
| mapping: Dict[str, str] = {} | |
| counters: Dict[str, int] = {} | |
| assigned: Dict[Tuple[str, str], str] = {} | |
| for category, item in flat: | |
| start, end = int(item["span"][0]), int(item["span"][1]) | |
| if start < 0 or end > len(text) or end <= start: | |
| continue | |
| if any(not (end <= cs or start >= ce) for cs, ce in claimed): | |
| continue | |
| claimed.append((start, end)) | |
| label = "{}_{}".format( | |
| prefixes.get(category, "RISK"), | |
| re.sub(r"[^A-Z0-9]+", "_", str(item.get("type", "UNKNOWN")).upper()).strip("_") or "UNKNOWN") | |
| original = text[start:end] | |
| placeholder = assigned.get((label, original)) | |
| if placeholder is None: | |
| counters[label] = counters.get(label, 0) + 1 | |
| placeholder = "<{}_{}>".format(label, counters[label]) | |
| assigned[(label, original)] = placeholder | |
| mapping[placeholder] = original | |
| planned.append((start, end, placeholder)) | |
| masked = text | |
| for start, end, placeholder in sorted(planned, key=lambda t: -t[0]): | |
| masked = masked[:start] + placeholder + masked[end:] | |
| return masked, mapping | |
| except Exception: | |
| return text if isinstance(text, str) else "", {} | |
| # ------------------------------------------------------ assessment | |
| def assess(self, text: str, thresholds: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: | |
| """Guarantees: always returns a verdict of allow/sanitize/block plus scored findings.""" | |
| result: Dict[str, Any] = { | |
| "verdict": "allow", "severity": 0.0, | |
| "scores": {"injection": 0.0, "pii": 0.0, "secrets": 0.0}, | |
| "counts": {"injection": 0, "pii": 0, "secrets": 0}, | |
| "risk_types": [], "findings": {"injection": [], "pii": [], "secrets": []}, | |
| "thresholds": dict(self.thresholds), "rules_version": self.version, "error": None, | |
| } | |
| try: | |
| merged = dict(self.thresholds) | |
| if isinstance(thresholds, dict): | |
| for key, value in thresholds.items(): | |
| try: | |
| merged[str(key)] = float(value) | |
| except Exception: | |
| continue | |
| result["thresholds"] = merged | |
| min_score = _clip(merged.get("min_score", 0.30)) | |
| text = text if isinstance(text, str) else "" | |
| injection = [f for f in self.detect_injection(text) if f["score"] >= min_score] | |
| pii = [f for f in self.detect_pii(text) if f["score"] >= min(min_score, 0.5)] | |
| secrets = [f for f in self.detect_secrets(text) if f["score"] >= min(min_score, 0.4)] | |
| inj_max = max((f["score"] for f in injection), default=0.0) | |
| families = {f["type"] for f in injection} | |
| injection_risk = 0.0 | |
| if injection: | |
| injection_risk = _clip(inj_max * (1.0 + 0.06 * (len(families) - 1)) + 0.02 * (len(injection) - 1)) | |
| secret_risk = _clip(max((f["score"] for f in secrets), default=0.0)) | |
| pii_risk = _clip(min(0.50, 0.22 + 0.08 * len(pii))) if pii else 0.0 | |
| severity = _clip(max(injection_risk, secret_risk, pii_risk)) | |
| if injection_risk >= merged.get("block", 0.80): | |
| verdict = "block" | |
| elif secrets or pii or severity >= merged.get("sanitize", 0.35): | |
| verdict = "sanitize" | |
| else: | |
| verdict = "allow" | |
| result.update({ | |
| "verdict": verdict, | |
| "severity": round(severity, 3), | |
| "scores": {"injection": round(injection_risk, 3), "pii": round(pii_risk, 3), | |
| "secrets": round(secret_risk, 3)}, | |
| "counts": {"injection": len(injection), "pii": len(pii), "secrets": len(secrets)}, | |
| "risk_types": sorted(families | {f["type"] for f in pii} | {f["type"] for f in secrets}), | |
| "findings": {"injection": injection, "pii": pii, "secrets": secrets}, | |
| }) | |
| return result | |
| except Exception as exc: | |
| result["error"] = {"code": "assess_failed", "message": str(exc)[:300]} | |
| return result | |
| def _cli(argv: List[str]) -> int: | |
| """Guarantees: prints one JSON assessment for argv text or stdin and returns an exit code.""" | |
| text = " ".join(argv[1:]) if len(argv) > 1 else sys.stdin.read() | |
| engine = RuleEngine.load_default() | |
| report = engine.assess(text) | |
| print(json.dumps(report, ensure_ascii=False, indent=2)) | |
| return 0 if report["verdict"] != "block" else 2 | |
| if __name__ == "__main__": | |
| raise SystemExit(_cli(sys.argv)) | |