Compare commits
1 Commits
fix/839
...
fix/749-v2
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f7f89e15ff |
136
tests/test_batch_executor.py
Normal file
136
tests/test_batch_executor.py
Normal file
@@ -0,0 +1,136 @@
|
||||
"""Tests for batch tool execution — Issue #749."""
|
||||
import asyncio
|
||||
import sys
|
||||
from pathlib import Path
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
|
||||
from tools.batch_executor import (
|
||||
ToolSafety, ToolCall, BatchResult,
|
||||
classify_tool_safety, classify_calls,
|
||||
execute_batch_sync, get_tool_safety_report
|
||||
)
|
||||
|
||||
|
||||
class TestClassification:
|
||||
def test_parallel_safe_read(self):
|
||||
assert classify_tool_safety("file_read") == ToolSafety.PARALLEL_SAFE
|
||||
|
||||
def test_sequential_write(self):
|
||||
assert classify_tool_safety("file_write") == ToolSafety.SEQUENTIAL
|
||||
|
||||
def test_destructive_terminal(self):
|
||||
assert classify_tool_safety("terminal") == ToolSafety.DESTRUCTIVE
|
||||
|
||||
def test_unknown_defaults_sequential(self):
|
||||
assert classify_tool_safety("unknown_tool") == ToolSafety.SEQUENTIAL
|
||||
|
||||
def test_prefix_match(self):
|
||||
assert classify_tool_safety("file_read_special") == ToolSafety.PARALLEL_SAFE
|
||||
|
||||
|
||||
class TestClassifyCalls:
|
||||
def test_classifies_multiple(self):
|
||||
calls = [
|
||||
{"name": "file_read", "arguments": "{}"},
|
||||
{"name": "file_write", "arguments": "{}"},
|
||||
{"name": "terminal", "arguments": "{}"},
|
||||
]
|
||||
result = classify_calls(calls)
|
||||
assert len(result) == 3
|
||||
assert result[0].safety == ToolSafety.PARALLEL_SAFE
|
||||
assert result[1].safety == ToolSafety.SEQUENTIAL
|
||||
assert result[2].safety == ToolSafety.DESTRUCTIVE
|
||||
|
||||
|
||||
class TestBatchExecution:
|
||||
def test_parallel_execution(self):
|
||||
"""Parallel-safe calls should execute faster than sequential."""
|
||||
import time
|
||||
|
||||
def slow_executor(name, args):
|
||||
time.sleep(0.1)
|
||||
return f"result_{name}"
|
||||
|
||||
calls = [
|
||||
{"name": "file_read", "arguments": "{}"},
|
||||
{"name": "file_search", "arguments": "{}"},
|
||||
{"name": "web_search", "arguments": "{}"},
|
||||
]
|
||||
|
||||
start = time.time()
|
||||
result = execute_batch_sync(calls, slow_executor)
|
||||
duration = time.time() - start
|
||||
|
||||
# Should be faster than 0.3s (3 * 0.1) since parallel
|
||||
assert duration < 0.25
|
||||
assert result.parallel_count == 3
|
||||
assert len(result.errors) == 0
|
||||
|
||||
def test_sequential_execution(self):
|
||||
"""Sequential calls should execute one at a time."""
|
||||
import time
|
||||
|
||||
def slow_executor(name, args):
|
||||
time.sleep(0.05)
|
||||
return f"result_{name}"
|
||||
|
||||
calls = [
|
||||
{"name": "file_write", "arguments": "{}"},
|
||||
{"name": "file_patch", "arguments": "{}"},
|
||||
]
|
||||
|
||||
start = time.time()
|
||||
result = execute_batch_sync(calls, slow_executor)
|
||||
duration = time.time() - start
|
||||
|
||||
# Should take at least 0.1s (2 * 0.05) since sequential
|
||||
assert duration >= 0.1
|
||||
assert result.sequential_count == 2
|
||||
|
||||
def test_mixed_execution(self):
|
||||
"""Mixed calls: parallel first, then sequential."""
|
||||
calls = [
|
||||
{"name": "file_read", "arguments": "{}"},
|
||||
{"name": "file_write", "arguments": "{}"},
|
||||
{"name": "web_search", "arguments": "{}"},
|
||||
]
|
||||
|
||||
def executor(name, args):
|
||||
return f"result_{name}"
|
||||
|
||||
result = execute_batch_sync(calls, executor)
|
||||
assert result.parallel_count == 2
|
||||
assert result.sequential_count == 1
|
||||
assert len(result.errors) == 0
|
||||
|
||||
def test_error_handling(self):
|
||||
"""Errors in one call shouldn't stop others."""
|
||||
def failing_executor(name, args):
|
||||
if name == "file_write":
|
||||
raise Exception("Write failed")
|
||||
return "ok"
|
||||
|
||||
calls = [
|
||||
{"name": "file_read", "arguments": "{}"},
|
||||
{"name": "file_write", "arguments": "{}"},
|
||||
]
|
||||
|
||||
result = execute_batch_sync(calls, failing_executor)
|
||||
assert len(result.errors) == 1
|
||||
assert "file_write" in result.errors[0]
|
||||
|
||||
|
||||
class TestSafetyReport:
|
||||
def test_report_format(self):
|
||||
calls = [
|
||||
ToolCall(name="file_read", args={}, safety=ToolSafety.PARALLEL_SAFE, duration=0.1),
|
||||
ToolCall(name="file_write", args={}, safety=ToolSafety.SEQUENTIAL, duration=0.2),
|
||||
]
|
||||
report = get_tool_safety_report(calls)
|
||||
assert "Parallel-safe: 1" in report
|
||||
assert "Sequential: 1" in report
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import pytest
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -1,122 +0,0 @@
|
||||
"""Tests for credential redaction — Issue #839."""
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
|
||||
from tools.credential_redaction import (
|
||||
redact_credentials, should_auto_mask, mask_config_values,
|
||||
redact_tool_output, RedactionResult
|
||||
)
|
||||
|
||||
|
||||
class TestRedactCredentials:
|
||||
def test_openai_key(self):
|
||||
text = "API key: sk-abc123def456ghi789jkl012mno345pqr678stu901vwx"
|
||||
result = redact_credentials(text)
|
||||
assert result.was_redacted
|
||||
assert "sk-abc" not in result.text
|
||||
assert "[REDACTED" in result.text
|
||||
|
||||
def test_github_pat(self):
|
||||
text = "token: ghp_1234567890abcdefghijklmnopqrstuvwxyz"
|
||||
result = redact_credentials(text)
|
||||
assert result.was_redacted
|
||||
assert "ghp_" not in result.text
|
||||
|
||||
def test_bearer_token(self):
|
||||
text = "Authorization: Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9"
|
||||
result = redact_credentials(text)
|
||||
assert result.was_redacted
|
||||
assert "Bearer eyJ" not in result.text
|
||||
|
||||
def test_password_assignment(self):
|
||||
text = 'password: "supersecret123"'
|
||||
result = redact_credentials(text)
|
||||
assert result.was_redacted
|
||||
|
||||
def test_clean_text(self):
|
||||
text = "Hello world, no credentials here"
|
||||
result = redact_credentials(text)
|
||||
assert not result.was_redacted
|
||||
assert result.text == text
|
||||
|
||||
def test_empty_text(self):
|
||||
result = redact_credentials("")
|
||||
assert not result.was_redacted
|
||||
|
||||
|
||||
class TestShouldAutoMask:
|
||||
def test_env_file(self):
|
||||
assert should_auto_mask(".env") == True
|
||||
|
||||
def test_config_file(self):
|
||||
assert should_auto_mask("config.yaml") == True
|
||||
|
||||
def test_token_file(self):
|
||||
assert should_auto_mask("gitea_token") == True
|
||||
|
||||
def test_normal_file(self):
|
||||
assert should_auto_mask("readme.md") == False
|
||||
|
||||
|
||||
class TestMaskConfigValues:
|
||||
def test_env_api_key(self):
|
||||
text = "API_KEY=sk-abc123def456"
|
||||
result = mask_config_values(text)
|
||||
assert "sk-abc" not in result
|
||||
assert "[REDACTED]" in result
|
||||
|
||||
def test_yaml_token(self):
|
||||
text = 'token: "ghp_1234567890"'
|
||||
result = mask_config_values(text)
|
||||
assert "ghp_" not in result
|
||||
assert "[REDACTED]" in result
|
||||
|
||||
def test_preserves_structure(self):
|
||||
text = "API_KEY=secret\nOTHER=value"
|
||||
result = mask_config_values(text)
|
||||
assert "OTHER=value" in result # Non-credential preserved
|
||||
|
||||
|
||||
class TestRedactToolOutput:
|
||||
def test_string_output(self):
|
||||
output = "Result: sk-abc123def456ghi789jkl012mno345pqr678stu901vwx"
|
||||
redacted, notice = redact_tool_output("file_read", output)
|
||||
assert "sk-abc123" not in redacted
|
||||
assert notice is not None
|
||||
|
||||
def test_dict_output(self):
|
||||
output = {"content": "token: ghp_1234567890abcdefghijklmnopqrstuvwxyz"}
|
||||
redacted, notice = redact_tool_output("file_read", output)
|
||||
assert "ghp_" not in redacted["content"]
|
||||
|
||||
def test_clean_output(self):
|
||||
output = "No credentials here"
|
||||
redacted, notice = redact_tool_output("file_read", output)
|
||||
assert redacted == output
|
||||
assert notice is None
|
||||
|
||||
|
||||
class TestRedactionResult:
|
||||
def test_notice_singular(self):
|
||||
result = RedactionResult("redacted", "original", [{"pattern_name": "test"}])
|
||||
assert "1 credential pattern" in result.notice()
|
||||
|
||||
def test_notice_plural(self):
|
||||
result = RedactionResult("redacted", "original", [
|
||||
{"pattern_name": "test1"},
|
||||
{"pattern_name": "test2"},
|
||||
])
|
||||
assert "2 credential patterns" in result.notice()
|
||||
|
||||
def test_to_dict(self):
|
||||
result = RedactionResult("redacted", "original", [{"pattern_name": "test"}])
|
||||
d = result.to_dict()
|
||||
assert d["redacted"] == True
|
||||
assert d["count"] == 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import pytest
|
||||
pytest.main([__file__, "-v"])
|
||||
280
tools/batch_executor.py
Normal file
280
tools/batch_executor.py
Normal file
@@ -0,0 +1,280 @@
|
||||
"""Batch tool execution with parallel safety checks.
|
||||
|
||||
Classifies tool calls as parallel-safe vs sequential and executes
|
||||
parallel-safe calls concurrently while keeping destructive ops serialized.
|
||||
|
||||
Issue #749: feat: batch tool execution with parallel safety checks
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ToolSafety(Enum):
|
||||
"""Safety classification for tool calls."""
|
||||
PARALLEL_SAFE = "parallel_safe" # Can run concurrently
|
||||
SEQUENTIAL = "sequential" # Must run one at a time
|
||||
DESTRUCTIVE = "destructive" # Destructive, needs approval
|
||||
|
||||
|
||||
# Tool safety classifications
|
||||
_TOOL_SAFETY: Dict[str, ToolSafety] = {
|
||||
# Parallel-safe: reads, searches, non-destructive
|
||||
"file_read": ToolSafety.PARALLEL_SAFE,
|
||||
"file_search": ToolSafety.PARALLEL_SAFE,
|
||||
"web_search": ToolSafety.PARALLEL_SAFE,
|
||||
"web_extract": ToolSafety.PARALLEL_SAFE,
|
||||
"browser_snapshot": ToolSafety.PARALLEL_SAFE,
|
||||
"browser_vision": ToolSafety.PARALLEL_SAFE,
|
||||
"browser_get_images": ToolSafety.PARALLEL_SAFE,
|
||||
"skill_view": ToolSafety.PARALLEL_SAFE,
|
||||
"memory_search": ToolSafety.PARALLEL_SAFE,
|
||||
"memory_recall": ToolSafety.PARALLEL_SAFE,
|
||||
"session_search": ToolSafety.PARALLEL_SAFE,
|
||||
|
||||
# Sequential: writes, edits, state changes
|
||||
"file_write": ToolSafety.SEQUENTIAL,
|
||||
"file_patch": ToolSafety.SEQUENTIAL,
|
||||
"file_append": ToolSafety.SEQUENTIAL,
|
||||
"browser_navigate": ToolSafety.SEQUENTIAL,
|
||||
"browser_click": ToolSafety.SEQUENTIAL,
|
||||
"browser_type": ToolSafety.SEQUENTIAL,
|
||||
"browser_scroll": ToolSafety.SEQUENTIAL,
|
||||
"memory_store": ToolSafety.SEQUENTIAL,
|
||||
"memory_update": ToolSafety.SEQUENTIAL,
|
||||
"cronjob": ToolSafety.SEQUENTIAL,
|
||||
"send_message": ToolSafety.SEQUENTIAL,
|
||||
|
||||
# Destructive: needs approval
|
||||
"terminal": ToolSafety.DESTRUCTIVE,
|
||||
"execute_code": ToolSafety.DESTRUCTIVE,
|
||||
"browser_execute_js": ToolSafety.DESTRUCTIVE,
|
||||
"delegate_task": ToolSafety.DESTRUCTIVE,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCall:
|
||||
"""A single tool call with metadata."""
|
||||
name: str
|
||||
args: Dict[str, Any]
|
||||
call_id: str = ""
|
||||
safety: ToolSafety = ToolSafety.SEQUENTIAL
|
||||
result: Optional[Any] = None
|
||||
error: Optional[str] = None
|
||||
duration: float = 0.0
|
||||
started_at: float = 0.0
|
||||
completed_at: float = 0.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchResult:
|
||||
"""Result of batch tool execution."""
|
||||
calls: List[ToolCall] = field(default_factory=list)
|
||||
parallel_count: int = 0
|
||||
sequential_count: int = 0
|
||||
total_duration: float = 0.0
|
||||
errors: List[str] = field(default_factory=list)
|
||||
|
||||
|
||||
def classify_tool_safety(tool_name: str) -> ToolSafety:
|
||||
"""Classify a tool call's safety level."""
|
||||
# Check exact match first
|
||||
if tool_name in _TOOL_SAFETY:
|
||||
return _TOOL_SAFETY[tool_name]
|
||||
|
||||
# Check prefix matches
|
||||
for pattern, safety in _TOOL_SAFETY.items():
|
||||
if tool_name.startswith(pattern):
|
||||
return safety
|
||||
|
||||
# Default to sequential for unknown tools
|
||||
return ToolSafety.SEQUENTIAL
|
||||
|
||||
|
||||
def classify_calls(tool_calls: List[Dict[str, Any]]) -> List[ToolCall]:
|
||||
"""Classify a list of tool calls by safety level."""
|
||||
calls = []
|
||||
for i, tc in enumerate(tool_calls):
|
||||
name = tc.get("name", tc.get("function", {}).get("name", ""))
|
||||
args = tc.get("arguments", tc.get("function", {}).get("arguments", {}))
|
||||
if isinstance(args, str):
|
||||
import json
|
||||
try:
|
||||
args = json.loads(args)
|
||||
except Exception:
|
||||
args = {}
|
||||
|
||||
call_id = tc.get("id", f"call_{i}")
|
||||
safety = classify_tool_safety(name)
|
||||
|
||||
calls.append(ToolCall(
|
||||
name=name,
|
||||
args=args,
|
||||
call_id=call_id,
|
||||
safety=safety,
|
||||
))
|
||||
|
||||
return calls
|
||||
|
||||
|
||||
async def execute_parallel(
|
||||
calls: List[ToolCall],
|
||||
executor: Callable[[str, Dict[str, Any]], Any],
|
||||
) -> List[ToolCall]:
|
||||
"""Execute parallel-safe calls concurrently."""
|
||||
async def run_call(call: ToolCall) -> ToolCall:
|
||||
call.started_at = time.time()
|
||||
try:
|
||||
# Run in thread pool to avoid blocking
|
||||
loop = asyncio.get_event_loop()
|
||||
result = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: executor(call.name, call.args),
|
||||
)
|
||||
call.result = result
|
||||
except Exception as e:
|
||||
call.error = str(e)
|
||||
logger.error(f"Parallel call {call.name} failed: {e}")
|
||||
finally:
|
||||
call.completed_at = time.time()
|
||||
call.duration = call.completed_at - call.started_at
|
||||
return call
|
||||
|
||||
# Execute all parallel-safe calls concurrently
|
||||
tasks = [run_call(call) for call in calls]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
# Handle exceptions from gather
|
||||
processed = []
|
||||
for i, result in enumerate(results):
|
||||
if isinstance(result, Exception):
|
||||
calls[i].error = str(result)
|
||||
calls[i].completed_at = time.time()
|
||||
calls[i].duration = calls[i].completed_at - calls[i].started_at
|
||||
processed.append(calls[i])
|
||||
else:
|
||||
processed.append(result)
|
||||
|
||||
return processed
|
||||
|
||||
|
||||
async def execute_sequential(
|
||||
calls: List[ToolCall],
|
||||
executor: Callable[[str, Dict[str, Any]], Any],
|
||||
) -> List[ToolCall]:
|
||||
"""Execute sequential/destructive calls one at a time."""
|
||||
for call in calls:
|
||||
call.started_at = time.time()
|
||||
try:
|
||||
result = executor(call.name, call.args)
|
||||
call.result = result
|
||||
except Exception as e:
|
||||
call.error = str(e)
|
||||
logger.error(f"Sequential call {call.name} failed: {e}")
|
||||
finally:
|
||||
call.completed_at = time.time()
|
||||
call.duration = call.completed_at - call.started_at
|
||||
|
||||
return calls
|
||||
|
||||
|
||||
async def execute_batch(
|
||||
tool_calls: List[Dict[str, Any]],
|
||||
executor: Callable[[str, Dict[str, Any]], Any],
|
||||
max_parallel: int = 5,
|
||||
) -> BatchResult:
|
||||
"""Execute a batch of tool calls with parallel safety checks.
|
||||
|
||||
Args:
|
||||
tool_calls: List of tool call dicts (OpenAI format)
|
||||
executor: Function to execute a single tool call (name, args) -> result
|
||||
max_parallel: Maximum concurrent parallel calls
|
||||
|
||||
Returns:
|
||||
BatchResult with all call results and timing info
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Classify all calls
|
||||
calls = classify_calls(tool_calls)
|
||||
|
||||
# Split by safety level
|
||||
parallel_calls = [c for c in calls if c.safety == ToolSafety.PARALLEL_SAFE]
|
||||
sequential_calls = [c for c in calls if c.safety != ToolSafety.PARALLEL_SAFE]
|
||||
|
||||
result = BatchResult(
|
||||
calls=calls,
|
||||
parallel_count=len(parallel_calls),
|
||||
sequential_count=len(sequential_calls),
|
||||
)
|
||||
|
||||
# Execute parallel calls concurrently
|
||||
if parallel_calls:
|
||||
logger.info(f"Executing {len(parallel_calls)} parallel-safe calls concurrently")
|
||||
|
||||
# Batch into chunks of max_parallel
|
||||
for i in range(0, len(parallel_calls), max_parallel):
|
||||
chunk = parallel_calls[i:i + max_parallel]
|
||||
await execute_parallel(chunk, executor)
|
||||
|
||||
# Execute sequential calls one at a time
|
||||
if sequential_calls:
|
||||
logger.info(f"Executing {len(sequential_calls)} sequential calls")
|
||||
await execute_sequential(sequential_calls, executor)
|
||||
|
||||
# Collect errors
|
||||
for call in calls:
|
||||
if call.error:
|
||||
result.errors.append(f"{call.name}: {call.error}")
|
||||
|
||||
result.total_duration = time.time() - start_time
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def execute_batch_sync(
|
||||
tool_calls: List[Dict[str, Any]],
|
||||
executor: Callable[[str, Dict[str, Any]], Any],
|
||||
max_parallel: int = 5,
|
||||
) -> BatchResult:
|
||||
"""Synchronous wrapper for execute_batch."""
|
||||
return asyncio.run(execute_batch(tool_calls, executor, max_parallel))
|
||||
|
||||
|
||||
def get_tool_safety_report(calls: List[ToolCall]) -> str:
|
||||
"""Generate a human-readable safety report."""
|
||||
parallel = [c for c in calls if c.safety == ToolSafety.PARALLEL_SAFE]
|
||||
sequential = [c for c in calls if c.safety == ToolSafety.SEQUENTIAL]
|
||||
destructive = [c for c in calls if c.safety == ToolSafety.DESTRUCTIVE]
|
||||
|
||||
lines = ["Tool Safety Report:"]
|
||||
lines.append(f" Parallel-safe: {len(parallel)}")
|
||||
lines.append(f" Sequential: {len(sequential)}")
|
||||
lines.append(f" Destructive: {len(destructive)}")
|
||||
|
||||
if parallel:
|
||||
lines.append("\nParallel-safe calls:")
|
||||
for c in parallel:
|
||||
status = "✓" if not c.error else "✗"
|
||||
lines.append(f" {status} {c.name} ({c.duration:.2f}s)")
|
||||
|
||||
if sequential:
|
||||
lines.append("\nSequential calls:")
|
||||
for c in sequential:
|
||||
status = "✓" if not c.error else "✗"
|
||||
lines.append(f" {status} {c.name} ({c.duration:.2f}s)")
|
||||
|
||||
if destructive:
|
||||
lines.append("\nDestructive calls:")
|
||||
for c in destructive:
|
||||
status = "✓" if not c.error else "✗"
|
||||
lines.append(f" {status} {c.name} ({c.duration:.2f}s)")
|
||||
|
||||
return "\n".join(lines)
|
||||
@@ -1,269 +0,0 @@
|
||||
"""Credential Redaction — Poka-yoke for tool outputs.
|
||||
|
||||
Blocks silent credential exposure by redacting API keys, tokens, and
|
||||
passwords from tool outputs before they enter agent context.
|
||||
|
||||
Issue #839: Poka-yoke: Block silent credential exposure in tool outputs
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Audit log path
|
||||
_AUDIT_DIR = Path.home() / ".hermes" / "audit"
|
||||
_AUDIT_LOG = _AUDIT_DIR / "redactions.jsonl"
|
||||
|
||||
# Credential patterns — order matters (most specific first)
|
||||
_CREDENTIAL_PATTERNS = [
|
||||
# API keys
|
||||
(r'sk-[a-zA-Z0-9]{20,}', '[REDACTED: OpenAI-style API key]'),
|
||||
(r'sk-ant-[a-zA-Z0-9-]{20,}', '[REDACTED: Anthropic API key]'),
|
||||
(r'ghp_[a-zA-Z0-9]{36}', '[REDACTED: GitHub PAT]'),
|
||||
(r'gho_[a-zA-Z0-9]{36}', '[REDACTED: GitHub OAuth token]'),
|
||||
(r'github_pat_[a-zA-Z0-9_]{82}', '[REDACTED: GitHub fine-grained PAT]'),
|
||||
(r'glpat-[a-zA-Z0-9-]{20,}', '[REDACTED: GitLab PAT]'),
|
||||
(r'syt_[a-zA-Z0-9_-]{40,}', '[REDACTED: Matrix access token]'),
|
||||
(r'xoxb-[0-9]{10,}-[a-zA-Z0-9]{20,}', '[REDACTED: Slack bot token]'),
|
||||
(r'xoxp-[0-9]{10,}-[a-zA-Z0-9]{20,}', '[REDACTED: Slack user token]'),
|
||||
|
||||
# Bearer tokens
|
||||
(r'Bearer\s+[a-zA-Z0-9_.-]{20,}', '[REDACTED: Bearer token]'),
|
||||
|
||||
# Generic tokens/passwords in assignments
|
||||
(r'(?:token|api_key|api_key|secret|password|passwd|pwd)\s*[:=]\s*["\']?([a-zA-Z0-9_.-]{8,})["\']?', '[REDACTED: credential]'),
|
||||
|
||||
# Environment variable assignments
|
||||
(r'(?:export\s+)?(?:TOKEN|KEY|SECRET|PASSWORD|API_KEY)\s*=\s*["\']?([a-zA-Z0-9_.-]{8,})["\']?', '[REDACTED: env credential]'),
|
||||
|
||||
# Base64 encoded credentials (high entropy strings)
|
||||
(r'(?:authorization|auth)\s*[:=]\s*(?:basic|bearer)\s+[a-zA-Z0-9+/=]{20,}', '[REDACTED: auth header]'),
|
||||
|
||||
# AWS credentials
|
||||
(r'AKIA[0-9A-Z]{16}', '[REDACTED: AWS access key]'),
|
||||
(r'(?<![A-Z0-9])[A-Za-z0-9/+=]{40}(?![A-Z0-9])', None), # Only match near context
|
||||
|
||||
# Private keys
|
||||
(r'-----BEGIN (?:RSA |EC |OPENSSH )?PRIVATE KEY-----', '[REDACTED: private key block]'),
|
||||
]
|
||||
|
||||
|
||||
class RedactionResult:
|
||||
"""Result of credential redaction."""
|
||||
|
||||
def __init__(self, text: str, original: str, redactions: List[Dict[str, Any]]):
|
||||
self.text = text
|
||||
self.original = original
|
||||
self.redactions = redactions
|
||||
|
||||
@property
|
||||
def was_redacted(self) -> bool:
|
||||
return len(self.redactions) > 0
|
||||
|
||||
@property
|
||||
def count(self) -> int:
|
||||
return len(self.redactions)
|
||||
|
||||
def notice(self) -> str:
|
||||
"""Generate compact redaction notice."""
|
||||
if not self.was_redacted:
|
||||
return ""
|
||||
return f"[REDACTED: {self.count} credential pattern{'s' if self.count > 1 else ''} found]"
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"redacted": self.was_redacted,
|
||||
"count": self.count,
|
||||
"notice": self.notice(),
|
||||
"patterns": [r["pattern_name"] for r in self.redactions],
|
||||
}
|
||||
|
||||
|
||||
def redact_credentials(text: str, source: str = "unknown") -> RedactionResult:
|
||||
"""Redact credentials from text.
|
||||
|
||||
Args:
|
||||
text: Text to redact
|
||||
source: Source identifier for audit logging
|
||||
|
||||
Returns:
|
||||
RedactionResult with redacted text and metadata
|
||||
"""
|
||||
if not text:
|
||||
return RedactionResult(text, text, [])
|
||||
|
||||
redactions = []
|
||||
result = text
|
||||
|
||||
for pattern, replacement in _CREDENTIAL_PATTERNS:
|
||||
if replacement is None:
|
||||
continue # Skip conditional patterns
|
||||
|
||||
matches = list(re.finditer(pattern, result, re.IGNORECASE))
|
||||
for match in matches:
|
||||
redactions.append({
|
||||
"pattern_name": replacement,
|
||||
"position": match.start(),
|
||||
"length": len(match.group()),
|
||||
"source": source,
|
||||
"timestamp": time.time(),
|
||||
})
|
||||
|
||||
result = re.sub(pattern, replacement, result, flags=re.IGNORECASE)
|
||||
|
||||
redaction_result = RedactionResult(result, text, redactions)
|
||||
|
||||
# Log to audit trail
|
||||
if redaction_result.was_redacted:
|
||||
_log_redaction(redaction_result, source)
|
||||
|
||||
return redaction_result
|
||||
|
||||
|
||||
def _log_redaction(result: RedactionResult, source: str) -> None:
|
||||
"""Log redaction event to audit trail."""
|
||||
try:
|
||||
_AUDIT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
entry = {
|
||||
"timestamp": time.time(),
|
||||
"source": source,
|
||||
"count": result.count,
|
||||
"patterns": [r["pattern_name"] for r in result.redactions],
|
||||
}
|
||||
with open(_AUDIT_LOG, "a") as f:
|
||||
f.write(json.dumps(entry) + "\n")
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to log redaction: {e}")
|
||||
|
||||
|
||||
def should_auto_mask(file_path: str) -> bool:
|
||||
"""Check if file should have credentials auto-masked."""
|
||||
path_lower = file_path.lower()
|
||||
sensitive_patterns = [
|
||||
".env", "config", "token", "secret", "credential",
|
||||
"key", "auth", "password", ".pem", ".key",
|
||||
]
|
||||
return any(p in path_lower for p in sensitive_patterns)
|
||||
|
||||
|
||||
def mask_config_values(text: str) -> str:
|
||||
"""Mask credential values in config/env files while preserving structure.
|
||||
|
||||
Transforms:
|
||||
API_KEY=sk-abc123 → API_KEY=[REDACTED]
|
||||
token: "ghp_xyz" → token: "[REDACTED]"
|
||||
"""
|
||||
lines = text.split("\n")
|
||||
result = []
|
||||
|
||||
for line in lines:
|
||||
# Match KEY=VALUE patterns
|
||||
match = re.match(r'^(\s*(?:export\s+)?[A-Z_][A-Z0-9_]*)\s*=\s*(.*)', line)
|
||||
if match:
|
||||
key = match.group(1)
|
||||
value = match.group(2).strip()
|
||||
|
||||
# Check if key looks credential-like
|
||||
key_lower = key.lower()
|
||||
if any(p in key_lower for p in ["key", "token", "secret", "password", "auth"]):
|
||||
if value and not value.startswith("[REDACTED]"):
|
||||
# Preserve quotes
|
||||
if value.startswith('"') and value.endswith('"'):
|
||||
result.append(f'{key}="[REDACTED]"')
|
||||
elif value.startswith("'") and value.endswith("'"):
|
||||
result.append(f"{key}='[REDACTED]'")
|
||||
else:
|
||||
result.append(f"{key}=[REDACTED]")
|
||||
continue
|
||||
|
||||
# Match YAML-style key: value
|
||||
match = re.match(r'^(\s*[a-z_][a-z0-9_]*)\s*:\s*["\']?(.*?)["\']?\s*$', line)
|
||||
if match:
|
||||
key = match.group(1)
|
||||
value = match.group(2).strip()
|
||||
|
||||
key_lower = key.lower()
|
||||
if any(p in key_lower for p in ["key", "token", "secret", "password", "auth"]):
|
||||
if value and not value.startswith("[REDACTED]"):
|
||||
result.append(f'{key}: "[REDACTED]"')
|
||||
continue
|
||||
|
||||
result.append(line)
|
||||
|
||||
return "\n".join(result)
|
||||
|
||||
|
||||
def redact_tool_output(
|
||||
tool_name: str,
|
||||
output: Any,
|
||||
source: str = None,
|
||||
) -> Tuple[Any, Optional[str]]:
|
||||
"""Redact credentials from tool output.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool
|
||||
output: Tool output (string or dict)
|
||||
source: Source identifier (defaults to tool_name)
|
||||
|
||||
Returns:
|
||||
Tuple of (redacted_output, notice)
|
||||
"""
|
||||
source = source or tool_name
|
||||
|
||||
if isinstance(output, str):
|
||||
result = redact_credentials(output, source)
|
||||
if result.was_redacted:
|
||||
return result.text, result.notice()
|
||||
return output, None
|
||||
|
||||
if isinstance(output, dict):
|
||||
# Redact string values in dict
|
||||
redacted = {}
|
||||
notices = []
|
||||
for key, value in output.items():
|
||||
if isinstance(value, str):
|
||||
r, n = redact_tool_output(tool_name, value, f"{source}.{key}")
|
||||
redacted[key] = r
|
||||
if n:
|
||||
notices.append(n)
|
||||
else:
|
||||
redacted[key] = value
|
||||
|
||||
notice = "; ".join(notices) if notices else None
|
||||
return redacted, notice
|
||||
|
||||
# Non-string, non-dict: pass through
|
||||
return output, None
|
||||
|
||||
|
||||
def get_redaction_stats() -> Dict[str, Any]:
|
||||
"""Get redaction statistics from audit log."""
|
||||
stats = {
|
||||
"total_redactions": 0,
|
||||
"by_source": {},
|
||||
"by_pattern": {},
|
||||
}
|
||||
|
||||
if not _AUDIT_LOG.exists():
|
||||
return stats
|
||||
|
||||
try:
|
||||
with open(_AUDIT_LOG, "r") as f:
|
||||
for line in f:
|
||||
entry = json.loads(line.strip())
|
||||
stats["total_redactions"] += entry.get("count", 0)
|
||||
|
||||
source = entry.get("source", "unknown")
|
||||
stats["by_source"][source] = stats["by_source"].get(source, 0) + 1
|
||||
|
||||
for pattern in entry.get("patterns", []):
|
||||
stats["by_pattern"][pattern] = stats["by_pattern"].get(pattern, 0) + 1
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return stats
|
||||
Reference in New Issue
Block a user