From 847e84578e079cd564bf0e083b006257b0c040e3 Mon Sep 17 00:00:00 2001 From: Chandan Kumar Date: Fri, 6 Nov 2026 18:50:46 +0530 Subject: [PATCH 1/2] fix: close privacy bypasses and preserve recovery vaults --- README.md | 7 +- SETUP.md | 78 ++++++--- examples/proxy_chat.py | 6 +- src/shieldprompt/cli.py | 103 ++++++++--- src/shieldprompt/decorator.py | 99 ++++++----- src/shieldprompt/detector/hybrid.py | 42 ++--- src/shieldprompt/detector/ner_engine.py | 8 +- src/shieldprompt/detector/presidio_engine.py | 42 +++-- src/shieldprompt/detector/regex_engine.py | 88 ++++++---- src/shieldprompt/entities.py | 3 +- src/shieldprompt/mcp_server.py | 121 +++++-------- src/shieldprompt/middleware.py | 170 ++++++++++-------- src/shieldprompt/shield.py | 15 +- src/shieldprompt/vault.py | 26 ++- tests/test_cli_security.py | 156 +++++++++++++++++ tests/test_detection_security.py | 175 +++++++++++++++++++ tests/test_integration_security.py | 173 ++++++++++++++++++ tests/test_mcp_security.py | 46 +++++ tests/test_vault_security.py | 22 +++ 19 files changed, 1050 insertions(+), 330 deletions(-) create mode 100644 tests/test_cli_security.py create mode 100644 tests/test_detection_security.py create mode 100644 tests/test_integration_security.py create mode 100644 tests/test_mcp_security.py create mode 100644 tests/test_vault_security.py diff --git a/README.md b/README.md index 04cd0e5..0332f60 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,11 @@ cards, and other sensitive values for reversible placeholders. When the model's answer comes back, the placeholders are swapped back so your users still see the real values. -**The model never sees the real data. Your users never see placeholders.** +**Mask locally before sending data to a model. Review output: detection can miss PII.** + +The MCP helper does not intercept prompts before the model sees them. Keep +originals and vaults outside model-accessible files and restore responses locally. +Streaming middleware responses retain placeholders for local-client restoration. --- @@ -67,3 +71,4 @@ requests reviewed by the project owner. ## License MIT + diff --git a/SETUP.md b/SETUP.md index af85fef..d0af5c8 100644 --- a/SETUP.md +++ b/SETUP.md @@ -17,8 +17,10 @@ Claude Code, and using it from the CLI or Python. ## 1. Setup Claude Code -ShieldPrompt includes an MCP (Model Context Protocol) server so Claude -Code can mask PII in any conversation, automatically. +ShieldPrompt includes an MCP helper, not a pre-send privacy filter. A model sees +its prompt before it invokes a tool. Do not paste secrets into Claude and expect +`shield_mask` to hide them retroactively. Use the local CLI or Python library +before text enters any model context; review masked output before sending. ### Prerequisites @@ -40,8 +42,14 @@ Code can mask PII in any conversation, automatically. pip install shieldprompt ``` -The first run will automatically download the spaCy `en_core_web_sm` -model (~12 MB). No separate download step needed. +Install the spaCy model explicitly before enabling contextual detection: + +```bash +python -m spacy download en_core_web_sm +``` + +ShieldPrompt does not install model packages during a request. Missing models +cause a clear failure rather than silently reducing privacy protection. ### Step 2 — Register the MCP server with Claude Code @@ -78,36 +86,31 @@ You should see `shieldprompt` in the output. ### Step 3 — Restart Claude Code Close and reopen your Claude Code session. The ShieldPrompt tools will -now be available: `shield_mask`, `shield_unmask`, `shield_inspect`, -`shield_vault`, `shield_clear`. +now be available: `shield_mask`, `shield_inspect`, `shield_clear`. -### Step 4 — Try it +### Step 4 — Keep originals outside the model context -In Claude Code, type: +Try MCP only with synthetic, non-sensitive examples. It does not automatically +protect a conversation, intercept files, or mask prompts before Claude receives +them. Inspection returns entity types and positions, never original values. +The MCP server deliberately does not expose unmask or vault-export tools. -> Mask the PII in this sentence: "Email alice@example.com and call -> +1-415-555-1234, my name is Chandan Kumar." +For real data, run masking in a local terminal first: -Claude will invoke `shield_mask` and reply with: - -``` -Email [EMAIL_ADDRESS_1] and call [PHONE_NUMBER_1], my name is [PERSON_1]. +```bash +shieldprompt mask --file input.txt --save-vault private.vault.json > masked.txt ``` -Now you can send that masked text to any external service safely. Ask -Claude to `shield_unmask` the reply when you're done to get the real -values back. +Review `masked.txt`, then send only that file's contents to the model. Keep +`input.txt` and `private.vault.json` out of the model's accessible files, prompts, +and tools. Restore the response locally, not through a model-visible tool: -Or if you want this in complete session. Just write - -~~~ -Use shield_mask for this complete session -~~~ - -This will work for enitre session. You write your prompt normally as you do. -Your ShieldPrompt will handle everything. +```bash +shieldprompt unmask --file response.txt --vault private.vault.json +``` -**You're live.** ShieldPrompt is active in Claude Code. +Detection is best effort: even a functioning detector can miss sensitive values. +Do not treat masked output as a guarantee that no personal data remains. --- @@ -272,6 +275,8 @@ Each request gets its own vault, so data never leaks across users. - `URL` - `AWS_KEY` - `API_KEY` (OpenAI, GitHub tokens) +- `DATE_OF_BIRTH` (common numeric date forms; not contextual birth-date inference) +- `MONEY` (common currency-symbol/code numeric amounts) **NER-based (Presidio + spaCy, contextual):** @@ -279,7 +284,6 @@ Each request gets its own vault, so data never leaks across users. - `ORGANIZATION` - `LOCATION` - `DATE` -- `MONEY` --- @@ -309,3 +313,23 @@ Worth post-processing if it matters. **Q: Does it slow things down?** Regex is microseconds per KB. Presidio + spaCy loads once (a couple of seconds) then runs fast. Network calls to the LLM dwarf both. + + +## Security behavior and compatibility + +- NER initialization or inference failure raises an error instead of returning + apparently protected text. `--no-ner` / `use_ner=False` explicitly selects + regex-only protection; names and locations are not covered in that mode. +- Vault outputs are created owner-only and existing vault files are not silently + overwritten. Choose a new vault path for each masking operation. In-place + masking saves the recovery vault before replacing the source; repeated masking + of tokenized files is refused. Keep the vault until restoration is verified. +- Never use the source path as the vault path. Vaults contain plaintext secrets: + filesystem permissions do not replace encryption or appropriate file isolation. +- Streaming responses are forwarded with placeholders intact. Restore them in a + trusted local client; automatic restoration applies only to bounded, complete, uncompressed JSON/text + responses (default limit 1 MiB). Oversized and other response types retain tokens. + Requests over the configurable 1 MiB default limit are rejected with HTTP 413. JSON responses are restored as structured values, then serialized. +- The decorator supports string arguments, including variadic wrappers, and + recursively restores strings in list/dict responses. SDK-specific objects + require a caller-owned adapter. diff --git a/examples/proxy_chat.py b/examples/proxy_chat.py index 904e24d..bd62a67 100644 --- a/examples/proxy_chat.py +++ b/examples/proxy_chat.py @@ -2,7 +2,8 @@ You talk to this script. It masks PII locally with ShieldPrompt, sends the masked prompt to Claude, then unmasks the reply before printing it. The -Anthropic API only ever sees tokens like [PERSON_1] and [PHONE_NUMBER_1]. +API receives detected values as tokens such as [PERSON_1]. Detection can +miss PII: review output and do not rely on this example as a privacy guarantee. Setup: pip install -e . # from repo root, installs shieldprompt @@ -36,7 +37,7 @@ def main() -> None: print("error: set ANTHROPIC_API_KEY", file=sys.stderr) sys.exit(1) - shield = Shield(use_ner=False) + shield = Shield(use_ner=True) client = Anthropic() history: list[dict] = [] @@ -73,3 +74,4 @@ def main() -> None: if __name__ == "__main__": main() + diff --git a/src/shieldprompt/cli.py b/src/shieldprompt/cli.py index a6db3d1..17c4a72 100644 --- a/src/shieldprompt/cli.py +++ b/src/shieldprompt/cli.py @@ -17,6 +17,9 @@ import argparse import json +import os +import re +import tempfile import sys from pathlib import Path @@ -31,15 +34,52 @@ def _load_vault(path: str) -> Vault: return Vault.from_dict(data) -def _save_vault(vault: Vault, path: str) -> None: - """Save vault mappings to a JSON file with owner-only permissions (0600).""" +def _sync_directory(path: Path) -> None: + """Make published directory entries durable before changing the source.""" + if os.name == "posix": + fd = os.open(str(path), os.O_RDONLY | getattr(os, "O_DIRECTORY", 0)) + try: + os.fsync(fd) + finally: + os.close(fd) + + +def _atomic_write(path: str, text: str, *, exclusive: bool = False) -> None: + """Publish a complete private file; never expose partially written secrets.""" p = Path(path) - p.write_text(json.dumps(vault.to_dict(), indent=2)) + fd, temporary = tempfile.mkstemp(prefix=".shieldprompt-", dir=str(p.parent)) try: - p.chmod(0o600) - except OSError: - # chmod isn't meaningful on every platform (e.g. Windows); ignore. - pass + with os.fdopen(fd, "w", encoding="utf-8", newline="") as stream: + stream.write(text) + stream.flush() + os.fsync(stream.fileno()) + if exclusive: + # Unlike replace(), link() fails atomically if any entry already exists, + # including a symlink. Keep old recovery mappings at all costs. + os.link(temporary, p) + else: + os.replace(temporary, p) + _sync_directory(p.parent) + finally: + if os.path.exists(temporary): + os.unlink(temporary) + + +def _save_vault(vault: Vault, path: str) -> None: + """Save a new owner-only vault without ever overwriting recovery mappings.""" + _atomic_write(path, json.dumps(vault.to_dict(), indent=2), exclusive=True) + + +def _validate_paths(file_path: str | None, vault_path: str, in_place: bool) -> None: + if not file_path: + return + source, vault = Path(file_path), Path(vault_path) + if source.resolve() == vault.resolve() or ( + source.exists() and vault.exists() and source.samefile(vault) + ): + raise ValueError("Source and vault must be different files") + if in_place and (source.is_symlink() or not source.is_file()): + raise ValueError("In-place operations require a regular, non-symlink source file") def _default_vault_path(file_path: str) -> str: @@ -125,14 +165,18 @@ def main(argv: list[str] | None = None, prog: str = "shieldprompt") -> None: "restore": "unmask", }.get(args.command, args.command) - if command == "mask": - _cmd_mask(args) - elif command == "unmask": - _cmd_unmask(args) - elif command == "inspect": - _cmd_inspect(args) - elif command == "map": - _cmd_map(args) + try: + if command == "mask": + _cmd_mask(args) + elif command == "unmask": + _cmd_unmask(args) + elif command == "inspect": + _cmd_inspect(args) + elif command == "map": + _cmd_map(args) + except (OSError, ValueError) as exc: + print(f"Error: {exc}", file=sys.stderr) + sys.exit(1) def _get_text(args) -> str: @@ -152,7 +196,16 @@ def _cmd_mask(args) -> None: print("Error: --in-place requires --file", file=sys.stderr) sys.exit(1) + vault_path = args.save_vault or ( + _default_vault_path(args.file) if args.in_place else None + ) + if vault_path: + _validate_paths(args.file, vault_path, args.in_place) + if os.path.lexists(vault_path): + raise ValueError("Vault already exists; preserve it and choose a new vault path") text = _get_text(args) + if args.in_place and re.search(r"\[[A-Z][A-Z_]*_\d+\]", text): + raise ValueError("Refusing to mask token-containing text in-place; restore it first") entities = {EntityType(e) for e in args.entities} if args.entities else None vault = Vault() shield = Shield( @@ -161,17 +214,15 @@ def _cmd_mask(args) -> None: vault=vault, ) masked = shield.mask(text) + # A durable recovery vault must exist before replacing or emitting masked text. + if vault_path: + _save_vault(vault, vault_path) if args.in_place: - Path(args.file).write_text(masked) + _atomic_write(args.file, masked) + print(f"Masked file in-place: {args.file}", file=sys.stderr) else: print(masked) - - should_save_vault = bool(args.save_vault or args.in_place) - if should_save_vault: - vault_path = args.save_vault or _default_vault_path(args.file) - _save_vault(vault, vault_path) - if args.in_place: - print(f"Masked file in-place: {args.file}", file=sys.stderr) + if vault_path: print(f"Vault saved to {vault_path}", file=sys.stderr) @@ -180,13 +231,14 @@ def _cmd_unmask(args) -> None: print("Error: --in-place requires --file", file=sys.stderr) sys.exit(1) - text = _get_text(args) vault_path = _resolve_vault_path(args.file, args.vault) + _validate_paths(args.file, vault_path, args.in_place) + text = _get_text(args) vault = _load_vault(vault_path) unmasked = vault.restore_text(text) if args.in_place: - Path(args.file).write_text(unmasked) + _atomic_write(args.file, unmasked) print(f"Unmasked file in-place: {args.file}", file=sys.stderr) else: print(unmasked) @@ -243,3 +295,4 @@ def retrace_main() -> None: if __name__ == "__main__": main() + diff --git a/src/shieldprompt/decorator.py b/src/shieldprompt/decorator.py index e1b933b..b9a4c7c 100644 --- a/src/shieldprompt/decorator.py +++ b/src/shieldprompt/decorator.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import copy import functools import inspect from typing import Any, Callable, Optional @@ -46,55 +47,73 @@ async def call_llm_async(prompt: str) -> str: return await anthropic_client.messages.create(prompt) """ entity_set = ( - {EntityType(e) for e in entities} if entities else DEFAULT_ENTITIES + {EntityType(e) for e in entities} if entities is not None else DEFAULT_ENTITIES ) def decorator(func: Callable) -> Callable: sig = inspect.signature(func) + # Detectors may load expensive models. Share only detector/configuration; + # every invocation gets its own vault, including concurrent calls. + template = Shield( + entities=entity_set, use_ner=use_ner, ner_model=ner_model, + ner_threshold=ner_threshold, + ) + def _create_shield() -> tuple[Shield, Vault]: vault = Vault() - shield = Shield( - entities=entity_set, - use_ner=use_ner, - ner_model=ner_model, - ner_threshold=ner_threshold, - vault=vault, - ) + shield = copy.copy(template) + shield._vault = vault return shield, vault def _mask_args( shield: Shield, args: tuple, kwargs: dict ) -> tuple[tuple, dict]: - """Mask the target parameter(s) in args/kwargs.""" bound = sig.bind(*args, **kwargs) bound.apply_defaults() - - # Mask the primary param - if param_name in bound.arguments: - val = bound.arguments[param_name] - if isinstance(val, str): - bound.arguments[param_name] = shield.mask(val) - else: - # Fall back: mask first string argument by position - for name, val in bound.arguments.items(): - if isinstance(val, str): - bound.arguments[name] = shield.mask(val) - break - - # Mask additional kwargs - if mask_kwargs: - for kw in mask_kwargs: - if kw in bound.arguments and isinstance( - bound.arguments[kw], str - ): - bound.arguments[kw] = shield.mask( - bound.arguments[kw] - ) - - new_args = tuple(bound.args) - new_kwargs = bound.kwargs - return new_args, new_kwargs + extra_name = next((name for name, p in sig.parameters.items() + if p.kind == inspect.Parameter.VAR_KEYWORD), None) + extras = bound.arguments.get(extra_name, {}) + targets = set() + + def select(name): + if name in bound.arguments and isinstance(bound.arguments[name], str): + targets.add((name, None)) + return True + if name in extras and isinstance(extras[name], str): + targets.add((extra_name, name)) + return True + return False + + if not select(param_name): + # Fallback is positional only, including strings inside *args. + for name, p in sig.parameters.items(): + value = bound.arguments.get(name) + if p.kind in (inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD): + if isinstance(value, str): + targets.add((name, None)) + break + elif p.kind == inspect.Parameter.VAR_POSITIONAL: + index = next((i for i, v in enumerate(value or ()) + if isinstance(v, str)), None) + if index is not None: + targets.add((name, index)) + break + for name in mask_kwargs or []: + select(name) + + # A target is masked exactly once even if it is listed repeatedly. + for name, key in sorted(targets, key=repr): + if key is None: + bound.arguments[name] = shield.mask(bound.arguments[name]) + elif isinstance(key, int): + values = list(bound.arguments[name]) + values[key] = shield.mask(values[key]) + bound.arguments[name] = tuple(values) + else: + bound.arguments[name][key] = shield.mask(bound.arguments[name][key]) + return tuple(bound.args), bound.kwargs def _unmask_result(shield: Shield, result: Any) -> Any: if not unmask_response: @@ -102,10 +121,11 @@ def _unmask_result(shield: Shield, result: Any) -> Any: if isinstance(result, str): return shield.unmask(result) if isinstance(result, dict): - return { - k: shield.unmask(v) if isinstance(v, str) else v - for k, v in result.items() - } + return {k: _unmask_result(shield, v) for k, v in result.items()} + if isinstance(result, list): + return [_unmask_result(shield, v) for v in result] + if isinstance(result, tuple): + return tuple(_unmask_result(shield, v) for v in result) return result if asyncio.iscoroutinefunction(func): @@ -138,3 +158,4 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any: return sync_wrapper return decorator + diff --git a/src/shieldprompt/detector/hybrid.py b/src/shieldprompt/detector/hybrid.py index 82aa331..a9b326e 100644 --- a/src/shieldprompt/detector/hybrid.py +++ b/src/shieldprompt/detector/hybrid.py @@ -17,7 +17,7 @@ class HybridDetector: NER backend selection order: 1. Presidio (presidio-analyzer + spaCy en_core_web_sm) — preferred 2. HuggingFace transformer (dslim/bert-base-NER) — fallback - 3. Regex only — if neither is installed + Initialization and inference fail closed if NER is requested. """ def __init__( @@ -31,7 +31,7 @@ def __init__( self._entities = entities self._ner_threshold = ner_threshold - regex_entities = entities & REGEX_ENTITIES if entities else None + regex_entities = entities & REGEX_ENTITIES if entities is not None else None self._regex = RegexDetector(entities=regex_entities) self._ner = None @@ -51,7 +51,7 @@ def _init_ner_backend( ner_model: str, ner_threshold: float, ): - ner_entities = entities & NER_ENTITIES if entities else None + ner_entities = entities & NER_ENTITIES if entities is not None else None # Try Presidio first. try: @@ -77,14 +77,15 @@ def _init_ner_backend( from .ner_engine import NERDetector logger.info("Using HuggingFace transformer NER backend: %s", ner_model) - return NERDetector(model_name=ner_model, entities=ner_entities) - except ImportError: - logger.info( - "No NER backend installed — regex-only detection. " - "Reinstall shieldprompt (Presidio is a core dependency) " - "or: pip install shieldprompt[ner]" - ) - return None + detector = NERDetector(model_name=ner_model, entities=ner_entities) + detector._load() + return detector + except Exception as exc: + raise RuntimeError( + "NER was requested but no usable backend could be loaded. " + "Install a supported NER model or explicitly select use_ner=False " + "for reduced, regex-only protection." + ) from exc def detect(self, text: str) -> list[Detection]: """Run all enabled engines and merge results.""" @@ -97,8 +98,8 @@ def detect(self, text: str) -> list[Detection]: d for d in ner_detections if d.score >= self._ner_threshold ] detections = self._merge(detections, ner_detections) - except Exception as e: - logger.warning("NER detection failed, using regex only: %s", e) + except Exception: + raise RuntimeError("NER detection failed; refusing to return partially masked text") from None return detections @@ -106,16 +107,5 @@ def detect(self, text: str) -> list[Detection]: def _merge( regex_dets: list[Detection], ner_dets: list[Detection] ) -> list[Detection]: - """Merge detections, preferring regex for overlapping spans.""" - merged = list(regex_dets) - regex_spans = {(d.start, d.end) for d in regex_dets} - - for nd in ner_dets: - overlaps = any( - nd.start < re and nd.end > rs for rs, re in regex_spans - ) - if not overlaps: - merged.append(nd) - - merged.sort(key=lambda d: d.start) - return merged + """Keep all spans; Shield merges their unions before replacement.""" + return sorted(regex_dets + ner_dets, key=lambda d: d.start) diff --git a/src/shieldprompt/detector/ner_engine.py b/src/shieldprompt/detector/ner_engine.py index cce6855..15cdefa 100644 --- a/src/shieldprompt/detector/ner_engine.py +++ b/src/shieldprompt/detector/ner_engine.py @@ -57,6 +57,11 @@ def _load(self) -> None: model=self._model_name, aggregation_strategy="simple", ) + labels = self._pipeline.model.config.id2label.values() + supported = {_LABEL_MAP[label] for label in labels if label in _LABEL_MAP} + if self._entities is not None and not self._entities <= supported: + self._pipeline = None + raise RuntimeError("NER model does not support all requested entity types") logger.info("NER model loaded: %s", self._model_name) except ImportError: raise ImportError( @@ -76,7 +81,7 @@ def detect(self, text: str) -> list[Detection]: mapped = _LABEL_MAP.get(label) if mapped is None: continue - if self._entities and mapped not in self._entities: + if self._entities is not None and mapped not in self._entities: continue value = text[ent["start"]:ent["end"]] @@ -95,3 +100,4 @@ def detect(self, text: str) -> list[Detection]: ) return detections + diff --git a/src/shieldprompt/detector/presidio_engine.py b/src/shieldprompt/detector/presidio_engine.py index 4f527c7..7f3217b 100644 --- a/src/shieldprompt/detector/presidio_engine.py +++ b/src/shieldprompt/detector/presidio_engine.py @@ -4,9 +4,8 @@ Presidio's built-in recognizers and context-aware scoring — handles PERSON, LOCATION, DATE, ORG, etc. far more reliably than a raw transformer. -Default model is `en_core_web_sm` (~12 MB). spaCy models aren't on PyPI, so -we auto-download on first use if the model isn't already installed — this -keeps `pip install shieldprompt` self-contained. +The configured spaCy model must be installed explicitly before use. +Detection never downloads or installs executable model packages. """ from __future__ import annotations @@ -69,11 +68,7 @@ def __init__( self._analyzer = None def _ensure_spacy_model(self) -> None: - """Make sure the spaCy model is installed; download it if not. - - spaCy models aren't on PyPI, so this is how we keep installation a - single `pip install` for the end user. - """ + """Require an already-installed spaCy model without network side effects.""" import spacy try: @@ -82,10 +77,10 @@ def _ensure_spacy_model(self) -> None: except OSError: pass - logger.info("spaCy model %r not found — downloading…", self._spacy_model) - from spacy.cli.download import download as spacy_download - - spacy_download(self._spacy_model) + raise RuntimeError( + "The configured spaCy model is not installed. Install it explicitly " + "before enabling NER." + ) def _load(self) -> None: if self._analyzer is not None: @@ -105,12 +100,32 @@ def _load(self) -> None: nlp_configuration={ "nlp_engine_name": "spacy", "models": [{"lang_code": "en", "model_name": self._spacy_model}], + # Do not inherit Presidio defaults that may ignore or lower + # ORGANIZATION below our threshold in some package versions. + "ner_model_configuration": { + "model_to_presidio_entity_mapping": { + "PER": "PERSON", "PERSON": "PERSON", + "LOC": "LOCATION", "LOCATION": "LOCATION", + "GPE": "LOCATION", "FAC": "LOCATION", + "ORG": "ORGANIZATION", "ORGANIZATION": "ORGANIZATION", + "DATE": "DATE_TIME", "TIME": "DATE_TIME", + "NORP": "NRP", + }, + "labels_to_ignore": [], + "low_score_entity_names": [], + "default_score": 0.85, + }, } ) self._analyzer = AnalyzerEngine( nlp_engine=provider.create_engine(), default_score_threshold=self._score_threshold, ) + supported = self._analyzer.get_supported_entities(language="en") + mapped = {_LABEL_MAP[label] for label in supported if label in _LABEL_MAP} + if self._entities is not None and not self._entities <= mapped: + self._analyzer = None + raise RuntimeError("Presidio does not support all requested entity types") logger.info("Presidio analyzer loaded (spacy=%s)", self._spacy_model) def detect(self, text: str) -> list[Detection]: @@ -124,7 +139,7 @@ def detect(self, text: str) -> list[Detection]: mapped = _LABEL_MAP.get(r.entity_type) if mapped is None: continue - if self._entities and mapped not in self._entities: + if self._entities is not None and mapped not in self._entities: continue value = text[r.start:r.end] if len(value.strip()) < 2: @@ -141,3 +156,4 @@ def detect(self, text: str) -> list[Detection]: ) ) return out + diff --git a/src/shieldprompt/detector/regex_engine.py b/src/shieldprompt/detector/regex_engine.py index 89976b2..0534695 100644 --- a/src/shieldprompt/detector/regex_engine.py +++ b/src/shieldprompt/detector/regex_engine.py @@ -15,29 +15,28 @@ _SPACE_CHARS = {"\u00a0", "\u2007", "\u202f", "\u2009", "\u200a", "\u2002", "\u2003"} -def _normalize_for_match(text: str) -> str: - """NFKC + strip zero-width + fold exotic dashes/spaces to ASCII. +def _normalize_with_offsets(text: str): + """Normalize each source character and retain its original span. - Preserves string length so detection offsets still index into the - original text. Zero-width chars are replaced with a space rather than - removed to keep positions aligned. + Compatibility expansions may change length. Format characters are removed + so inserting invisible characters cannot split a credential pattern. """ - out_chars: list[str] = [] - for ch in unicodedata.normalize("NFKC", text): - if ch in _ZERO_WIDTH: - out_chars.append(" ") - elif ch in _DASH_CHARS: - out_chars.append("-") - elif ch in _SPACE_CHARS: - out_chars.append(" ") - else: - out_chars.append(ch) - result = "".join(out_chars) - # NFKC can change length (e.g. ligatures). If so, fall back to original - # to keep offsets valid — we lose normalization but don't corrupt spans. - if len(result) != len(text): - return text - return result + chars, offsets = [], [] + for index, source in enumerate(text): + for ch in unicodedata.normalize("NFKC", source): + if ch in _ZERO_WIDTH or unicodedata.category(ch) == "Cf": + continue + if ch in _DASH_CHARS: + ch = "-" + elif ch in _SPACE_CHARS: + ch = " " + chars.append(ch) + offsets.append((index, index + 1)) + return "".join(chars), offsets + + +def _normalize_for_match(text: str) -> str: + return _normalize_with_offsets(text)[0] @dataclass @@ -129,10 +128,21 @@ def _luhn_check(number: str) -> bool: ( EntityType.API_KEY, re.compile( - r"\b(?:sk-[A-Za-z0-9]{20,}|ghp_[A-Za-z0-9]{36,}|gho_[A-Za-z0-9]{36,})\b" + r"(? None: def detect(self, text: str) -> list[Detection]: """Scan text and return all detected PII spans.""" detections: list[Detection] = [] - # Run regex over normalized text so Unicode dashes / zero-width chars - # don't let attackers slip past the patterns. Offsets stay aligned - # because _normalize_for_match preserves length. - search_text = _normalize_for_match(text) + # Match normalized text and map each span back to the original source. + search_text, offsets = _normalize_with_offsets(text) for entity_type, pattern, validator in _PATTERNS: - if self._entities and entity_type not in self._entities: + if self._entities is not None and entity_type not in self._entities: continue for match in pattern.finditer(search_text): # Store the original (un-normalized) substring so the vault # round-trips what the user actually wrote. - value = text[match.start():match.end()] - if validator and not validator(value): + start, end = offsets[match.start()][0], offsets[match.end() - 1][1] + value = text[start:end] + if validator and not validator(match.group()): continue detections.append( Detection( entity_type=entity_type, value=value, - start=match.start(), - end=match.end(), + start=start, + end=end, ) ) return detections + + + +def merge_overlaps(text: str, detections: list[Detection]) -> list[Detection]: + """Mask the union of every overlapping span; never discard exposed tails.""" + merged: list[Detection] = [] + for det in sorted(detections, key=lambda d: (d.start, -d.end)): + if not (0 <= det.start < det.end <= len(text)): + raise ValueError("Detector returned an invalid source span") + if merged and det.start < merged[-1].end: + previous = merged[-1] + previous.end = max(previous.end, det.end) + previous.value = text[previous.start:previous.end] + else: + merged.append(Detection(det.entity_type, text[det.start:det.end], + det.start, det.end, det.score)) + return merged diff --git a/src/shieldprompt/entities.py b/src/shieldprompt/entities.py index d2eff51..56ff9fe 100644 --- a/src/shieldprompt/entities.py +++ b/src/shieldprompt/entities.py @@ -32,6 +32,7 @@ class EntityType(str, Enum): EntityType.IP_ADDRESS, EntityType.IBAN, EntityType.DATE_OF_BIRTH, + EntityType.MONEY, EntityType.URL, EntityType.AWS_KEY, EntityType.API_KEY, @@ -42,7 +43,6 @@ class EntityType(str, Enum): EntityType.ORGANIZATION, EntityType.LOCATION, EntityType.DATE, - EntityType.MONEY, } ALL_ENTITIES = REGEX_ENTITIES | NER_ENTITIES @@ -59,3 +59,4 @@ class EntityType(str, Enum): EntityType.ORGANIZATION, EntityType.LOCATION, } + diff --git a/src/shieldprompt/mcp_server.py b/src/shieldprompt/mcp_server.py index 1810271..d593ece 100644 --- a/src/shieldprompt/mcp_server.py +++ b/src/shieldprompt/mcp_server.py @@ -20,6 +20,7 @@ import sys from typing import Any +from . import __version__ from .entities import EntityType from .shield import Shield from .vault import Vault @@ -55,7 +56,7 @@ def _handle_initialize(id: Any, params: dict) -> dict: "capabilities": {"tools": {}}, "serverInfo": { "name": "shieldprompt", - "version": "0.1.4", + "version": __version__, }, }) @@ -67,7 +68,8 @@ def _handle_tools_list(id: Any, params: dict) -> dict: "description": ( "Mask PII in text. Replaces emails, phone numbers, SSNs, " "credit cards, IPs, names, etc. with tokens like [EMAIL_ADDRESS_1]. " - "Use this BEFORE sending sensitive user data to any LLM or external service." + "This helper does not intercept model inputs. Mask locally with the CLI " + "before sending secrets to the model. Never pass secrets already excluded from its context." ), "inputSchema": { "type": "object", @@ -80,24 +82,6 @@ def _handle_tools_list(id: Any, params: dict) -> dict: "required": ["text"], }, }, - { - "name": "shield_unmask", - "description": ( - "Restore masked tokens back to their original PII values. " - "Use this to convert [EMAIL_ADDRESS_1] etc. back to real data " - "after receiving a response from an LLM." - ), - "inputSchema": { - "type": "object", - "properties": { - "text": { - "type": "string", - "description": "Text containing tokens like [PERSON_1] to restore.", - }, - }, - "required": ["text"], - }, - }, { "name": "shield_inspect", "description": ( @@ -115,17 +99,6 @@ def _handle_tools_list(id: Any, params: dict) -> dict: "required": ["text"], }, }, - { - "name": "shield_vault", - "description": ( - "Show the current vault mappings (token -> real value). " - "Useful for debugging or auditing what PII has been masked." - ), - "inputSchema": { - "type": "object", - "properties": {}, - }, - }, { "name": "shield_clear", "description": "Clear all vault mappings. Start fresh.", @@ -141,6 +114,10 @@ def _handle_tools_list(id: Any, params: dict) -> dict: def _handle_tools_call(id: Any, params: dict) -> dict: name = params.get("name", "") args = params.get("arguments", {}) + if not isinstance(args, dict): + return _make_error(id, -32602, "arguments must be an object") + if name in {"shield_mask", "shield_inspect"} and not isinstance(args.get("text"), str): + return _make_error(id, -32602, "text must be a string") try: if name == "shield_mask": @@ -150,36 +127,19 @@ def _handle_tools_call(id: Any, params: dict) -> dict: "content": [{"type": "text", "text": masked}], }) - elif name == "shield_unmask": - text = _session_vault.restore_text(args["text"]) - return _make_response(id, { - "content": [{"type": "text", "text": text}], - }) - elif name == "shield_inspect": shield = _get_shield() detections = shield._detector.detect(args["text"]) lines = [] for d in detections: lines.append( - f"- {d.entity_type}: {d.value!r} (pos {d.start}-{d.end})" + f"- {d.entity_type} (pos {d.start}-{d.end})" ) result = "\n".join(lines) if lines else "No PII detected." return _make_response(id, { "content": [{"type": "text", "text": result}], }) - elif name == "shield_vault": - mappings = _session_vault.mappings - if not mappings: - text = "Vault is empty." - else: - lines = [f" {tok} -> {val}" for tok, val in mappings.items()] - text = "Current vault mappings:\n" + "\n".join(lines) - return _make_response(id, { - "content": [{"type": "text", "text": text}], - }) - elif name == "shield_clear": _session_vault.clear() global _session_shield @@ -192,7 +152,7 @@ def _handle_tools_call(id: Any, params: dict) -> dict: return _make_error(id, -32601, f"Unknown tool: {name}") except Exception as e: - return _make_error(id, -32603, str(e)) + return _make_error(id, -32603, "Tool failed; no successful masking result is available.") def _handle_notifications_initialized(params: dict) -> None: @@ -210,38 +170,45 @@ def _handle_notifications_initialized(params: dict) -> None: } +def _dispatch(msg: Any) -> dict | None: + """Validate messages before dispatch; never expose exception data or PII.""" + if not isinstance(msg, dict) or msg.get("jsonrpc") != "2.0": + return _make_error(None, -32600, "Invalid Request") + msg_id = msg.get("id") + method = msg.get("method") + if not isinstance(method, str) or isinstance(msg_id, (dict, list, bool)): + return _make_error(None, -32600, "Invalid Request") + params = msg.get("params", {}) + if not isinstance(params, dict): + return _make_error(msg_id, -32602, "params must be an object") if "id" in msg else None + if "id" not in msg: + handler = _NOTIFICATIONS.get(method) + if handler: + handler(params) + return None + if method == "ping": + return _make_response(msg_id, {}) + handler = _HANDLERS.get(method) + if handler is None: + return _make_error(msg_id, -32601, "Method not found") + try: + return handler(msg_id, params) + except Exception: + return _make_error(msg_id, -32603, "Internal error") + + def main() -> None: - """Run the MCP server over stdio.""" + """Run the MCP helper; raw vault values are never returned as tool output.""" for line in sys.stdin: - line = line.strip() - if not line: + if not line.strip(): continue - try: - msg = json.loads(line) + response = _dispatch(json.loads(line)) except json.JSONDecodeError: - continue - - method = msg.get("method", "") - msg_id = msg.get("id") - params = msg.get("params", {}) - - # Handle notifications (no id, no response) - if msg_id is None: - handler = _NOTIFICATIONS.get(method) - if handler: - handler(params) - continue - - # Handle requests - handler = _HANDLERS.get(method) - if handler: - response = handler(msg_id, params) - else: - response = _make_error(msg_id, -32601, f"Method not found: {method}") - - sys.stdout.write(json.dumps(response) + "\n") - sys.stdout.flush() + response = _make_error(None, -32700, "Parse error") + if response is not None: + sys.stdout.write(json.dumps(response) + "\n") + sys.stdout.flush() if __name__ == "__main__": diff --git a/src/shieldprompt/middleware.py b/src/shieldprompt/middleware.py index 6d1bba4..b86839a 100644 --- a/src/shieldprompt/middleware.py +++ b/src/shieldprompt/middleware.py @@ -2,13 +2,14 @@ from __future__ import annotations +import copy import json import logging from typing import Optional from .entities import ALL_ENTITIES, DEFAULT_ENTITIES, EntityType from .shield import Shield -from .vault import Vault +from .vault import Vault, reset_vault, set_vault logger = logging.getLogger("shieldprompt") @@ -30,9 +31,18 @@ def ShieldPromptMiddleware( exclude_paths: Optional[list[str]] = None, use_ner: Optional[bool] = None, ner_model: str = "dslim/bert-base-NER", + max_request_bytes: int = 1024 * 1024, + max_response_bytes: int = 1024 * 1024, ): """Create and return an ASGI middleware that masks PII in requests - and unmasks it in responses.""" + and unmasks it in finite JSON/text responses. + + Requests larger than max_request_bytes are rejected with 413. Streaming, + compressed, non-text, and oversized responses are forwarded unchanged + (still tokenized). Streams are never buffered or partially restored. + """ + if max_request_bytes <= 0 or max_response_bytes <= 0: + raise ValueError("Body size limits must be positive") try: from starlette.types import ASGIApp, Receive, Scope, Send, Message except ImportError: @@ -41,7 +51,7 @@ def ShieldPromptMiddleware( "Install with: pip install shieldprompt[fastapi]" ) - if entities: + if entities is not None: entity_set = {EntityType(e) for e in entities} else: entity_set = _SENSITIVITY_MAP.get(sensitivity, DEFAULT_ENTITIES) @@ -51,6 +61,9 @@ def ShieldPromptMiddleware( class _ShieldMiddleware: def __init__(self, app: ASGIApp) -> None: self.app = app + self.template = Shield( + entities=entity_set, use_ner=use_ner, ner_model=ner_model, + ) async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: if scope["type"] != "http": @@ -63,106 +76,105 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: return vault = Vault() - shield = Shield( - entities=entity_set, - use_ner=use_ner, - ner_model=ner_model, - vault=vault, - ) + shield = copy.copy(self.template) + shield._vault = vault - # Read the entire request body first body_parts: list[bytes] = [] + size = 0 while True: message = await receive() - body_parts.append(message.get("body", b"")) + if message["type"] == "http.disconnect": + return + body = message.get("body", b"") + size += len(body) + if size > max_request_bytes: + await send({"type": "http.response.start", "status": 413, + "headers": [(b"content-length", b"0")]}) + await send({"type": "http.response.body", "body": b""}) + return + body_parts.append(body) if not message.get("more_body", False): break - full_body = b"".join(body_parts) - - # Try to mask JSON body masked_body = full_body try: - data = json.loads(full_body) - masked_data = _mask_json(shield, data) - masked_body = json.dumps(masked_data).encode() + masked_body = json.dumps(_mask_json(shield, json.loads(full_body))).encode() except (json.JSONDecodeError, UnicodeDecodeError): pass - # Update content-length in scope headers - new_headers = [] - for k, v in scope.get("headers", []): - if k == b"content-length": - new_headers.append((k, str(len(masked_body)).encode())) - else: - new_headers.append((k, v)) - scope["headers"] = new_headers - - # Provide masked body via a new receive + scope = dict(scope) + scope["headers"] = [ + (k, str(len(masked_body)).encode() if k.lower() == b"content-length" else v) + for k, v in scope.get("headers", []) + ] body_sent = False async def masked_receive() -> Message: nonlocal body_sent if not body_sent: body_sent = True - return { - "type": "http.request", - "body": masked_body, - "more_body": False, - } - # After body is consumed, return disconnect - return {"type": "http.disconnect"} - - # Buffer response body for unmasking - response_started = False + return {"type": "http.request", "body": masked_body, "more_body": False} + # A synthetic disconnect cancels Starlette streaming responses. + return await receive() + start_message = None - response_parts: list[bytes] = [] + passthrough = False async def masked_send(message: Message) -> None: - nonlocal response_started, start_message - + nonlocal start_message, passthrough if message["type"] == "http.response.start": - response_started = True start_message = message return - - if message["type"] == "http.response.body": + if message["type"] == "http.response.body" and not passthrough: body = message.get("body", b"") - more = message.get("more_body", False) - response_parts.append(body) - - if not more: - full_resp = b"".join(response_parts) + headers = dict(start_message.get("headers", [])) if start_message else {} + content_type = headers.get(b"content-type", b"").split(b";", 1)[0].lower() + is_json = content_type == b"application/json" or content_type.endswith(b"+json") + is_text = content_type.startswith(b"text/") + passthrough = (message.get("more_body", False) + or len(body) > max_response_bytes + or b"content-encoding" in headers + or not (is_json or is_text)) + changed = False + if not passthrough: try: - text = full_resp.decode() - unmasked = shield.unmask(text) - full_resp = unmasked.encode() - except (UnicodeDecodeError, Exception): + if is_json: + data = json.loads(body) + restored = _unmask_json(shield, data) + if restored != data: + body = json.dumps(restored).encode() + changed = True + else: + restored = shield.unmask(body.decode()).encode() + changed = restored != body + body = restored + except (UnicodeDecodeError, json.JSONDecodeError): pass - - # Send start with updated content-length - if start_message: - sm = dict(start_message) - new_h = [] - for k, v in sm.get("headers", []): - if k == b"content-length": - new_h.append( - (k, str(len(full_resp)).encode()) - ) - else: - new_h.append((k, v)) - sm["headers"] = new_h - await send(sm) - - msg = dict(message) - msg["body"] = full_resp - msg["more_body"] = False - await send(msg) + if start_message is not None: + sm = dict(start_message) + if changed: + # Body-dependent validators no longer describe this body. + sm["headers"] = [ + (k, str(len(body)).encode() if k.lower() == b"content-length" else v) + for k, v in sm.get("headers", []) + if k.lower() not in (b"etag", b"content-md5", b"digest") + ] + await send(sm) + start_message = None + await send(dict(message, body=body)) return - + # Preserve ASGI extensions (e.g. pathsend) without delaying start. + if start_message is not None: + await send(start_message) + start_message = None + passthrough = True await send(message) - await self.app(scope, masked_receive, masked_send) + token = set_vault(vault) + try: + await self.app(scope, masked_receive, masked_send) + finally: + reset_vault(token) return _ShieldMiddleware(app) @@ -176,3 +188,15 @@ def _mask_json(shield: Shield, data): if isinstance(data, list): return [_mask_json(shield, item) for item in data] return data + + + +def _unmask_json(shield: Shield, data): + """Restore JSON values before serialization, preserving escaping and types.""" + if isinstance(data, str): + return shield.unmask(data) + if isinstance(data, dict): + return {k: _unmask_json(shield, v) for k, v in data.items()} + if isinstance(data, list): + return [_unmask_json(shield, item) for item in data] + return data diff --git a/src/shieldprompt/shield.py b/src/shieldprompt/shield.py index a79589f..62339cd 100644 --- a/src/shieldprompt/shield.py +++ b/src/shieldprompt/shield.py @@ -6,7 +6,7 @@ from typing import Optional from .detector.hybrid import HybridDetector -from .detector.regex_engine import RegexDetector +from .detector.regex_engine import RegexDetector, merge_overlaps from .entities import DEFAULT_ENTITIES, NER_ENTITIES, EntityType from .vault import Vault, get_vault @@ -30,7 +30,7 @@ class Shield: Args: entities: Set of entity types to detect. Defaults to common PII types. - use_ner: Whether to use the NER model. Auto-detected if None. + use_ner: Whether to use the NER model. Enabled by default; failures raise instead of exposing text. ner_model: HuggingFace model name for NER. ner_threshold: Minimum confidence for NER detections. vault: Optional vault instance. Uses context-local vault if None. @@ -46,13 +46,13 @@ def __init__( ) -> None: self._entities = ( {EntityType(e) if isinstance(e, str) else e for e in entities} - if entities - else DEFAULT_ENTITIES + if entities is not None + else set(DEFAULT_ENTITIES) ) self._vault = vault - # Use HybridDetector whenever any NER-category entity is requested, so - # the name-dictionary recognizer runs even if transformer NER is off. + # Use HybridDetector whenever a contextual entity is requested. + # Explicit use_ner=False opts into reduced regex-only protection. needs_ner_category = bool(self._entities & NER_ENTITIES) if not needs_ner_category: self._detector = RegexDetector(entities=self._entities) @@ -71,7 +71,7 @@ def vault(self) -> Vault: def mask(self, text: str) -> str: """Detect PII and replace with tokens. Returns masked text.""" text = _neutralize_user_tokens(text) - detections = self._detector.detect(text) + detections = merge_overlaps(text, self._detector.detect(text)) if not detections: return text @@ -97,3 +97,4 @@ def mask_and_track(self, text: str) -> tuple[str, dict[str, str]]: def clear(self) -> None: """Clear the vault mappings.""" self.vault.clear() + diff --git a/src/shieldprompt/vault.py b/src/shieldprompt/vault.py index b3e205d..e1d76b9 100644 --- a/src/shieldprompt/vault.py +++ b/src/shieldprompt/vault.py @@ -3,6 +3,7 @@ from __future__ import annotations import contextvars +import re import threading from typing import Optional @@ -47,13 +48,14 @@ def restore(self, token: str) -> Optional[str]: def restore_text(self, text: str) -> str: """Replace all tokens in text with their real values.""" - result = text - # Sort by token length descending to avoid partial replacement issues - for token, real in sorted( - self._token_to_real.items(), key=lambda x: len(x[0]), reverse=True - ): - result = result.replace(token, real) - return result + with self._lock: + mappings = dict(self._token_to_real) + if not mappings: + return text + pattern = re.compile("|".join(re.escape(token) for token in + sorted(mappings, key=len, reverse=True))) + # A single pass never interprets token-like text inside original values. + return pattern.sub(lambda match: mappings[match.group(0)], text) @property def mappings(self) -> dict[str, str]: @@ -67,9 +69,18 @@ def to_dict(self) -> dict[str, str]: @classmethod def from_dict(cls, mappings: dict[str, str]) -> "Vault": """Rebuild a Vault from a token->real mapping dict.""" + if not isinstance(mappings, dict): + raise ValueError("Vault mappings must be a dictionary") vault = cls() with vault._lock: for token, real in mappings.items(): + if not isinstance(token, str) or not isinstance(real, str): + raise ValueError("Vault tokens and values must be strings") + match = re.fullmatch(r"\[([A-Z][A-Z_]*?)_([1-9][0-9]*)\]", token) + if not match: + raise ValueError("Invalid vault token format") + entity, number = match.groups() + vault._counters[entity] = max(vault._counters.get(entity, 0), int(number)) vault._token_to_real[token] = real vault._real_to_token[real] = token return vault @@ -99,3 +110,4 @@ def set_vault(vault: Vault) -> contextvars.Token: def reset_vault(token: contextvars.Token) -> None: """Reset the vault to its previous state.""" _vault_var.reset(token) + diff --git a/tests/test_cli_security.py b/tests/test_cli_security.py new file mode 100644 index 0000000..66abbd5 --- /dev/null +++ b/tests/test_cli_security.py @@ -0,0 +1,156 @@ +"""Regression coverage for recovery-preserving CLI writes.""" +import json +import os +import stat + +import pytest + +from shieldprompt import cli +from shieldprompt.vault import Vault + + +def mask(source, *extra): + cli.main(["mask", "--file", str(source), "--in-place", "--no-ner", "--entities", "EMAIL_ADDRESS", *extra]) + + +def test_vault_failure_leaves_source_untouched(tmp_path, monkeypatch): + source = tmp_path / "source" + source.write_text("fixture@example.test") + def fail(*args): + raise OSError("vault write failed") + monkeypatch.setattr(cli, "_save_vault", fail) + with pytest.raises(SystemExit): + mask(source) + assert source.read_text() == "fixture@example.test" + + +def test_source_failure_keeps_recovery_vault(tmp_path, monkeypatch): + source = tmp_path / "source" + source.write_text("fixture@example.test") + def fail(*args): + raise OSError("source replacement failed") + monkeypatch.setattr(cli.os, "replace", fail) + with pytest.raises(SystemExit): + mask(source) + assert source.read_text() == "fixture@example.test" + vault = json.loads((tmp_path / "source.shieldprompt.vault.json").read_text()) + assert "fixture@example.test" in vault.values() + assert not list(tmp_path.glob(".shieldprompt-*")) + + +def test_repeat_mask_preserves_both_files(tmp_path): + source = tmp_path / "source" + source.write_text("fixture@example.test") + mask(source) + sidecar = tmp_path / "source.shieldprompt.vault.json" + before = (source.read_bytes(), sidecar.read_bytes()) + with pytest.raises(SystemExit): + mask(source) + assert (source.read_bytes(), sidecar.read_bytes()) == before + with pytest.raises(SystemExit): + mask(source, "--save-vault", str(tmp_path / "new-vault")) + assert source.read_bytes() == before[0] + cli.main(["unmask", "--file", str(source), "--in-place"]) + assert source.read_text() == "fixture@example.test" + + +@pytest.mark.parametrize("alias", ["same", "symlink", "hardlink"]) +def test_source_vault_alias_rejected(tmp_path, alias): + source = tmp_path / "source" + source.write_text("fixture@example.test") + vault = source if alias == "same" else tmp_path / "vault" + if alias == "symlink": + vault.symlink_to(source) + elif alias == "hardlink": + os.link(source, vault) + with pytest.raises(SystemExit): + mask(source, "--save-vault", str(vault)) + assert source.read_text() == "fixture@example.test" + + +def test_vault_cannot_overwrite_existing_or_symlink(tmp_path): + existing = tmp_path / "existing" + existing.write_text("keep") + alias = tmp_path / "alias" + alias.symlink_to(existing) + for path in (existing, alias): + with pytest.raises(FileExistsError): + cli._save_vault(Vault(), str(path)) + assert existing.read_text() == "keep" + assert not list(tmp_path.glob(".shieldprompt-*")) + + +@pytest.mark.skipif(os.name != "posix", reason="POSIX permission modes") +def test_private_permissions_at_publication_and_after_restore(tmp_path, monkeypatch): + source = tmp_path / "source" + source.write_text("fixture@example.test") + source.chmod(0o644) + real_link, real_replace = os.link, os.replace + def check_link(src, dst): + assert stat.S_IMODE(os.stat(src).st_mode) == 0o600 + return real_link(src, dst) + def check_replace(src, dst): + assert stat.S_IMODE(os.stat(src).st_mode) == 0o600 + return real_replace(src, dst) + monkeypatch.setattr(cli.os, "link", check_link) + monkeypatch.setattr(cli.os, "replace", check_replace) + mask(source) + cli.main(["unmask", "--file", str(source), "--in-place"]) + assert stat.S_IMODE(source.stat().st_mode) == 0o600 + + +def test_unmask_rejects_source_equal_to_vault(tmp_path): + source = tmp_path / "vault" + original = '{"[PERSON_1]": "Fixture"}' + source.write_text(original) + with pytest.raises(SystemExit): + cli.main(["unmask", "--file", str(source), "--vault", str(source), "--in-place"]) + assert source.read_text() == original + + +@pytest.mark.parametrize("stage", ["file", "directory"]) +def test_vault_sync_failure_never_replaces_source(tmp_path, monkeypatch, stage): + source = tmp_path / "source" + source.write_text("fixture@example.test") + def fail(*args): + raise OSError("sync failed") + if stage == "file": + monkeypatch.setattr(cli.os, "fsync", fail) + else: + monkeypatch.setattr(cli, "_sync_directory", fail) + with pytest.raises(SystemExit): + mask(source) + assert source.read_text() == "fixture@example.test" + sidecar = tmp_path / "source.shieldprompt.vault.json" + if stage == "directory": + assert "fixture@example.test" in json.loads(sidecar.read_text()).values() + else: + assert not sidecar.exists() + assert not list(tmp_path.glob(".shieldprompt-*")) + + +def test_vault_created_during_mask_cannot_be_overwritten(tmp_path, monkeypatch): + source = tmp_path / "source" + source.write_text("fixture@example.test") + sidecar = tmp_path / "source.shieldprompt.vault.json" + real_link = os.link + def concurrent_link(src, dst): + sidecar.write_text("concurrent recovery vault") + return real_link(src, dst) + monkeypatch.setattr(cli.os, "link", concurrent_link) + with pytest.raises(SystemExit): + mask(source) + assert sidecar.read_text() == "concurrent recovery vault" + assert source.read_text() == "fixture@example.test" + assert not list(tmp_path.glob(".shieldprompt-*")) + + +def test_symlink_source_is_not_replaced(tmp_path): + source = tmp_path / "source" + source.write_text("fixture@example.test") + alias = tmp_path / "alias" + alias.symlink_to(source) + with pytest.raises(SystemExit): + mask(alias) + assert alias.is_symlink() + assert source.read_text() == "fixture@example.test" diff --git a/tests/test_detection_security.py b/tests/test_detection_security.py new file mode 100644 index 0000000..e321229 --- /dev/null +++ b/tests/test_detection_security.py @@ -0,0 +1,175 @@ +"""Security regression tests use synthetic data and offline fake NER backends.""" +import pytest + +from shieldprompt import Shield, Vault +from shieldprompt.detector.hybrid import HybridDetector +from shieldprompt.detector.regex_engine import Detection, RegexDetector, merge_overlaps +from shieldprompt.detector.ner_engine import NERDetector +from shieldprompt.detector.presidio_engine import PresidioDetector + + +def make_shield(entities): + return Shield(entities=set(entities), use_ner=False, vault=Vault()) + + +@pytest.mark.parametrize('text', [ + 'before https://192.0.2.1/path after', + 'before https://test@example.com/192.0.2.1 after', +]) +def test_overlapping_url_spans_roundtrip(text): + shield = make_shield({'URL', 'IP_ADDRESS', 'EMAIL_ADDRESS'}) + masked = shield.mask(text) + assert masked == 'before [URL_1] after' + assert shield.unmask(masked) == text + + +def test_partial_overlap_union_preserves_all_sensitive_characters(): + merged = merge_overlaps('abcdefgh', [Detection('URL', 'abcd', 0, 4), + Detection('PERSON', 'cdef', 2, 6)]) + assert len(merged) == 1 + assert (merged[0].start, merged[0].end, merged[0].value) == (0, 6, 'abcdef') + + +@pytest.mark.parametrize('key', [ + 'sk-proj-' + 'aB1_' * 20 + '-end', + 'sk-svcacct-' + 'aB1_' * 20, + 'github_pat_' + 'abc123_' * 12, + 'ghp_' + 'a' * 36, +]) +def test_modern_synthetic_credentials(key): + shield = make_shield({'API_KEY'}) + masked = shield.mask('key=' + key) + assert masked == 'key=[API_KEY_1]' + assert shield.unmask(masked) == 'key=' + key + + +@pytest.mark.parametrize('text, value', [ + ('ligature ffi: alice@example.com', 'alice@example.com'), + ('contact ali\u200bce@exam\u2060ple.com', 'ali\u200bce@exam\u2060ple.com'), + ('contact ａｌｉｃｅ＠ｅｘａｍｐｌｅ．ｃｏｍ', 'ａｌｉｃｅ＠ｅｘａｍｐｌｅ．ｃｏｍ'), + ('ffi contact alice@example.com\u200b\u200b', 'alice@example.com'), +]) +def test_unicode_normalization_original_offsets(text, value): + shield = make_shield({'EMAIL_ADDRESS'}) + masked = shield.mask(text) + assert value not in masked + assert shield.unmask(masked) == text + assert value in shield.vault.mappings.values() + + +def test_normalized_card_validation(): + shield = make_shield({'CREDIT_CARD'}) + text = '４１１１\u200b１１１１\u200b１１１１\u200b１１１１' + assert shield.mask(text) == '[CREDIT_CARD_1]' + assert shield.unmask('[CREDIT_CARD_1]') == text + + +def test_empty_entity_selection_detects_nothing(): + text = 'alice@example.com 192.0.2.1' + assert RegexDetector(set()).detect(text) == [] + assert HybridDetector(set()).detect(text) == [] + assert make_shield(set()).mask(text) == text + + +def test_ner_only_selection_does_not_enable_regex(): + assert HybridDetector({'PERSON'}, use_ner=False).detect('alice@example.com') == [] + + +def test_requested_ner_initialization_fails_closed(monkeypatch): + attempts = [] + def fail(self): + attempts.append(type(self).__name__) + raise ImportError('offline synthetic missing dependency') + monkeypatch.setattr(PresidioDetector, '_load', fail) + monkeypatch.setattr(NERDetector, '_load', fail) + with pytest.raises(RuntimeError, match='no usable backend'): + Shield(entities={'PERSON'}) + assert attempts == ['PresidioDetector', 'NERDetector'] + assert make_shield({'PERSON'}).mask('Synthetic Person') == 'Synthetic Person' + + +def test_ner_runtime_failure_fails_closed(monkeypatch): + class BrokenNER: + def detect(self, text): + raise RuntimeError('synthetic inference failure') + monkeypatch.setattr(HybridDetector, '_init_ner_backend', staticmethod(lambda **kw: BrokenNER())) + shield = Shield(entities={'PERSON', 'EMAIL_ADDRESS'}, vault=Vault()) + with pytest.raises(RuntimeError, match='refusing'): + shield.mask('Synthetic Person alice@example.com') + assert shield.vault.mappings == {} + + +@pytest.mark.parametrize('entity,text', [('DATE_OF_BIRTH', '2000-01-02'), ('MONEY', '$123.45')]) +def test_previously_unimplemented_entities(entity, text): + shield = make_shield({entity}) + masked = shield.mask(text) + assert text not in masked + assert shield.unmask(masked) == text + + +def test_invalid_detector_span_fails_closed(): + with pytest.raises(ValueError, match='invalid source span'): + merge_overlaps('abc', [Detection('PERSON', 'secret', 0, 10)]) + + +def test_transformer_rejects_unsupported_entity_configuration(monkeypatch): + import sys + from types import SimpleNamespace + fake_pipeline = SimpleNamespace(model=SimpleNamespace( + config=SimpleNamespace(id2label={0: 'O', 1: 'B-PER', 2: 'I-PER'}))) + monkeypatch.setitem(sys.modules, 'transformers', SimpleNamespace( + pipeline=lambda *args, **kwargs: fake_pipeline)) + detector = NERDetector(entities={'DATE'}) + with pytest.raises(RuntimeError, match='does not support'): + detector._load() + assert detector._pipeline is None + + +def test_missing_spacy_model_never_triggers_download(monkeypatch): + import sys + from types import SimpleNamespace + def missing_model(name): + raise OSError('synthetic missing model') + monkeypatch.setitem(sys.modules, 'spacy', SimpleNamespace(load=missing_model)) + with pytest.raises(RuntimeError, match='Install it explicitly'): + PresidioDetector()._ensure_spacy_model() + + +def test_presidio_explicit_configuration_keeps_default_organization(monkeypatch): + """Exercise initialization without inheriting version-dependent ORG defaults.""" + import sys + from types import SimpleNamespace + configuration = {} + + class Provider: + def __init__(self, nlp_configuration): + configuration.update(nlp_configuration) + + def create_engine(self): + return configuration['ner_model_configuration'] + + class Analyzer: + def __init__(self, nlp_engine, default_score_threshold): + self.config = nlp_engine + self.threshold = default_score_threshold + + def get_supported_entities(self, language): + return set(self.config['model_to_presidio_entity_mapping'].values()) - set( + self.config['labels_to_ignore']) + + def analyze(self, text, language): + label = self.config['model_to_presidio_entity_mapping']['ORG'] + score = self.config['default_score'] + assert label not in self.config['low_score_entity_names'] + assert score >= self.threshold + return [SimpleNamespace(entity_type=label, start=0, end=len(text), score=score)] + + monkeypatch.setitem(sys.modules, 'presidio_analyzer', SimpleNamespace(AnalyzerEngine=Analyzer)) + monkeypatch.setitem(sys.modules, 'presidio_analyzer.nlp_engine', SimpleNamespace(NlpEngineProvider=Provider)) + monkeypatch.setattr(PresidioDetector, '_ensure_spacy_model', lambda self: None) + shield = Shield(vault=Vault()) + assert isinstance(shield._detector._ner, PresidioDetector) + text = 'Synthetic Example Company' + masked = shield.mask(text) + assert masked == '[ORGANIZATION_1]' + assert shield.unmask(masked) == text diff --git a/tests/test_integration_security.py b/tests/test_integration_security.py new file mode 100644 index 0000000..a30173d --- /dev/null +++ b/tests/test_integration_security.py @@ -0,0 +1,173 @@ +"""Regression coverage for integration boundaries using synthetic PII only.""" +import asyncio +import json + +import pytest + +from shieldprompt.decorator import mask_pii +from shieldprompt.middleware import ShieldPromptMiddleware +from shieldprompt.vault import get_vault + +EMAIL = "synthetic@example.com" + + +def test_variadic_keywords_deduplicated_and_nested_result(): + @mask_pii(entities=["EMAIL_ADDRESS"], use_ner=False, + mask_kwargs=["prompt", "context", "context"]) + def echo(**kwargs): + assert EMAIL not in repr(kwargs) + assert "\u200b" not in repr(kwargs) + return {"nested": [kwargs["prompt"], {"context": kwargs["context"]}]} + + assert echo(prompt=EMAIL, context=EMAIL) == {"nested": [EMAIL, {"context": EMAIL}]} + + +def test_variadic_positional_fallback(): + @mask_pii(entities=["EMAIL_ADDRESS"], use_ner=False) + def echo(*args): + assert EMAIL not in repr(args) + return args + + assert echo(123, EMAIL) == (123, EMAIL) + + +@pytest.mark.asyncio +async def test_async_decorator_isolates_concurrent_vaults(): + @mask_pii(entities=["EMAIL_ADDRESS"], use_ner=False) + async def echo(prompt): + await asyncio.sleep(0) + assert len(get_vault().mappings) == 1 + return [prompt] + + values = ["synthetic1@example.com", "synthetic2@example.com"] + assert await asyncio.gather(*(echo(v) for v in values)) == [[v] for v in values] + + +async def invoke(app, body=b"{}", extra_messages=None): + received = [{"type": "http.request", "body": body, "more_body": False}] + received += extra_messages or [] + sent = [] + + async def receive(): + if received: + return received.pop(0) + await asyncio.Future() + + async def send(message): + sent.append(message) + + await app({"type": "http", "path": "/", "headers": []}, receive, send) + return sent + + +@pytest.mark.asyncio +async def test_json_restoration_escapes_values_and_preserves_types(): + async def app(scope, receive, send): + await receive() + # Synthetic value deliberately exercises JSON escaping, independent of detectors. + token = get_vault().store("PHONE_NUMBER", '555\n123\t"4567') + body = json.dumps({"nested": [token, 3, None]}).encode() + await send({"type": "http.response.start", "status": 200, + "headers": [(b"content-type", b"application/json"), + (b"content-length", str(len(body)).encode()), (b"etag", b"old")]}) + await send({"type": "http.response.body", "body": body}) + + sent = await invoke(ShieldPromptMiddleware(app, entities=["EMAIL_ADDRESS"], use_ner=False)) + assert json.loads(sent[1]["body"]) == {"nested": ['555\n123\t"4567', 3, None]} + assert dict(sent[0]["headers"])[b"content-length"] == str(len(sent[1]["body"])).encode() + assert b"etag" not in dict(sent[0]["headers"]) + + +@pytest.mark.asyncio +async def test_streams_forward_immediately_without_fake_disconnect(): + sent = [] + messages = iter([{"type": "http.request", "body": b"{}"}, + {"type": "http.request", "body": b"", "more_body": False}]) + + async def receive(): + return next(messages) + + async def send(message): + sent.append(message) + + async def app(scope, receive, send): + await receive() + assert (await receive())["type"] == "http.request" + token = get_vault().store("EMAIL_ADDRESS", EMAIL) + await send({"type": "http.response.start", "status": 200, + "headers": [(b"content-type", b"text/event-stream")]}) + await send({"type": "http.response.body", "body": token.encode(), "more_body": True}) + assert len(sent) == 2 # no waiting for stream completion + await send({"type": "http.response.body", "body": b"done", "more_body": False}) + + await ShieldPromptMiddleware(app, entities=["EMAIL_ADDRESS"], use_ner=False)( + {"type": "http", "path": "/", "headers": []}, receive, send) + assert sent[1]["body"] == b"[EMAIL_ADDRESS_1]" + assert sent[2]["body"] == b"done" + + +@pytest.mark.asyncio +async def test_request_limit_rejects_before_app_runs(): + async def app(*args): + pytest.fail("Oversized request must not reach app") + + wrapped = ShieldPromptMiddleware(app, entities=["EMAIL_ADDRESS"], use_ner=False, max_request_bytes=4) + sent = await invoke(wrapped, b"12345") + assert sent[0]["status"] == 413 + + +@pytest.mark.asyncio +async def test_response_limit_preserves_tokenized_body(): + async def app(scope, receive, send): + await receive() + token = get_vault().store("EMAIL_ADDRESS", EMAIL) + await send({"type": "http.response.start", "status": 200, + "headers": [(b"content-type", b"text/plain")]}) + await send({"type": "http.response.body", "body": token.encode()}) + + wrapped = ShieldPromptMiddleware(app, entities=["EMAIL_ADDRESS"], use_ner=False, max_response_bytes=4) + sent = await invoke(wrapped) + assert sent[1]["body"] == b"[EMAIL_ADDRESS_1]" + + +@pytest.mark.asyncio +async def test_middleware_reuses_detector_and_isolates_concurrent_requests(): + detectors = [] + + async def app(scope, receive, send): + data = json.loads((await receive())["body"]) + assert "@" not in data["prompt"] + await asyncio.sleep(0) + assert len(get_vault().mappings) == 1 + await send({"type": "http.response.start", "status": 200, + "headers": [(b"content-type", b"application/json")]}) + await send({"type": "http.response.body", "body": json.dumps(data).encode()}) + + wrapped = ShieldPromptMiddleware(app, entities=["EMAIL_ADDRESS"], use_ner=False) + original_detect = wrapped.template._detector.detect + + def detect(text): + detectors.append(wrapped.template._detector) + return original_detect(text) + + wrapped.template._detector.detect = detect + values = ["synthetic1@example.com", "synthetic2@example.com"] + results = await asyncio.gather(*(invoke(wrapped, json.dumps({"prompt": v}).encode()) for v in values)) + assert [json.loads(r[1]["body"])["prompt"] for r in results] == values + assert len(detectors) == 2 + assert detectors[0] is detectors[1] + + +@pytest.mark.asyncio +async def test_request_disconnect_does_not_call_app(): + async def app(*args): + pytest.fail("Disconnected request must not reach app") + + async def receive(): + return {"type": "http.disconnect"} + + async def send(message): + pytest.fail("No response should be sent after disconnect") + + wrapped = ShieldPromptMiddleware(app, entities=["EMAIL_ADDRESS"], use_ner=False) + await wrapped({"type": "http", "path": "/", "headers": []}, receive, send) diff --git a/tests/test_mcp_security.py b/tests/test_mcp_security.py new file mode 100644 index 0000000..e48b203 --- /dev/null +++ b/tests/test_mcp_security.py @@ -0,0 +1,46 @@ +"""Protocol and model-visible privacy boundary regressions.""" +import io +import json + +from shieldprompt import Shield, Vault, __version__ +from shieldprompt import mcp_server as server + + +def test_no_secret_returning_tools(monkeypatch): + vault = Vault() + vault.store('EMAIL_ADDRESS', 'alice@example.com') + monkeypatch.setattr(server, '_session_vault', vault) + names = {t['name'] for t in server._handle_tools_list(1, {})['result']['tools']} + assert names == {'shield_mask', 'shield_inspect', 'shield_clear'} + for name in ('shield_unmask', 'shield_vault'): + response = server._handle_tools_call(1, {'name': name, 'arguments': {'text': '[EMAIL_ADDRESS_1]'}}) + assert 'error' in response + assert 'alice@example.com' not in json.dumps(response) + + +def test_inspection_does_not_return_originals(monkeypatch): + monkeypatch.setattr(server, '_session_shield', Shield(entities={'EMAIL_ADDRESS'}, use_ner=False, vault=Vault())) + result = server._handle_tools_call(1, {'name': 'shield_inspect', 'arguments': {'text': 'alice@example.com'}}) + assert 'EMAIL_ADDRESS' in json.dumps(result) + assert 'alice@example.com' not in json.dumps(result) + + +def test_dispatch_survives_bad_shapes(): + for request in ([1], None, {'jsonrpc': '2.0', 'id': 1, 'method': 'tools/call', 'params': None}): + assert 'error' in server._dispatch(request) + for args in (None, [], {'text': 123}): + response = server._dispatch({'jsonrpc': '2.0', 'id': 1, 'method': 'tools/call', 'params': {'name': 'shield_mask', 'arguments': args}}) + assert response['error']['code'] == -32602 + assert server._dispatch({'jsonrpc': '2.0', 'id': 2, 'method': 'ping'})['result'] == {} + assert server._handle_initialize(1, {})['result']['serverInfo']['version'] == __version__ + + +def test_parse_errors_do_not_stop_stdio(monkeypatch): + monkeypatch.setattr(server.sys, 'stdin', io.StringIO('invalid\n[1]\n{"jsonrpc":"2.0","id":3,"method":"ping"}\n')) + output = io.StringIO() + monkeypatch.setattr(server.sys, 'stdout', output) + server.main() + results = [json.loads(line) for line in output.getvalue().splitlines()] + assert results[0]['error']['code'] == -32700 + assert results[1]['error']['code'] == -32600 + assert results[2] == {'jsonrpc': '2.0', 'id': 3, 'result': {}} diff --git a/tests/test_vault_security.py b/tests/test_vault_security.py new file mode 100644 index 0000000..c7696d9 --- /dev/null +++ b/tests/test_vault_security.py @@ -0,0 +1,22 @@ +"""Recovery mappings must remain stable across loading and restoration.""" +import pytest + +from shieldprompt.vault import Vault + + +def test_loaded_counters_never_overwrite_existing_tokens(): + vault = Vault.from_dict({"[EMAIL_ADDRESS_7]": "old@example.test", "[PERSON_2]": "Example"}) + assert vault.store("EMAIL_ADDRESS", "new@example.test") == "[EMAIL_ADDRESS_8]" + assert vault.store("PERSON", "Other") == "[PERSON_3]" + assert vault.restore("[EMAIL_ADDRESS_7]") == "old@example.test" + + +def test_restore_is_single_pass_even_when_original_contains_another_token(): + vault = Vault.from_dict({"[EMAIL_ADDRESS_1]": "literal [PERSON_1]", "[PERSON_1]": "Example"}) + assert vault.restore_text("[EMAIL_ADDRESS_1] [PERSON_1]") == "literal [PERSON_1] Example" + + +@pytest.mark.parametrize("mappings", [[], {"": "value"}, {"[PERSON_0]": "value"}, {"[PERSON_1]": 42}, {42: "value"}]) +def test_invalid_vaults_fail_closed(mappings): + with pytest.raises(ValueError): + Vault.from_dict(mappings) From 0b4827adc2d3c2dd8b2090bbd2706e0e15cbe4bb Mon Sep 17 00:00:00 2001 From: Chandan Kumar Date: Fri, 6 Nov 2026 18:52:14 +0530 Subject: [PATCH 2/2] ci: pin Ubuntu runner for Python 3.9 compatibility --- .github/workflows/tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 28f92a1..d56260c 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -10,7 +10,7 @@ permissions: jobs: pytest: - runs-on: ubuntu-latest + runs-on: ubuntu-22.04 strategy: fail-fast: false matrix: