diff --git a/app/core/config.py b/app/core/config.py index 5fdd318..6f1c621 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -38,6 +38,15 @@ class Settings(BaseSettings): placeholder_email: str = "" placeholder_phone: str = "" placeholder_card: str = "" + placeholder_sensitive_word: str = "" + + # Sensitive word detection (comma-separated list) + sensitive_words: str = "" + + @property + def sensitive_word_list(self) -> list[str]: + """Parse comma-separated sensitive words.""" + return [w.strip() for w in self.sensitive_words.split(",") if w.strip()] # ======================================== # Safe-to-LLM Proxy Settings diff --git a/app/models/schemas.py b/app/models/schemas.py index 4508d12..78204b2 100644 --- a/app/models/schemas.py +++ b/app/models/schemas.py @@ -7,7 +7,7 @@ from app.core.config import settings # Entity types -EntityType = Literal["EMAIL", "PHONE", "CARD", "PERSON"] +EntityType = Literal["EMAIL", "PHONE", "CARD", "PERSON", "SENSITIVE_WORD"] class TextRequest(BaseModel): diff --git a/app/services/pii_detector.py b/app/services/pii_detector.py index afa9dbe..a879658 100644 --- a/app/services/pii_detector.py +++ b/app/services/pii_detector.py @@ -32,6 +32,7 @@ class PIIDetector: - PHONE: Phone numbers (international formats) - CARD: Credit/debit card numbers - PERSON: Person names (via spaCy NER) + - SENSITIVE_WORD: User-defined sensitive words (via MASKER_SENSITIVE_WORDS env var) """ # Regex patterns for PII detection @@ -154,6 +155,24 @@ def _detect_by_ner(self, text: str, language: str) -> list[DetectedEntity]: return entities + def _detect_sensitive_words(self, text: str) -> list[DetectedEntity]: + """Detect configured sensitive words using word boundary regex.""" + from app.core.config import settings + + entities = [] + for word in settings.sensitive_word_list: + pattern = re.compile(rf"\b{re.escape(word)}\b", re.IGNORECASE) + for match in pattern.finditer(text): + entities.append( + DetectedEntity( + type="SENSITIVE_WORD", + value=match.group(), + start=match.start(), + end=match.end(), + ) + ) + return entities + def _remove_overlaps(self, entities: list[DetectedEntity]) -> list[DetectedEntity]: """Remove overlapping entities, preferring regex matches. @@ -171,7 +190,7 @@ def _remove_overlaps(self, entities: list[DetectedEntity]) -> list[DetectedEntit # Sort by start position, then by priority (more specific types first) # CARD has higher priority than PHONE to avoid card numbers being detected as phones - priority = {"EMAIL": 0, "CARD": 1, "PHONE": 2, "PERSON": 3} + priority = {"EMAIL": 0, "CARD": 1, "PHONE": 2, "PERSON": 3, "SENSITIVE_WORD": 4} sorted_entities = sorted(entities, key=lambda e: (e.start, priority.get(e.type, 99))) result = [] @@ -207,8 +226,11 @@ def detect( # Then, detect using NER ner_entities = self._detect_by_ner(text, language) + # Detect sensitive words + sensitive_entities = self._detect_sensitive_words(text) + # Combine and remove overlaps - all_entities = regex_entities + ner_entities + all_entities = regex_entities + ner_entities + sensitive_entities unique_entities = self._remove_overlaps(all_entities) # Filter by entity types if specified diff --git a/app/services/redaction.py b/app/services/redaction.py index 0024a7a..069fdb6 100644 --- a/app/services/redaction.py +++ b/app/services/redaction.py @@ -26,6 +26,7 @@ class RedactedEntity: "EMAIL": "", "PHONE": "", "CARD": "", + "SENSITIVE_WORD": "", } # Default mask token diff --git a/policies/default.yaml b/policies/default.yaml index 3f366de..5477fcd 100644 --- a/policies/default.yaml +++ b/policies/default.yaml @@ -14,6 +14,7 @@ categories: phone: mask card: drop person: placeholder + sensitive_word: mask # Behavior when redaction fails: # - closed: block the request (do not forward to upstream) diff --git a/tests/test_sensitive_words.py b/tests/test_sensitive_words.py new file mode 100644 index 0000000..378326b --- /dev/null +++ b/tests/test_sensitive_words.py @@ -0,0 +1,127 @@ +"""Tests for sensitive word detection.""" + +from fastapi.testclient import TestClient + +from app.core.config import settings + + +class TestSensitiveWords: + """Tests for SENSITIVE_WORD entity detection and masking.""" + + def test_detect_sensitive_word(self, client: TestClient, monkeypatch): + """Should detect configured sensitive words.""" + monkeypatch.setattr(settings, "sensitive_words", "secret,confidential") + response = client.post("/api/v1/detect", json={"text": "This is a secret document"}) + + assert response.status_code == 200 + data = response.json() + sw = [e for e in data["entities"] if e["type"] == "SENSITIVE_WORD"] + assert len(sw) == 1 + assert sw[0]["value"] == "secret" + + def test_detect_multiple_sensitive_words(self, client: TestClient, monkeypatch): + """Should detect multiple different sensitive words.""" + monkeypatch.setattr(settings, "sensitive_words", "secret,confidential") + response = client.post("/api/v1/detect", json={"text": "This is secret and confidential"}) + + assert response.status_code == 200 + data = response.json() + sw = [e for e in data["entities"] if e["type"] == "SENSITIVE_WORD"] + assert len(sw) == 2 + values = {e["value"].lower() for e in sw} + assert "secret" in values + assert "confidential" in values + + def test_case_insensitive_detection(self, client: TestClient, monkeypatch): + """Should detect words regardless of case.""" + monkeypatch.setattr(settings, "sensitive_words", "secret") + response = client.post("/api/v1/detect", json={"text": "This is SECRET and Secret"}) + + assert response.status_code == 200 + data = response.json() + sw = [e for e in data["entities"] if e["type"] == "SENSITIVE_WORD"] + assert len(sw) == 2 + + def test_word_boundary_prevents_substring_match(self, client: TestClient, monkeypatch): + """Should not match sensitive words inside other words.""" + monkeypatch.setattr(settings, "sensitive_words", "secret") + response = client.post("/api/v1/detect", json={"text": "This is a secretory gland"}) + + assert response.status_code == 200 + data = response.json() + sw = [e for e in data["entities"] if e["type"] == "SENSITIVE_WORD"] + assert len(sw) == 0 + + def test_empty_config_no_detection(self, client: TestClient, monkeypatch): + """Should not detect sensitive words when config is empty.""" + monkeypatch.setattr(settings, "sensitive_words", "") + response = client.post("/api/v1/detect", json={"text": "This is a secret document"}) + + assert response.status_code == 200 + data = response.json() + sw = [e for e in data["entities"] if e["type"] == "SENSITIVE_WORD"] + assert len(sw) == 0 + + def test_mask_sensitive_word(self, client: TestClient, monkeypatch): + """Should mask sensitive words with asterisks.""" + monkeypatch.setattr(settings, "sensitive_words", "secret") + response = client.post("/api/v1/mask", json={"text": "This is a secret document"}) + + assert response.status_code == 200 + data = response.json() + assert "***" in data["text"] + assert "secret" not in data["text"] + + def test_redact_sensitive_word(self, client: TestClient, monkeypatch): + """Should redact sensitive words with [REDACTED].""" + monkeypatch.setattr(settings, "sensitive_words", "secret") + response = client.post("/api/v1/redact", json={"text": "This is a secret document"}) + + assert response.status_code == 200 + data = response.json() + assert "[REDACTED]" in data["text"] + assert "secret" not in data["text"] + + def test_overlap_existing_entity_wins(self, client: TestClient, monkeypatch): + """Should prefer EMAIL over SENSITIVE_WORD on overlap.""" + monkeypatch.setattr(settings, "sensitive_words", "test") + response = client.post("/api/v1/detect", json={"text": "Contact test@example.com"}) + + assert response.status_code == 200 + data = response.json() + types = [e["type"] for e in data["entities"]] + assert "EMAIL" in types + assert "SENSITIVE_WORD" not in types + + def test_cyrillic_word_boundary(self, client: TestClient, monkeypatch): + """Should detect Cyrillic sensitive words.""" + monkeypatch.setattr(settings, "sensitive_words", "СВО,спецоперация") + response = client.post("/api/v1/detect", json={"text": "Новости о СВО и спецоперации"}) + + assert response.status_code == 200 + data = response.json() + sw = [e for e in data["entities"] if e["type"] == "SENSITIVE_WORD"] + assert len(sw) >= 1 + + def test_sensitive_word_at_text_boundaries(self, client: TestClient, monkeypatch): + """Should detect words at start and end of text.""" + monkeypatch.setattr(settings, "sensitive_words", "secret") + response = client.post("/api/v1/detect", json={"text": "secret is the secret"}) + + assert response.status_code == 200 + data = response.json() + sw = [e for e in data["entities"] if e["type"] == "SENSITIVE_WORD"] + assert len(sw) == 2 + assert sw[0]["start"] == 0 + + def test_sensitive_word_entity_positions(self, client: TestClient, monkeypatch): + """Should return correct positions for sensitive words.""" + monkeypatch.setattr(settings, "sensitive_words", "secret") + text = "The secret is here" + response = client.post("/api/v1/detect", json={"text": text}) + + assert response.status_code == 200 + data = response.json() + sw = [e for e in data["entities"] if e["type"] == "SENSITIVE_WORD"] + assert len(sw) == 1 + assert text[sw[0]["start"] : sw[0]["end"]] == "secret"