surrogateshield 0.1.0__py3-none-any.whl

surrogateshield/__init__.py

@@ -0,0 +1,264 @@
+"""
+SurrogateShield — Privacy-preserving PII proxy for LLMs.
+
+Intercepts text before it reaches any LLM, replaces all PII with realistic
+fake surrogates, and restores the real values in the LLM response.
+
+Public API
+──────────
+    import surrogateshield as shield
+
+    shield.config(pii_off=["phone", "location"])
+    sanitized = shield.mask(user_text)
+    response  = llm.chat(sanitized)
+    restored  = shield.unmask(response)
+    shield.flush()
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Dict, List, Optional
+
+from ._state import cfg, session
+from . import _display, _response_parser
+from .core.detection import pipeline as _pipeline
+from .core.detection import service_query as _service_query
+from .core.reconstruction.resolve import ResolvePass as _ResolvePass
+
+__version__ = "0.1.0"
+__all__ = ["config", "scan", "pii_finder", "mask", "unmask", "flush"]
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# config()
+# ─────────────────────────────────────────────────────────────────────────────
+
+def config(
+    detailed_view: bool = True,
+    pii_mem: str = "temp",
+    pii_off=None,
+    service: bool = True,
+    spacy_model: str = "en_core_web_lg",
+    context_guard_enabled: bool = True,
+    entity_trace_high_threshold: float = 0.85,
+    entity_trace_low_threshold: float = 0.60,
+    context_guard_threshold: float = 0.70,
+    entity_trace_fallback_threshold: float = 0.65,
+    fuzzy_threshold: int = 85,
+) -> None:
+    """
+    Configure SurrogateShield.
+
+    Args:
+        detailed_view:                  Print detection/masking tables to stdout.
+        pii_mem:                        "temp" for in-memory session (default), or
+                                        a directory path for encrypted persistent storage.
+        pii_off:                        PII types to detect but NOT replace.
+                                        Accepts type names or aliases:
+                                        "phone", "name", "location", "org", "email",
+                                        "ssn", "dob", "address", "zip", "postcode",
+                                        "credit_card", "ip_address", "api_key",
+                                        "crypto", "bank", "license", "gender_indicator".
+        service:                        Enable service-query detection (address fuzzing
+                                        instead of full replacement for map queries).
+        spacy_model:                    spaCy model name for named entity recognition.
+        context_guard_enabled:          Enable the HuggingFace NER second-pass.
+        entity_trace_high_threshold:    spaCy score ≥ this → confirmed entity.
+        entity_trace_low_threshold:     spaCy score ≥ this → borderline entity.
+        context_guard_threshold:        ContextGuard score ≥ this → confirmed.
+        entity_trace_fallback_threshold: Promotion threshold when ContextGuard is off.
+        fuzzy_threshold:                rapidfuzz partial_ratio threshold for unmask().
+
+    Raises:
+        ValueError: If pii_mem is not "temp" and the path does not exist or is
+                    not a directory.
+    """
+    if pii_off is None:
+        pii_off = []
+
+    if pii_mem != "temp":
+        if not os.path.isdir(pii_mem):
+            raise ValueError(
+                f"pii_mem path does not exist or is not a directory: {pii_mem!r}"
+            )
+
+    cfg.detailed_view = detailed_view
+    cfg.pii_mem = pii_mem
+    cfg.pii_off = list(pii_off)
+    cfg.service = service
+    cfg.spacy_model = spacy_model
+    cfg.context_guard_enabled = context_guard_enabled
+    cfg.entity_trace_high_threshold = entity_trace_high_threshold
+    cfg.entity_trace_low_threshold = entity_trace_low_threshold
+    cfg.context_guard_threshold = context_guard_threshold
+    cfg.entity_trace_fallback_threshold = entity_trace_fallback_threshold
+    cfg.fuzzy_threshold = fuzzy_threshold
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# scan()  /  pii_finder
+# ─────────────────────────────────────────────────────────────────────────────
+
+def scan(text: str) -> Dict[str, str]:
+    """
+    Detect all PII in *text* without modifying anything.
+
+    Runs the full detection cascade (PatternScan → EntityTrace → ContextGuard)
+    and returns every detected entity regardless of pii_off settings.
+    Does NOT update the session shadow map.
+
+    Args:
+        text: Any string to scan for PII.
+
+    Returns:
+        Dict mapping detected_value → pii_type_string.
+        Example: {"john@example.com": "email", "John Smith": "PERSON"}
+    """
+    confirmed, _ = _pipeline.run_cascade(
+        text=text,
+        skip_values=None,
+        skip_location_entities=False,
+        pii_off=None,  # scan is always comprehensive
+        spacy_model=cfg.spacy_model,
+        context_guard_enabled=cfg.context_guard_enabled,
+        entity_trace_high_threshold=cfg.entity_trace_high_threshold,
+        entity_trace_low_threshold=cfg.entity_trace_low_threshold,
+        context_guard_threshold=cfg.context_guard_threshold,
+        entity_trace_fallback_threshold=cfg.entity_trace_fallback_threshold,
+    )
+
+    if cfg.detailed_view:
+        _display.show_scan_results(confirmed, cfg.pii_off)
+
+    return {ent.text: ent.type for ent in confirmed}
+
+
+# Alias
+pii_finder = scan
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# mask()
+# ─────────────────────────────────────────────────────────────────────────────
+
+def mask(text: str) -> str:
+    """
+    Replace all PII in *text* with realistic fake surrogates.
+
+    The original→surrogate mapping is stored in the session shadow map so
+    that unmask() can restore the real values from the LLM response.
+
+    Args:
+        text: The text to sanitize before sending to an LLM.
+
+    Returns:
+        Sanitized text with PII replaced by surrogates.
+    """
+    skip_values = None
+    skip_location_entities = False
+
+    # Service-query detection: fuzz addresses, suppress location entities
+    if cfg.service and _service_query.is_service_query(text):
+        text, fuzz_map = _service_query.fuzz_addresses(text, verify=True)
+        skip_location_entities = True
+        skip_values = set(fuzz_map.values())
+
+    # Run detection cascade
+    confirmed, _ = _pipeline.run_cascade(
+        text=text,
+        skip_values=skip_values,
+        skip_location_entities=skip_location_entities,
+        pii_off=cfg.pii_off,
+        spacy_model=cfg.spacy_model,
+        context_guard_enabled=cfg.context_guard_enabled,
+        entity_trace_high_threshold=cfg.entity_trace_high_threshold,
+        entity_trace_low_threshold=cfg.entity_trace_low_threshold,
+        context_guard_threshold=cfg.context_guard_threshold,
+        entity_trace_fallback_threshold=cfg.entity_trace_fallback_threshold,
+    )
+
+    # Deduplicate
+    confirmed = _pipeline.deduplicate(confirmed)
+
+    if not confirmed:
+        if cfg.detailed_view:
+            _display.show_mask_results([], {})
+        return text
+
+    # Reuse surrogates for originals already seen in this session
+    existing_shadow = session.get_shadow_map().get_all()  # {surrogate: original}
+    original_to_surrogate = {v: k for k, v in existing_shadow.items()}
+
+    already_mapped = [e for e in confirmed if e.text in original_to_surrogate]
+    new_entities   = [e for e in confirmed if e.text not in original_to_surrogate]
+
+    surrogate_map = {e.text: original_to_surrogate[e.text] for e in already_mapped}
+    if new_entities:
+        surrogate_map.update(session.get_mimic().generate_all(new_entities))
+
+    # Apply substitutions (longest match first to avoid substring collisions)
+    sanitized = text
+    for original, surrogate in sorted(surrogate_map.items(), key=lambda x: len(x[0]), reverse=True):
+        sanitized = sanitized.replace(original, surrogate)
+
+    # Store inverted map (surrogate → original) in session shadow map
+    inverted = {v: k for k, v in surrogate_map.items()}
+    session.get_shadow_map().update(inverted)
+
+    if cfg.detailed_view:
+        _display.show_mask_results(confirmed, surrogate_map)
+
+    return sanitized
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# unmask()
+# ─────────────────────────────────────────────────────────────────────────────
+
+def unmask(response) -> str:
+    """
+    Restore original PII values in the LLM *response*.
+
+    Extracts text from any major LLM SDK response object (Anthropic, OpenAI,
+    Gemini) or accepts a plain string, then replaces surrogates with the
+    originals stored in the session shadow map.
+
+    Args:
+        response: An LLM SDK response object or a plain string.
+
+    Returns:
+        Response text with surrogates replaced by the original PII values.
+    """
+    text = _response_parser.extract_text(response)
+    shadow_map = session.get_shadow_map().get_all()
+
+    resolver = _ResolvePass()
+    restored = resolver.resolve(
+        response_text=text,
+        shadow_map=shadow_map,
+        fuzzy_threshold=cfg.fuzzy_threshold,
+    )
+
+    if cfg.detailed_view:
+        # Count how many surrogates were actually replaced
+        replaced_count = sum(1 for s in shadow_map if s not in restored or s not in text)
+        _display.show_unmask_results(len(shadow_map))
+
+    return restored
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# flush()
+# ─────────────────────────────────────────────────────────────────────────────
+
+def flush() -> None:
+    """
+    Clear the session: discard all surrogate mappings and reset the session id.
+
+    Call this after a conversation ends to ensure surrogate mappings from
+    one session cannot bleed into the next.
+    """
+    session.reset()
+    if cfg.detailed_view:
+        _display.show_flush()

surrogateshield/_display.py

@@ -0,0 +1,101 @@
+"""
+surrogateshield/_display.py — Output display helpers
+
+Uses Rich tables when available; falls back to plain print otherwise.
+"""
+
+from __future__ import annotations
+
+try:
+    from rich.console import Console as _Console
+    from rich.table import Table as _Table
+    _console = _Console()
+    HAS_RICH = True
+except ImportError:
+    HAS_RICH = False
+
+
+def show_scan_results(entities: list, pii_off: list) -> None:
+    """Print a table of detected PII entities from scan()."""
+    if not entities:
+        print("[SurrogateShield] No PII detected.")
+        return
+
+    pii_off_lower = {t.lower() for t in (pii_off or [])}
+
+    if HAS_RICH:
+        table = _Table(title="[bold cyan]SurrogateShield — Scan Results[/bold cyan]", show_lines=True)
+        table.add_column("Detected Value", style="red bold")
+        table.add_column("Type", style="yellow")
+        table.add_column("Score", style="white")
+        table.add_column("Source", style="dim")
+        for ent in entities:
+            skipped = ent.type.lower() in pii_off_lower
+            note = "[dim](skipped — pii_off)[/dim]" if skipped else ""
+            table.add_row(
+                ent.text,
+                ent.type,
+                f"{ent.score:.2f}",
+                ent.source,
+                note,
+            )
+        _console.print(table)
+    else:
+        print("\n[SurrogateShield] Scan Results")
+        print(f"{'Detected Value':<30} {'Type':<20} {'Score':<8} {'Source':<10}")
+        print("-" * 70)
+        for ent in entities:
+            skipped = ent.type.lower() in pii_off_lower
+            note = " (skipped — pii_off)" if skipped else ""
+            print(f"{ent.text:<30} {ent.type:<20} {ent.score:<8.2f} {ent.source:<10}{note}")
+        print()
+
+
+def show_mask_results(entities: list, surrogate_map: dict) -> None:
+    """Print a table showing original PII and their surrogates."""
+    if not entities:
+        return
+
+    if HAS_RICH:
+        table = _Table(title="[bold cyan]SurrogateShield — Masked[/bold cyan]", show_lines=True)
+        table.add_column("Original", style="red bold")
+        table.add_column("Type", style="yellow")
+        table.add_column("Score", style="white")
+        table.add_column("Source", style="dim")
+        table.add_column("Surrogate", style="green bold")
+        for ent in entities:
+            surrogate = surrogate_map.get(ent.text, "[dim]—[/dim]")
+            table.add_row(
+                ent.text,
+                ent.type,
+                f"{ent.score:.2f}",
+                ent.source,
+                surrogate,
+            )
+        _console.print(table)
+    else:
+        print("\n[SurrogateShield] Mask Results")
+        print(f"{'Original':<30} {'Type':<20} {'Score':<8} {'Source':<10} {'Surrogate':<30}")
+        print("-" * 100)
+        for ent in entities:
+            surrogate = surrogate_map.get(ent.text, "—")
+            print(f"{ent.text:<30} {ent.type:<20} {ent.score:<8.2f} {ent.source:<10} {surrogate:<30}")
+        print()
+
+
+def show_unmask_results(restored_count: int) -> None:
+    """Print a one-liner confirming how many surrogates were restored."""
+    msg = f"[SurrogateShield] Restored {restored_count} surrogate(s)"
+    if HAS_RICH:
+        _console.print(f"[green]{msg}[/green]")
+    else:
+        print(msg)
+
+
+def show_flush() -> None:
+    """Print a one-liner confirming the session was cleared."""
+    msg = "[SurrogateShield] Session memory cleared"
+    if HAS_RICH:
+        _console.print(f"[yellow]{msg}[/yellow]")
+    else:
+        print(msg)

surrogateshield/_response_parser.py

@@ -0,0 +1,50 @@
+"""
+surrogateshield/_response_parser.py — LLM response text extractor
+
+Extracts plain text from any major LLM SDK response object.
+"""
+
+from __future__ import annotations
+
+
+def extract_text(response) -> str:
+    """
+    Extract the text content from an LLM response object.
+
+    Tries in order:
+    1. Anthropic style: response.content[0].text
+    2. OpenAI style:    response.choices[0].message.content
+    3. Gemini style:    response.text  (has .text but no .choices)
+    4. Fallback:        str(response)
+
+    Args:
+        response: Any LLM SDK response, or a plain string.
+
+    Returns:
+        The extracted text as a string.
+    """
+    if isinstance(response, str):
+        return response
+
+    # Anthropic: response.content is a list of content blocks
+    try:
+        if hasattr(response, "content") and isinstance(response.content, list):
+            return response.content[0].text
+    except (AttributeError, IndexError):
+        pass
+
+    # OpenAI: response.choices[0].message.content
+    try:
+        if hasattr(response, "choices"):
+            return response.choices[0].message.content
+    except (AttributeError, IndexError):
+        pass
+
+    # Gemini: response.text (but no .choices attribute)
+    try:
+        if hasattr(response, "text") and not hasattr(response, "choices"):
+            return response.text
+    except AttributeError:
+        pass
+
+    return str(response)

surrogateshield/_state.py

@@ -0,0 +1,65 @@
+"""
+surrogateshield/_state.py — Module-level singletons
+
+Holds cfg (Config) and session (Session) as module-level singletons
+so all public API calls share the same state within a Python process.
+"""
+
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+
+@dataclass
+class _Config:
+    """Holds all library-wide configuration values."""
+    detailed_view: bool = True
+    pii_mem: str = "temp"
+    pii_off: List[str] = field(default_factory=list)
+    service: bool = True
+    spacy_model: str = "en_core_web_lg"
+    context_guard_enabled: bool = True
+    entity_trace_high_threshold: float = 0.85
+    entity_trace_low_threshold: float = 0.60
+    context_guard_threshold: float = 0.70
+    entity_trace_fallback_threshold: float = 0.65
+    fuzzy_threshold: int = 85
+
+
+class _Session:
+    """Holds per-session state: session id, shadow map, and mimic generator."""
+
+    def __init__(self) -> None:
+        self.id: str = str(uuid.uuid4())
+        self._shadow_map = None
+        self._mimic = None
+
+    def reset(self) -> None:
+        """Clear all session state and generate a new session id."""
+        if self._shadow_map is not None:
+            self._shadow_map.flush()
+        self._shadow_map = None
+        self._mimic = None
+        self.id = str(uuid.uuid4())
+
+    def get_mimic(self):
+        """Return the MimicGen for this session, creating it if needed."""
+        if self._mimic is None:
+            from .core.generation.mimic import MimicGen
+            self._mimic = MimicGen()
+        return self._mimic
+
+    def get_shadow_map(self):
+        """Return the ShadowMap for this session, creating it if needed."""
+        if self._shadow_map is None:
+            from .core.storage.shadow_map import ShadowMap
+            storage_dir = None if cfg.pii_mem == "temp" else cfg.pii_mem
+            self._shadow_map = ShadowMap(self.id, storage_dir=storage_dir)
+        return self._shadow_map
+
+
+# Module-level singletons
+cfg = _Config()
+session = _Session()

surrogateshield/core/detection/context_guard.py

@@ -0,0 +1,178 @@
+"""
+detection/context_guard.py — ContextGuard
+
+NER-based detection of named entities using a local HuggingFace model
+(dslim/distilbert-NER by default, ~250 MB).
+
+This module detects named entities in the text that PatternScan and
+EntityTrace missed.  It does NOT decide whether a geographic entity is
+PII — that decision is made in detection/pipeline.py.
+
+Tokenization artefact handling:
+  distilbert word-piece tokenisation sometimes produces tokens like ". Sun"
+  or "##wick".  Both are stripped before emitting entities.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+from typing import Dict, List, Tuple
+
+from ..entities import DetectedEntity
+
+logger = logging.getLogger(__name__)
+
+# Cache pipelines keyed by model name
+_ner_pipelines: Dict[str, object] = {}
+
+
+def _get_ner(model_name: str = "dslim/distilbert-NER"):
+    """Lazy-load and cache the HuggingFace NER pipeline by model name."""
+    global _ner_pipelines
+    if model_name in _ner_pipelines:
+        return _ner_pipelines[model_name]
+    try:
+        from transformers import pipeline as hf_pipeline
+        pipeline = hf_pipeline(
+            "ner",
+            model=model_name,
+            aggregation_strategy="simple",
+            device=-1,
+        )
+        _ner_pipelines[model_name] = pipeline
+        logger.info(f"[ContextGuard] Loaded NER model: {model_name}")
+    except ImportError:
+        logger.warning(
+            "[ContextGuard] transformers not installed — skipping. "
+            "Run: pip install transformers torch"
+        )
+        _ner_pipelines[model_name] = None
+    except Exception as exc:
+        logger.warning(f"[ContextGuard] Failed to load NER model: {exc}")
+        _ner_pipelines[model_name] = None
+    return _ner_pipelines[model_name]
+
+
+_LABEL_MAP = {
+    "PER":    "PERSON",
+    "PERSON": "PERSON",
+    "ORG":    "ORG",
+    "LOC":    "LOC",
+    "GPE":    "GPE",
+    "MISC":   "MISC",
+}
+
+_KEEP_LABELS = {"PER", "PERSON", "ORG", "LOC", "GPE"}
+
+_CG_BLOCKLIST: frozenset = frozenset({
+    "dr", "mr", "mrs", "ms", "prof", "professor", "rev", "sr", "jr",
+    "sir", "lord", "dame", "capt", "lt", "sgt", "col", "gen",
+    "de", "le", "la", "el", "al", "van", "von",
+})
+
+
+def _clean_token(raw: str) -> str:
+    """Strip HuggingFace word-piece artefacts and leading punctuation."""
+    text = raw.replace("##", "")
+    text = re.sub(r'^[^A-Za-z0-9]+', '', text)
+    return text.strip()
+
+
+def guard(
+    remaining_text: str,
+    borderline_entities: List[DetectedEntity],
+    model_name: str = "dslim/distilbert-NER",
+    enabled: bool = True,
+    confidence_threshold: float = 0.70,
+) -> Tuple[List[DetectedEntity], List[DetectedEntity]]:
+    """
+    Run NER on remaining_text and verify borderline_entities.
+
+    Args:
+        remaining_text:      Text not covered by PatternScan / EntityTrace.
+        borderline_entities: Entities EntityTrace was uncertain about.
+        model_name:          HuggingFace model to use for NER inference.
+        enabled:             If False, skip NER inference entirely.
+        confidence_threshold: Minimum score to promote a borderline entity.
+
+    Returns:
+        Tuple of (confirmed_entities, uncertain_entities).
+    """
+    confirmed: List[DetectedEntity] = []
+    uncertain: List[DetectedEntity] = []
+
+    # Verify borderline entities from EntityTrace against the threshold
+    for ent in borderline_entities:
+        if ent.score >= confidence_threshold:
+            confirmed.append(ent)
+            logger.debug(
+                f"[ContextGuard] Verified borderline: {ent.text!r} "
+                f"({ent.type}, score={ent.score:.2f})"
+            )
+        else:
+            uncertain.append(ent)
+
+    if not enabled:
+        return confirmed, uncertain
+
+    # Run NER on remaining text
+    clean = remaining_text.replace("█", " ").strip()
+    if not clean:
+        return confirmed, uncertain
+
+    ner = _get_ner(model_name)
+    if ner is None:
+        return confirmed, uncertain
+
+    try:
+        results = ner(clean)
+    except Exception as exc:
+        logger.warning(f"[ContextGuard] NER inference failed: {exc}")
+        return confirmed, uncertain
+
+    for r in results:
+        label = r.get("entity_group", r.get("entity", ""))
+        if label not in _KEEP_LABELS:
+            continue
+
+        entity_type = _LABEL_MAP.get(label, label)
+        score = float(r.get("score", 0.0))
+
+        raw_word = r.get("word", "")
+        text = _clean_token(raw_word)
+
+        if len(text) < 3:
+            logger.debug(
+                f"[ContextGuard] Skipping too-short token: {raw_word!r} → {text!r}"
+            )
+            continue
+
+        if text.lower() in _CG_BLOCKLIST:
+            logger.debug(f"[ContextGuard] Skipping blocklisted token: {text!r}")
+            continue
+
+        entity = DetectedEntity(
+            text=text,
+            start=r.get("start", 0),
+            end=r.get("end", len(text)),
+            type=entity_type,
+            score=score,
+            source="slm",
+        )
+
+        if score >= confidence_threshold:
+            confirmed.append(entity)
+            logger.debug(
+                f"[ContextGuard] Confirmed: {text!r} ({entity_type}, {score:.2f})"
+            )
+        else:
+            uncertain.append(entity)
+            logger.debug(
+                f"[ContextGuard] Uncertain: {text!r} ({entity_type}, {score:.2f})"
+            )
+
+    logger.info(
+        f"[ContextGuard] confirmed={len(confirmed)}, uncertain={len(uncertain)}"
+    )
+    return confirmed, uncertain