This commit is contained in:
Alexander Myasoedov committed 2026-01-28 21:04:29 +02:00
1 parent 8d42a84a9d
commit bc7fdd7cfa
19 files changed
+491 -77

No files matched your search

+9 -3
View File
@@ -83,7 +83,9 @@ class YAMLRulesDatasetLoader:
severity_enums = None severity_enums = None
if self.severities: if self.severities:
severity_enums = [AttackRuleSeverity.from_string(s) for s in self.severities] severity_enums = [
AttackRuleSeverity.from_string(s) for s in self.severities
]
filtered = self._loader.filter_rules( filtered = self._loader.filter_rules(
rules, types=self.types, severities=severity_enums rules, types=self.types, severities=severity_enums
@@ -113,10 +115,14 @@ class YAMLRulesDatasetLoader:
severity_enums = None severity_enums = None
if self.severities: if self.severities:
severity_enums = [AttackRuleSeverity.from_string(s) for s in self.severities] severity_enums = [
AttackRuleSeverity.from_string(s) for s in self.severities
]
filtered = self._loader.filter_rules( filtered = self._loader.filter_rules(
all_rules, types=self.types, severities=severity_enums all_rules, types=self.types, severities=severity_enums
) )
return rules_to_dataset(filtered, name="YAML Rules (merged)", variables=variables) return rules_to_dataset(
filtered, name="YAML Rules (merged)", variables=variables
)
+5 -10
View File
@@ -1,4 +1,3 @@
import os
from pathlib import Path from pathlib import Path
import yaml import yaml
@@ -81,9 +80,7 @@ class RuleLoader:
return None return None
def load_rules_from_directory( def load_rules_from_directory(
self, self, directory: str | Path | None = None, recursive: bool = True
directory: str | Path | None = None,
recursive: bool = True
) -> list[AttackRule]: ) -> list[AttackRule]:
directory = Path(directory) if directory else self.rules_dir directory = Path(directory) if directory else self.rules_dir
if not directory or not directory.exists(): if not directory or not directory.exists():
@@ -91,7 +88,7 @@ class RuleLoader:
return [] return []
rules = [] rules = []
pattern = "**/*.yaml" if recursive else "*.yaml" # pattern = "**/*.yaml" if recursive else "*.yaml"
for ext in [".yaml", ".yml"]: for ext in [".yaml", ".yml"]:
glob_pattern = f"**/*{ext}" if recursive else f"*{ext}" glob_pattern = f"**/*{ext}" if recursive else f"*{ext}"
@@ -105,9 +102,7 @@ class RuleLoader:
return rules return rules
def load_multiple_directories( def load_multiple_directories(
self, self, directories: list[str | Path], recursive: bool = True
directories: list[str | Path],
recursive: bool = True
) -> list[AttackRule]: ) -> list[AttackRule]:
all_rules = [] all_rules = []
for directory in directories: for directory in directories:
@@ -133,6 +128,7 @@ class RuleLoader:
if name_pattern: if name_pattern:
import re import re
pattern = re.compile(name_pattern, re.IGNORECASE) pattern = re.compile(name_pattern, re.IGNORECASE)
result = [r for r in result if pattern.search(r.name)] result = [r for r in result if pattern.search(r.name)]
@@ -154,8 +150,7 @@ class RuleLoader:
def load_rules_from_directory( def load_rules_from_directory(
directory: str | Path, directory: str | Path, recursive: bool = True
recursive: bool = True
) -> list[AttackRule]: ) -> list[AttackRule]:
loader = RuleLoader() loader = RuleLoader()
return loader.load_rules_from_directory(directory, recursive) return loader.load_rules_from_directory(directory, recursive)
+14 -4
View File
@@ -38,10 +38,20 @@ class AttackRule:
pass_conditions=data.get("pass_conditions", []), pass_conditions=data.get("pass_conditions", []),
fail_conditions=data.get("fail_conditions", []), fail_conditions=data.get("fail_conditions", []),
source=data.get("source"), source=data.get("source"),
metadata={k: v for k, v in data.items() if k not in { metadata={
"name", "type", "prompt", "severity", k: v
"pass_conditions", "fail_conditions", "source" for k, v in data.items()
}}, if k
not in {
"name",
"type",
"prompt",
"severity",
"pass_conditions",
"fail_conditions",
"source",
}
},
) )
def to_dict(self) -> dict[str, Any]: def to_dict(self) -> dict[str, Any]:
+179
View File
@@ -0,0 +1,179 @@
"""Security utilities and validation for agentic_security."""
from functools import wraps
from collections.abc import Callable
from urllib.parse import urlparse
import hashlib
import hmac
import os
import re
class SecurityValidator:
"""Input validation and sanitization."""
ALLOWED_URL_SCHEMES = {"http", "https"}
MAX_URL_LENGTH = 2048
MAX_FILE_SIZE = 10 * 1024 * 1024 # 10MB
@staticmethod
def validate_url(url: str, allowed_hosts: list[str] | None = None) -> bool:
"""Validate URL for SSRF prevention."""
if len(url) > SecurityValidator.MAX_URL_LENGTH:
return False
try:
parsed = urlparse(url)
if parsed.scheme not in SecurityValidator.ALLOWED_URL_SCHEMES:
return False
if not parsed.netloc:
return False
if parsed.netloc in ["localhost", "127.0.0.1", "0.0.0.0"]:
return False
if parsed.netloc.startswith("169.254."):
return False
if parsed.netloc.startswith("10.") or parsed.netloc.startswith("192.168."):
return False
if allowed_hosts and parsed.netloc not in allowed_hosts:
return False
return True
except Exception:
return False
@staticmethod
def sanitize_filename(filename: str) -> str:
"""Sanitize filename to prevent path traversal."""
filename = os.path.basename(filename)
filename = re.sub(r"[^\w\s.-]", "", filename)
filename = filename.strip()
if not filename or filename in [".", ".."]:
raise ValueError("Invalid filename")
return filename
@staticmethod
def validate_file_size(size: int) -> bool:
"""Validate file size."""
return 0 < size <= SecurityValidator.MAX_FILE_SIZE
@staticmethod
def validate_csv_content(content: str) -> bool:
"""Basic CSV validation."""
if not content or len(content) > SecurityValidator.MAX_FILE_SIZE:
return False
lines = content.split("\n", 2)
if not lines:
return False
return True
class SecretManager:
"""Secure secret handling."""
@staticmethod
def get_secret(key: str, default: str | None = None) -> str | None:
"""Get secret from environment."""
value = os.getenv(key, default)
if value and value.startswith("$"):
env_key = value[1:]
value = os.getenv(env_key, default)
return value
@staticmethod
def hash_secret(secret: str, salt: str | None = None) -> str:
"""Hash a secret value."""
if salt is None:
salt = os.urandom(32).hex()
hashed = hashlib.pbkdf2_hmac("sha256", secret.encode(), salt.encode(), 100000)
return f"{salt}${hashed.hex()}"
@staticmethod
def verify_secret(secret: str, hashed: str) -> bool:
"""Verify a secret against its hash."""
try:
salt, expected = hashed.split("$", 1)
actual = hashlib.pbkdf2_hmac(
"sha256", secret.encode(), salt.encode(), 100000
)
return hmac.compare_digest(actual.hex(), expected)
except Exception:
return False
class RateLimiter:
"""Simple in-memory rate limiter."""
def __init__(self, max_requests: int, window_seconds: int):
self.max_requests = max_requests
self.window_seconds = window_seconds
self._requests: dict[str, list[float]] = {}
def is_allowed(self, key: str) -> bool:
"""Check if request is allowed."""
import time
now = time.time()
if key not in self._requests:
self._requests[key] = []
self._requests[key] = [
ts for ts in self._requests[key] if now - ts < self.window_seconds
]
if len(self._requests[key]) >= self.max_requests:
return False
self._requests[key].append(now)
return True
def reset(self, key: str):
"""Reset rate limit for key."""
self._requests.pop(key, None)
def require_auth(func: Callable) -> Callable:
"""Decorator to require authentication."""
@wraps(func)
async def wrapper(*args, **kwargs):
# TODO: Implement actual auth check
# For now, check if API key is present
api_key = kwargs.get("api_key") or os.getenv("API_KEY")
if not api_key:
from fastapi import HTTPException
raise HTTPException(status_code=401, detail="Authentication required")
return await func(*args, **kwargs)
return wrapper
def sanitize_log_output(data: str | dict) -> str:
"""Remove sensitive data from logs."""
if isinstance(data, dict):
data = str(data)
patterns = [
(r'(api[_-]?key["\s:=]+)["\']?[\w-]+', r"\1***"),
(r'(token["\s:=]+)["\']?[\w-]+', r"\1***"),
(r'(password["\s:=]+)["\']?[\w-]+', r"\1***"),
(r'(secret["\s:=]+)["\']?[\w-]+', r"\1***"),
(r"Bearer\s+[\w-]+", "Bearer ***"),
]
for pattern, replacement in patterns:
data = re.sub(pattern, replacement, data, flags=re.IGNORECASE)
return data
+1 -2
View File
@@ -8,8 +8,7 @@ logger = logging.getLogger(__name__)
class FuzzRunnable(Protocol): class FuzzRunnable(Protocol):
"""Protocol for objects that can be run in a fuzzing chain.""" """Protocol for objects that can be run in a fuzzing chain."""
async def run(self, **kwargs: Any) -> str: async def run(self, **kwargs: Any) -> str: ...
...
class FuzzNode: class FuzzNode:
@@ -36,6 +36,7 @@ class AnthropicProvider(BaseLLMProvider):
def _get_client(self) -> Any: def _get_client(self) -> Any:
if self._client is None: if self._client is None:
import anthropic import anthropic
kwargs: dict[str, Any] = {"api_key": self.api_key} kwargs: dict[str, Any] = {"api_key": self.api_key}
if self.base_url: if self.base_url:
kwargs["base_url"] = self.base_url kwargs["base_url"] = self.base_url
@@ -45,6 +46,7 @@ class AnthropicProvider(BaseLLMProvider):
def _get_async_client(self) -> Any: def _get_async_client(self) -> Any:
if self._async_client is None: if self._async_client is None:
import anthropic import anthropic
kwargs: dict[str, Any] = {"api_key": self.api_key} kwargs: dict[str, Any] = {"api_key": self.api_key}
if self.base_url: if self.base_url:
kwargs["base_url"] = self.base_url kwargs["base_url"] = self.base_url
@@ -95,6 +97,7 @@ class AnthropicProvider(BaseLLMProvider):
def _handle_error(self, e: Exception) -> None: def _handle_error(self, e: Exception) -> None:
import anthropic import anthropic
if isinstance(e, anthropic.RateLimitError): if isinstance(e, anthropic.RateLimitError):
raise LLMRateLimitError(str(e)) from e raise LLMRateLimitError(str(e)) from e
if isinstance(e, anthropic.APIError): if isinstance(e, anthropic.APIError):
+2
View File
@@ -20,6 +20,7 @@ class LLMRateLimitError(LLMProviderError):
@dataclass @dataclass
class LLMMessage: class LLMMessage:
"""A message in a chat conversation.""" """A message in a chat conversation."""
role: str # "system", "user", or "assistant" role: str # "system", "user", or "assistant"
content: str content: str
@@ -27,6 +28,7 @@ class LLMMessage:
@dataclass @dataclass
class LLMResponse: class LLMResponse:
"""Response from an LLM provider.""" """Response from an LLM provider."""
content: str content: str
model: str | None = None model: str | None = None
finish_reason: str | None = None finish_reason: str | None = None
@@ -14,6 +14,7 @@ def _ensure_registered() -> None:
return return
from agentic_security.llm_providers.openai_provider import OpenAIProvider from agentic_security.llm_providers.openai_provider import OpenAIProvider
from agentic_security.llm_providers.anthropic_provider import AnthropicProvider from agentic_security.llm_providers.anthropic_provider import AnthropicProvider
_PROVIDERS["openai"] = OpenAIProvider _PROVIDERS["openai"] = OpenAIProvider
_PROVIDERS["anthropic"] = AnthropicProvider _PROVIDERS["anthropic"] = AnthropicProvider
@@ -36,13 +36,17 @@ class OpenAIProvider(BaseLLMProvider):
def _get_client(self) -> Any: def _get_client(self) -> Any:
if self._client is None: if self._client is None:
import openai import openai
self._client = openai.OpenAI(api_key=self.api_key, base_url=self.base_url) self._client = openai.OpenAI(api_key=self.api_key, base_url=self.base_url)
return self._client return self._client
def _get_async_client(self) -> Any: def _get_async_client(self) -> Any:
if self._async_client is None: if self._async_client is None:
import openai import openai
self._async_client = openai.AsyncOpenAI(api_key=self.api_key, base_url=self.base_url)
self._async_client = openai.AsyncOpenAI(
api_key=self.api_key, base_url=self.base_url
)
return self._async_client return self._async_client
@classmethod @classmethod
@@ -79,6 +83,7 @@ class OpenAIProvider(BaseLLMProvider):
def _handle_error(self, e: Exception) -> None: def _handle_error(self, e: Exception) -> None:
import openai import openai
if isinstance(e, openai.RateLimitError): if isinstance(e, openai.RateLimitError):
raise LLMRateLimitError(str(e)) from e raise LLMRateLimitError(str(e)) from e
raise LLMProviderError(str(e)) from e raise LLMProviderError(str(e)) from e
@@ -96,11 +96,13 @@ class HybridRefusalClassifier:
self for method chaining self for method chaining
""" """
detector_name = name or detector.__class__.__name__ detector_name = name or detector.__class__.__name__
self._detectors.append(DetectorConfig( self._detectors.append(
detector=detector, DetectorConfig(
weight=weight, detector=detector,
name=detector_name, weight=weight,
)) name=detector_name,
)
)
return self return self
def classify(self, response: str) -> HybridResult: def classify(self, response: str) -> HybridResult:
@@ -117,11 +119,13 @@ class HybridRefusalClassifier:
is_refusal = config.detector.is_refusal(response) is_refusal = config.detector.is_refusal(response)
except Exception: except Exception:
continue # Skip failed detectors continue # Skip failed detectors
results.append(DetectionResult( results.append(
method=config.name, DetectionResult(
is_refusal=is_refusal, method=config.name,
weight=config.weight, is_refusal=is_refusal,
)) weight=config.weight,
)
)
if not results: if not results:
return HybridResult(is_refusal=False, confidence=0.0) return HybridResult(is_refusal=False, confidence=0.0)
@@ -134,7 +138,9 @@ class HybridRefusalClassifier:
# Check unanimous requirement # Check unanimous requirement
if self.require_unanimous: if self.require_unanimous:
all_agree = all(r.is_refusal for r in results) or all(not r.is_refusal for r in results) all_agree = all(r.is_refusal for r in results) or all(
not r.is_refusal for r in results
)
if not all_agree: if not all_agree:
# Disagreement - return uncertain result # Disagreement - return uncertain result
return HybridResult( return HybridResult(
+8 -4
View File
@@ -56,16 +56,20 @@ class TestRulesToDataset:
class TestLoadRulesAsDataset: class TestLoadRulesAsDataset:
def test_basic_load(self): def test_basic_load(self):
with tempfile.TemporaryDirectory() as tmpdir: with tempfile.TemporaryDirectory() as tmpdir:
(Path(tmpdir) / "rule1.yaml").write_text(""" (Path(tmpdir) / "rule1.yaml").write_text(
"""
name: test1 name: test1
type: jailbreak type: jailbreak
prompt: Jailbreak prompt prompt: Jailbreak prompt
""") """
(Path(tmpdir) / "rule2.yaml").write_text(""" )
(Path(tmpdir) / "rule2.yaml").write_text(
"""
name: test2 name: test2
type: harmful type: harmful
prompt: Harmful prompt prompt: Harmful prompt
""") """
)
dataset = load_rules_as_dataset(tmpdir) dataset = load_rules_as_dataset(tmpdir)
assert len(dataset.prompts) == 2 assert len(dataset.prompts) == 2
+14 -12
View File
@@ -77,15 +77,15 @@ severity: high
def test_load_rule_from_file(self): def test_load_rule_from_file(self):
loader = RuleLoader() loader = RuleLoader()
with tempfile.NamedTemporaryFile( with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
mode="w", suffix=".yaml", delete=False f.write(
) as f: """
f.write("""
name: file_test name: file_test
type: harmful type: harmful
severity: medium severity: medium
prompt: Test prompt from file prompt: Test prompt from file
""") """
)
f.flush() f.flush()
rule = loader.load_rule_from_file(f.name) rule = loader.load_rule_from_file(f.name)
@@ -96,9 +96,7 @@ prompt: Test prompt from file
def test_load_rule_from_file_wrong_extension(self): def test_load_rule_from_file_wrong_extension(self):
loader = RuleLoader() loader = RuleLoader()
with tempfile.NamedTemporaryFile( with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
mode="w", suffix=".txt", delete=False
) as f:
f.write("name: test\nprompt: test") f.write("name: test\nprompt: test")
f.flush() f.flush()
rule = loader.load_rule_from_file(f.name) rule = loader.load_rule_from_file(f.name)
@@ -110,16 +108,20 @@ prompt: Test prompt from file
with tempfile.TemporaryDirectory() as tmpdir: with tempfile.TemporaryDirectory() as tmpdir:
rule1_path = Path(tmpdir) / "rule1.yaml" rule1_path = Path(tmpdir) / "rule1.yaml"
rule2_path = Path(tmpdir) / "rule2.yml" rule2_path = Path(tmpdir) / "rule2.yml"
rule1_path.write_text(""" rule1_path.write_text(
"""
name: rule1 name: rule1
type: jailbreak type: jailbreak
prompt: First rule prompt: First rule
""") """
rule2_path.write_text(""" )
rule2_path.write_text(
"""
name: rule2 name: rule2
type: harmful type: harmful
prompt: Second rule prompt: Second rule
""") """
)
loader = RuleLoader() loader = RuleLoader()
rules = loader.load_rules_from_directory(tmpdir) rules = loader.load_rules_from_directory(tmpdir)
+6 -1
View File
@@ -87,7 +87,12 @@ class TestAttackRule:
rule = AttackRule(name="test", type="jailbreak", prompt="Test") rule = AttackRule(name="test", type="jailbreak", prompt="Test")
result = rule.to_dict() result = rule.to_dict()
assert result == snapshot( assert result == snapshot(
{"name": "test", "type": "jailbreak", "prompt": "Test", "severity": "medium"} {
"name": "test",
"type": "jailbreak",
"prompt": "Test",
"severity": "medium",
}
) )
def test_render_prompt_no_variables(self): def test_render_prompt_no_variables(self):
+27 -19
View File
@@ -116,11 +116,13 @@ class TestFuzzChain:
result = await chain.run(input="initial") result = await chain.run(input="initial")
assert result == "final result" assert result == "final result"
assert llm.prompts == snapshot([ assert llm.prompts == snapshot(
"First: initial", [
"Second: step1 result", "First: initial",
"Third: step2 result", "Second: step1 result",
]) "Third: step2 result",
]
)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_chain_with_custom_variables(self): async def test_chain_with_custom_variables(self):
@@ -131,10 +133,12 @@ class TestFuzzChain:
result = await chain.run(topic="security", input="test prompt") result = await chain.run(topic="security", input="test prompt")
assert result == "evaluated" assert result == "evaluated"
assert llm.prompts == snapshot([ assert llm.prompts == snapshot(
"Analyze security: test prompt", [
"Evaluate: analyzed", "Analyze security: test prompt",
]) "Evaluate: analyzed",
]
)
def test_pipe_chain_to_node(self): def test_pipe_chain_to_node(self):
llm = MockLLMProvider() llm = MockLLMProvider()
@@ -158,11 +162,13 @@ class TestFuzzChain:
def test_len(self): def test_len(self):
llm = MockLLMProvider() llm = MockLLMProvider()
chain = FuzzChain([ chain = FuzzChain(
FuzzNode(llm, "A"), [
FuzzNode(llm, "B"), FuzzNode(llm, "A"),
FuzzNode(llm, "C"), FuzzNode(llm, "B"),
]) FuzzNode(llm, "C"),
]
)
assert len(chain) == 3 assert len(chain) == 3
def test_repr(self): def test_repr(self):
@@ -185,11 +191,13 @@ class TestPipeOperatorChaining:
result = await chain.run(input="start") result = await chain.run(input="start")
assert result == "c" assert result == "c"
assert llm.prompts == snapshot([ assert llm.prompts == snapshot(
"Step1: start", [
"Step2: a", "Step1: start",
"Step3: b", "Step2: a",
]) "Step3: b",
]
)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_chain_with_different_providers(self): async def test_chain_with_different_providers(self):
@@ -5,7 +5,11 @@ from unittest.mock import MagicMock, AsyncMock, patch
from inline_snapshot import snapshot from inline_snapshot import snapshot
from agentic_security.llm_providers.anthropic_provider import AnthropicProvider from agentic_security.llm_providers.anthropic_provider import AnthropicProvider
from agentic_security.llm_providers.base import LLMMessage, LLMProviderError, LLMRateLimitError from agentic_security.llm_providers.base import (
LLMMessage,
LLMProviderError,
LLMRateLimitError,
)
class TestAnthropicProviderInit: class TestAnthropicProviderInit:
@@ -209,13 +213,19 @@ class TestAnthropicProviderErrors:
def test_handle_rate_limit_error(self, provider): def test_handle_rate_limit_error(self, provider):
import anthropic import anthropic
with pytest.raises(LLMRateLimitError): with pytest.raises(LLMRateLimitError):
provider._handle_error(anthropic.RateLimitError("rate limited", response=MagicMock(), body={})) provider._handle_error(
anthropic.RateLimitError("rate limited", response=MagicMock(), body={})
)
def test_handle_api_error(self, provider): def test_handle_api_error(self, provider):
import anthropic import anthropic
with pytest.raises(LLMProviderError): with pytest.raises(LLMProviderError):
provider._handle_error(anthropic.APIError("api error", request=MagicMock(), body={})) provider._handle_error(
anthropic.APIError("api error", request=MagicMock(), body={})
)
def test_handle_generic_error(self, provider): def test_handle_generic_error(self, provider):
with pytest.raises(LLMProviderError): with pytest.raises(LLMProviderError):
+7 -1
View File
@@ -9,7 +9,11 @@ from agentic_security.llm_providers.factory import (
list_providers, list_providers,
register_provider, register_provider,
) )
from agentic_security.llm_providers.base import BaseLLMProvider, LLMProviderError, LLMResponse from agentic_security.llm_providers.base import (
BaseLLMProvider,
LLMProviderError,
LLMResponse,
)
class TestListProviders: class TestListProviders:
@@ -26,11 +30,13 @@ class TestListProviders:
class TestGetProviderClass: class TestGetProviderClass:
def test_get_openai(self): def test_get_openai(self):
from agentic_security.llm_providers.openai_provider import OpenAIProvider from agentic_security.llm_providers.openai_provider import OpenAIProvider
cls = get_provider_class("openai") cls = get_provider_class("openai")
assert cls is OpenAIProvider assert cls is OpenAIProvider
def test_get_anthropic(self): def test_get_anthropic(self):
from agentic_security.llm_providers.anthropic_provider import AnthropicProvider from agentic_security.llm_providers.anthropic_provider import AnthropicProvider
cls = get_provider_class("anthropic") cls = get_provider_class("anthropic")
assert cls is AnthropicProvider assert cls is AnthropicProvider
@@ -5,7 +5,11 @@ from unittest.mock import MagicMock, AsyncMock, patch
from inline_snapshot import snapshot from inline_snapshot import snapshot
from agentic_security.llm_providers.openai_provider import OpenAIProvider from agentic_security.llm_providers.openai_provider import OpenAIProvider
from agentic_security.llm_providers.base import LLMMessage, LLMProviderError, LLMRateLimitError from agentic_security.llm_providers.base import (
LLMMessage,
LLMProviderError,
LLMRateLimitError,
)
class TestOpenAIProviderInit: class TestOpenAIProviderInit:
@@ -111,7 +115,9 @@ class TestOpenAIProviderSync:
mock_response.usage = None mock_response.usage = None
with patch.object(provider, "_get_client") as mock_client: with patch.object(provider, "_get_client") as mock_client:
mock_client.return_value.chat.completions.create.return_value = mock_response mock_client.return_value.chat.completions.create.return_value = (
mock_response
)
result = provider.sync_generate("Hello") result = provider.sync_generate("Hello")
assert result.content == snapshot("Response") assert result.content == snapshot("Response")
@@ -126,7 +132,9 @@ class TestOpenAIProviderSync:
messages = [LLMMessage(role="user", content="Hi")] messages = [LLMMessage(role="user", content="Hi")]
with patch.object(provider, "_get_client") as mock_client: with patch.object(provider, "_get_client") as mock_client:
mock_client.return_value.chat.completions.create.return_value = mock_response mock_client.return_value.chat.completions.create.return_value = (
mock_response
)
result = provider.sync_chat(messages) result = provider.sync_chat(messages)
assert result.content == snapshot("Chat response") assert result.content == snapshot("Chat response")
@@ -196,8 +204,11 @@ class TestOpenAIProviderErrors:
def test_handle_rate_limit_error(self, provider): def test_handle_rate_limit_error(self, provider):
import openai import openai
with pytest.raises(LLMRateLimitError): with pytest.raises(LLMRateLimitError):
provider._handle_error(openai.RateLimitError("rate limited", response=MagicMock(), body={})) provider._handle_error(
openai.RateLimitError("rate limited", response=MagicMock(), body={})
)
def test_handle_generic_error(self, provider): def test_handle_generic_error(self, provider):
with pytest.raises(LLMProviderError): with pytest.raises(LLMProviderError):
@@ -45,7 +45,9 @@ class TestDetectionResult:
def test_weighted_score_cases(self): def test_weighted_score_cases(self):
for is_refusal, weight, expected in detection_result_cases: for is_refusal, weight, expected in detection_result_cases:
result = DetectionResult(method="test", is_refusal=is_refusal, weight=weight) result = DetectionResult(
method="test", is_refusal=is_refusal, weight=weight
)
assert result.weighted_score == expected assert result.weighted_score == expected
def test_default_weight(self): def test_default_weight(self):
@@ -234,7 +236,14 @@ factory_cases = [
({"ml_detector": MockDetector(True)}, 1), ({"ml_detector": MockDetector(True)}, 1),
({"llm_detector": MockDetector(True)}, 1), ({"llm_detector": MockDetector(True)}, 1),
({"marker_detector": MockDetector(True), "ml_detector": MockDetector(False)}, 2), ({"marker_detector": MockDetector(True), "ml_detector": MockDetector(False)}, 2),
({"marker_detector": MockDetector(True), "ml_detector": MockDetector(False), "llm_detector": MockDetector(True)}, 3), (
{
"marker_detector": MockDetector(True),
"ml_detector": MockDetector(False),
"llm_detector": MockDetector(True),
},
3,
),
] ]
+153
View File
@@ -0,0 +1,153 @@
"""Unit tests for security module."""
import pytest
from agentic_security.core.security import (
SecurityValidator,
SecretManager,
RateLimiter,
sanitize_log_output,
)
class TestSecurityValidator:
def test_validate_url_valid(self):
assert SecurityValidator.validate_url("https://example.com/path")
assert SecurityValidator.validate_url("http://api.example.com")
def test_validate_url_invalid_scheme(self):
assert not SecurityValidator.validate_url("ftp://example.com")
assert not SecurityValidator.validate_url("file:///etc/passwd")
def test_validate_url_localhost(self):
assert not SecurityValidator.validate_url("http://localhost/api")
assert not SecurityValidator.validate_url("http://127.0.0.1/api")
assert not SecurityValidator.validate_url("http://0.0.0.0/api")
def test_validate_url_private_ip(self):
assert not SecurityValidator.validate_url("http://10.0.0.1")
assert not SecurityValidator.validate_url("http://192.168.1.1")
assert not SecurityValidator.validate_url("http://169.254.1.1")
def test_validate_url_allowed_hosts(self):
allowed = ["api.example.com"]
assert SecurityValidator.validate_url("https://api.example.com", allowed)
assert not SecurityValidator.validate_url("https://evil.com", allowed)
def test_validate_url_too_long(self):
long_url = "https://example.com/" + "a" * 3000
assert not SecurityValidator.validate_url(long_url)
def test_sanitize_filename(self):
assert SecurityValidator.sanitize_filename("test.csv") == "test.csv"
assert SecurityValidator.sanitize_filename("../../../etc/passwd") == "passwd"
assert SecurityValidator.sanitize_filename("test/file.txt") == "file.txt"
assert (
SecurityValidator.sanitize_filename("file with spaces.txt")
== "file with spaces.txt"
)
def test_sanitize_filename_invalid(self):
with pytest.raises(ValueError):
SecurityValidator.sanitize_filename(".")
with pytest.raises(ValueError):
SecurityValidator.sanitize_filename("..")
with pytest.raises(ValueError):
SecurityValidator.sanitize_filename("")
def test_validate_file_size(self):
assert SecurityValidator.validate_file_size(1024)
assert SecurityValidator.validate_file_size(1024 * 1024)
assert not SecurityValidator.validate_file_size(0)
assert not SecurityValidator.validate_file_size(-1)
assert not SecurityValidator.validate_file_size(20 * 1024 * 1024)
def test_validate_csv_content(self):
assert SecurityValidator.validate_csv_content("col1,col2\nval1,val2")
assert not SecurityValidator.validate_csv_content("")
assert not SecurityValidator.validate_csv_content("x" * (11 * 1024 * 1024))
class TestSecretManager:
def test_hash_and_verify_secret(self):
secret = "my-secret-key"
hashed = SecretManager.hash_secret(secret)
assert SecretManager.verify_secret(secret, hashed)
assert not SecretManager.verify_secret("wrong-secret", hashed)
def test_hash_secret_with_salt(self):
secret = "my-secret"
hashed1 = SecretManager.hash_secret(secret, "salt1")
hashed2 = SecretManager.hash_secret(secret, "salt2")
assert hashed1 != hashed2
def test_verify_secret_invalid_format(self):
assert not SecretManager.verify_secret("secret", "invalid-hash")
class TestRateLimiter:
def test_rate_limiter_allows_requests(self):
limiter = RateLimiter(max_requests=3, window_seconds=60)
assert limiter.is_allowed("user1")
assert limiter.is_allowed("user1")
assert limiter.is_allowed("user1")
def test_rate_limiter_blocks_excess(self):
limiter = RateLimiter(max_requests=2, window_seconds=60)
assert limiter.is_allowed("user1")
assert limiter.is_allowed("user1")
assert not limiter.is_allowed("user1")
def test_rate_limiter_separate_keys(self):
limiter = RateLimiter(max_requests=2, window_seconds=60)
assert limiter.is_allowed("user1")
assert limiter.is_allowed("user1")
assert limiter.is_allowed("user2")
assert limiter.is_allowed("user2")
def test_rate_limiter_reset(self):
limiter = RateLimiter(max_requests=1, window_seconds=60)
assert limiter.is_allowed("user1")
assert not limiter.is_allowed("user1")
limiter.reset("user1")
assert limiter.is_allowed("user1")
class TestSanitizeLogOutput:
def test_sanitize_api_key(self):
data = 'api_key="sk-1234567890"'
result = sanitize_log_output(data)
assert "sk-1234567890" not in result
assert "***" in result
def test_sanitize_token(self):
data = "token: abc123xyz"
result = sanitize_log_output(data)
assert "abc123xyz" not in result
def test_sanitize_password(self):
data = {"password": "secret123"}
result = sanitize_log_output(data)
assert "secret123" not in result
def test_sanitize_bearer_token(self):
data = "Authorization: Bearer eyJhbGc..."
result = sanitize_log_output(data)
assert "eyJhbGc" not in result
assert "Bearer ***" in result
def test_preserves_non_sensitive(self):
data = "user_id=123 name=John"
result = sanitize_log_output(data)
assert "user_id=123" in result
assert "name=John" in result