mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-28 05:12:26 +00:00
143 lines
4.5 KiB
Python
143 lines
4.5 KiB
Python
"""Tests for taint tracking system (Phase 14.5)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from openjarvis.security.taint import (
|
|
SINK_POLICY,
|
|
TaintLabel,
|
|
TaintSet,
|
|
auto_detect_taint,
|
|
check_taint,
|
|
declassify,
|
|
propagate_taint,
|
|
)
|
|
|
|
|
|
class TestTaintSet:
|
|
def test_empty_taint(self):
|
|
ts = TaintSet()
|
|
assert not ts
|
|
assert not ts.labels
|
|
|
|
def test_from_labels(self):
|
|
ts = TaintSet.from_labels(TaintLabel.PII, TaintLabel.SECRET)
|
|
assert ts.has(TaintLabel.PII)
|
|
assert ts.has(TaintLabel.SECRET)
|
|
assert not ts.has(TaintLabel.EXTERNAL)
|
|
|
|
def test_union(self):
|
|
a = TaintSet.from_labels(TaintLabel.PII)
|
|
b = TaintSet.from_labels(TaintLabel.SECRET)
|
|
merged = a.union(b)
|
|
assert merged.has(TaintLabel.PII)
|
|
assert merged.has(TaintLabel.SECRET)
|
|
|
|
def test_frozen(self):
|
|
ts = TaintSet.from_labels(TaintLabel.PII)
|
|
# TaintSet is frozen dataclass
|
|
with pytest.raises(AttributeError):
|
|
ts.labels = frozenset()
|
|
|
|
def test_bool_true_when_has_labels(self):
|
|
ts = TaintSet.from_labels(TaintLabel.PII)
|
|
assert bool(ts)
|
|
|
|
def test_bool_false_when_empty(self):
|
|
ts = TaintSet()
|
|
assert not bool(ts)
|
|
|
|
|
|
class TestCheckTaint:
|
|
def test_clean_data_passes(self):
|
|
ts = TaintSet()
|
|
assert check_taint("web_search", ts) is None
|
|
|
|
def test_pii_blocked_for_web_search(self):
|
|
ts = TaintSet.from_labels(TaintLabel.PII)
|
|
result = check_taint("web_search", ts)
|
|
assert result is not None
|
|
assert "pii" in result.lower()
|
|
|
|
def test_secret_blocked_for_web_search(self):
|
|
ts = TaintSet.from_labels(TaintLabel.SECRET)
|
|
result = check_taint("web_search", ts)
|
|
assert result is not None
|
|
assert "secret" in result.lower()
|
|
|
|
def test_secret_blocked_for_channel_send(self):
|
|
ts = TaintSet.from_labels(TaintLabel.SECRET)
|
|
result = check_taint("channel_send", ts)
|
|
assert result is not None
|
|
|
|
def test_external_allowed_for_web_search(self):
|
|
ts = TaintSet.from_labels(TaintLabel.EXTERNAL)
|
|
assert check_taint("web_search", ts) is None
|
|
|
|
def test_unknown_tool_allowed(self):
|
|
ts = TaintSet.from_labels(TaintLabel.PII, TaintLabel.SECRET)
|
|
assert check_taint("calculator", ts) is None
|
|
|
|
def test_sink_policy_has_expected_tools(self):
|
|
assert "web_search" in SINK_POLICY
|
|
assert "channel_send" in SINK_POLICY
|
|
assert "code_interpreter" in SINK_POLICY
|
|
|
|
|
|
class TestDeclassify:
|
|
def test_remove_label(self):
|
|
ts = TaintSet.from_labels(TaintLabel.PII, TaintLabel.SECRET)
|
|
result = declassify(ts, TaintLabel.PII, "User consent given")
|
|
assert not result.has(TaintLabel.PII)
|
|
assert result.has(TaintLabel.SECRET)
|
|
|
|
def test_remove_nonexistent_label(self):
|
|
ts = TaintSet.from_labels(TaintLabel.PII)
|
|
result = declassify(ts, TaintLabel.SECRET, "Not present")
|
|
assert result.has(TaintLabel.PII)
|
|
|
|
|
|
class TestAutoDetect:
|
|
def test_detect_email(self):
|
|
ts = auto_detect_taint("Contact: user@example.com")
|
|
assert ts.has(TaintLabel.PII)
|
|
|
|
def test_detect_ssn(self):
|
|
ts = auto_detect_taint("SSN: 123-45-6789")
|
|
assert ts.has(TaintLabel.PII)
|
|
|
|
def test_detect_api_key(self):
|
|
ts = auto_detect_taint("Key: sk-abc123def456ghi789jkl012mno")
|
|
assert ts.has(TaintLabel.SECRET)
|
|
|
|
def test_detect_github_token(self):
|
|
ts = auto_detect_taint("Token: ghp_abcdefghijklmnopqrstuvwxyz0123456789")
|
|
assert ts.has(TaintLabel.SECRET)
|
|
|
|
def test_clean_text(self):
|
|
ts = auto_detect_taint("Hello, this is a normal message.")
|
|
assert not ts
|
|
|
|
def test_detect_private_key(self):
|
|
ts = auto_detect_taint("-----BEGIN RSA PRIVATE KEY-----\nMIIE...")
|
|
assert ts.has(TaintLabel.SECRET)
|
|
|
|
|
|
class TestPropagate:
|
|
def test_propagate_input_taint(self):
|
|
input_taint = TaintSet.from_labels(TaintLabel.EXTERNAL)
|
|
result = propagate_taint(input_taint, "Normal output")
|
|
assert result.has(TaintLabel.EXTERNAL)
|
|
|
|
def test_propagate_detects_new_taint(self):
|
|
input_taint = TaintSet()
|
|
result = propagate_taint(input_taint, "Found: user@example.com")
|
|
assert result.has(TaintLabel.PII)
|
|
|
|
def test_propagate_merges(self):
|
|
input_taint = TaintSet.from_labels(TaintLabel.EXTERNAL)
|
|
result = propagate_taint(input_taint, "Key: sk-abc123def456ghi789jkl012mno")
|
|
assert result.has(TaintLabel.EXTERNAL)
|
|
assert result.has(TaintLabel.SECRET)
|