Initial commit: Agent Governance System Phase 8
Phase 8 Production Hardening with complete governance infrastructure: - Vault integration with tiered policies (T0-T4) - DragonflyDB state management - SQLite audit ledger - Pipeline DSL and templates - Promotion/revocation engine - Checkpoint system for session persistence - Health manager and circuit breaker for fault tolerance - GitHub/Slack integrations - Architectural test pipeline with bug watcher, suggestion engine, council review - Multi-agent chaos testing framework Test Results: - Governance tests: 68/68 passing - E2E workflow: 16/16 passing - Phase 2 Vault: 14/14 passing - Integration tests: 27/27 passing Coverage: 57.6% average across 12 phases Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,36 @@
|
||||
# Mocks
|
||||
|
||||
> Mock implementations for testing
|
||||
|
||||
## Overview
|
||||
|
||||
This directory contains mock implementations for testing.
|
||||
|
||||
## Key Files
|
||||
|
||||
| File | Description |
|
||||
|------|-------------|
|
||||
| `mock_dragonfly.py` | Python module |
|
||||
| `mock_llm.py` | Python module |
|
||||
| `mock_vault.py` | Python module |
|
||||
| `mock_blackboard.py` | Python module |
|
||||
| `__init__.py` | Python module |
|
||||
|
||||
## Interfaces / APIs
|
||||
|
||||
*Document any APIs, CLI commands, or interfaces here.*
|
||||
|
||||
## Status
|
||||
|
||||
** Not Started**
|
||||
|
||||
See [STATUS.md](./STATUS.md) for detailed progress tracking.
|
||||
|
||||
## Architecture Reference
|
||||
|
||||
Part of the [Agent Governance System](/opt/agent-governance/docs/ARCHITECTURE.md).
|
||||
|
||||
Parent: [Tests](..)
|
||||
|
||||
---
|
||||
*Last updated: 2026-01-23 23:25:09 UTC*
|
||||
@@ -0,0 +1,30 @@
|
||||
# Status: Mocks
|
||||
|
||||
## Current Phase
|
||||
|
||||
** NOT STARTED**
|
||||
|
||||
## Tasks
|
||||
|
||||
| Status | Task | Updated |
|
||||
|--------|------|---------|
|
||||
| ☐ | *No tasks defined* | - |
|
||||
|
||||
## Dependencies
|
||||
|
||||
*No external dependencies.*
|
||||
|
||||
## Issues / Blockers
|
||||
|
||||
*No current issues or blockers.*
|
||||
|
||||
## Activity Log
|
||||
|
||||
### 2026-01-23 23:25:09 UTC
|
||||
- **Phase**: NOT STARTED
|
||||
- **Action**: Initialized
|
||||
- **Details**: Status tracking initialized for this directory.
|
||||
|
||||
|
||||
---
|
||||
*Last updated: 2026-01-23 23:25:09 UTC*
|
||||
@@ -0,0 +1,21 @@
|
||||
"""
|
||||
Mock infrastructure for agent governance testing.
|
||||
|
||||
Provides deterministic mock implementations of:
|
||||
- Vault (policy engine)
|
||||
- DragonflyDB (runtime state)
|
||||
- LLM (language model responses)
|
||||
- Blackboard (shared memory)
|
||||
"""
|
||||
|
||||
from .mock_vault import MockVault
|
||||
from .mock_dragonfly import MockDragonfly
|
||||
from .mock_llm import MockLLM
|
||||
from .mock_blackboard import MockBlackboard
|
||||
|
||||
__all__ = [
|
||||
'MockVault',
|
||||
'MockDragonfly',
|
||||
'MockLLM',
|
||||
'MockBlackboard'
|
||||
]
|
||||
@@ -0,0 +1,394 @@
|
||||
"""
|
||||
MockBlackboard - Simulates shared memory for multi-agent coordination testing.
|
||||
|
||||
Provides deterministic blackboard operations for testing multi-agent
|
||||
scenarios without DragonflyDB.
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Any, Optional, Set
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
import json
|
||||
import threading
|
||||
|
||||
|
||||
class BlackboardSection(Enum):
|
||||
"""Standard blackboard sections"""
|
||||
PROBLEM = "problem"
|
||||
SOLUTIONS = "solutions"
|
||||
PROGRESS = "progress"
|
||||
CONSENSUS = "consensus"
|
||||
|
||||
|
||||
@dataclass
|
||||
class BlackboardEntry:
|
||||
"""A single entry in a blackboard section"""
|
||||
key: str
|
||||
value: Any
|
||||
author: str
|
||||
timestamp: datetime
|
||||
version: int = 1
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConsensusVote:
|
||||
"""A vote in the consensus section"""
|
||||
agent: str
|
||||
proposal_id: str
|
||||
vote: str # ACCEPT, REJECT, ABSTAIN
|
||||
reasoning: str
|
||||
timestamp: datetime
|
||||
|
||||
|
||||
class MockBlackboard:
|
||||
"""
|
||||
Mock Blackboard implementation for multi-agent testing.
|
||||
|
||||
Simulates:
|
||||
- Section-based shared memory
|
||||
- Read/write operations
|
||||
- Consensus voting
|
||||
- History tracking
|
||||
- Conflict detection
|
||||
"""
|
||||
|
||||
def __init__(self, task_id: str = "test-task"):
|
||||
self.task_id = task_id
|
||||
self._sections: Dict[str, Dict[str, BlackboardEntry]] = {
|
||||
section.value: {} for section in BlackboardSection
|
||||
}
|
||||
self._votes: List[ConsensusVote] = []
|
||||
self._history: List[Dict[str, Any]] = []
|
||||
self._watchers: Dict[str, List[callable]] = {}
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def write(self, section: str, key: str, value: Any, author: str) -> bool:
|
||||
"""
|
||||
Write to a blackboard section.
|
||||
|
||||
Args:
|
||||
section: Section name (problem, solutions, progress, consensus)
|
||||
key: Entry key
|
||||
value: Entry value
|
||||
author: Agent ID writing the entry
|
||||
"""
|
||||
with self._lock:
|
||||
if section not in self._sections:
|
||||
self._sections[section] = {}
|
||||
|
||||
existing = self._sections[section].get(key)
|
||||
version = existing.version + 1 if existing else 1
|
||||
|
||||
entry = BlackboardEntry(
|
||||
key=key,
|
||||
value=value,
|
||||
author=author,
|
||||
timestamp=datetime.utcnow(),
|
||||
version=version
|
||||
)
|
||||
self._sections[section][key] = entry
|
||||
|
||||
# Record in history
|
||||
self._history.append({
|
||||
"operation": "write",
|
||||
"section": section,
|
||||
"key": key,
|
||||
"value": value,
|
||||
"author": author,
|
||||
"version": version,
|
||||
"timestamp": datetime.utcnow().isoformat()
|
||||
})
|
||||
|
||||
# Notify watchers
|
||||
self._notify_watchers(section, key, value, author)
|
||||
|
||||
return True
|
||||
|
||||
def read(self, section: str, key: str = None) -> Optional[Any]:
|
||||
"""
|
||||
Read from a blackboard section.
|
||||
|
||||
Args:
|
||||
section: Section name
|
||||
key: Optional specific key (if None, returns all entries)
|
||||
"""
|
||||
with self._lock:
|
||||
if section not in self._sections:
|
||||
return None
|
||||
|
||||
if key:
|
||||
entry = self._sections[section].get(key)
|
||||
return entry.value if entry else None
|
||||
else:
|
||||
return {
|
||||
k: v.value
|
||||
for k, v in self._sections[section].items()
|
||||
}
|
||||
|
||||
def read_with_metadata(self, section: str, key: str) -> Optional[Dict[str, Any]]:
|
||||
"""Read entry with full metadata"""
|
||||
with self._lock:
|
||||
if section not in self._sections:
|
||||
return None
|
||||
|
||||
entry = self._sections[section].get(key)
|
||||
if not entry:
|
||||
return None
|
||||
|
||||
return {
|
||||
"key": entry.key,
|
||||
"value": entry.value,
|
||||
"author": entry.author,
|
||||
"version": entry.version,
|
||||
"timestamp": entry.timestamp.isoformat()
|
||||
}
|
||||
|
||||
def delete(self, section: str, key: str, author: str) -> bool:
|
||||
"""Delete an entry from a section"""
|
||||
with self._lock:
|
||||
if section not in self._sections:
|
||||
return False
|
||||
|
||||
if key not in self._sections[section]:
|
||||
return False
|
||||
|
||||
del self._sections[section][key]
|
||||
|
||||
self._history.append({
|
||||
"operation": "delete",
|
||||
"section": section,
|
||||
"key": key,
|
||||
"author": author,
|
||||
"timestamp": datetime.utcnow().isoformat()
|
||||
})
|
||||
|
||||
return True
|
||||
|
||||
def list_keys(self, section: str) -> List[str]:
|
||||
"""List all keys in a section"""
|
||||
with self._lock:
|
||||
if section not in self._sections:
|
||||
return []
|
||||
return list(self._sections[section].keys())
|
||||
|
||||
# === Consensus Operations ===
|
||||
|
||||
def submit_proposal(self, proposal_id: str, proposal: Any, author: str) -> bool:
|
||||
"""Submit a proposal for consensus"""
|
||||
return self.write(
|
||||
BlackboardSection.SOLUTIONS.value,
|
||||
proposal_id,
|
||||
{"proposal": proposal, "status": "pending", "votes": []},
|
||||
author
|
||||
)
|
||||
|
||||
def vote(self, proposal_id: str, agent: str, vote: str, reasoning: str = "") -> bool:
|
||||
"""
|
||||
Vote on a proposal.
|
||||
|
||||
Args:
|
||||
proposal_id: ID of the proposal
|
||||
agent: Voting agent ID
|
||||
vote: ACCEPT, REJECT, or ABSTAIN
|
||||
reasoning: Optional explanation
|
||||
"""
|
||||
with self._lock:
|
||||
if vote not in ["ACCEPT", "REJECT", "ABSTAIN"]:
|
||||
return False
|
||||
|
||||
consensus_vote = ConsensusVote(
|
||||
agent=agent,
|
||||
proposal_id=proposal_id,
|
||||
vote=vote,
|
||||
reasoning=reasoning,
|
||||
timestamp=datetime.utcnow()
|
||||
)
|
||||
self._votes.append(consensus_vote)
|
||||
|
||||
# Update proposal with vote
|
||||
solutions = self._sections.get(BlackboardSection.SOLUTIONS.value, {})
|
||||
if proposal_id in solutions:
|
||||
proposal_entry = solutions[proposal_id]
|
||||
proposal_data = proposal_entry.value
|
||||
if "votes" not in proposal_data:
|
||||
proposal_data["votes"] = []
|
||||
proposal_data["votes"].append({
|
||||
"agent": agent,
|
||||
"vote": vote,
|
||||
"reasoning": reasoning
|
||||
})
|
||||
|
||||
self._history.append({
|
||||
"operation": "vote",
|
||||
"proposal_id": proposal_id,
|
||||
"agent": agent,
|
||||
"vote": vote,
|
||||
"timestamp": datetime.utcnow().isoformat()
|
||||
})
|
||||
|
||||
return True
|
||||
|
||||
def check_consensus(self, proposal_id: str, required_agents: List[str] = None) -> Dict[str, Any]:
|
||||
"""
|
||||
Check consensus status for a proposal.
|
||||
|
||||
Returns:
|
||||
{
|
||||
"reached": bool,
|
||||
"votes": {"ACCEPT": n, "REJECT": n, "ABSTAIN": n},
|
||||
"missing_votes": [agent_ids],
|
||||
"result": "ACCEPT"|"REJECT"|"PENDING"
|
||||
}
|
||||
"""
|
||||
with self._lock:
|
||||
proposal_votes = [v for v in self._votes if v.proposal_id == proposal_id]
|
||||
|
||||
vote_counts = {"ACCEPT": 0, "REJECT": 0, "ABSTAIN": 0}
|
||||
voted_agents = set()
|
||||
|
||||
for v in proposal_votes:
|
||||
vote_counts[v.vote] += 1
|
||||
voted_agents.add(v.agent)
|
||||
|
||||
missing = []
|
||||
if required_agents:
|
||||
missing = [a for a in required_agents if a not in voted_agents]
|
||||
|
||||
all_voted = len(missing) == 0 if required_agents else True
|
||||
|
||||
# Determine result
|
||||
if not all_voted:
|
||||
result = "PENDING"
|
||||
elif vote_counts["ACCEPT"] > vote_counts["REJECT"]:
|
||||
result = "ACCEPT"
|
||||
elif vote_counts["REJECT"] > vote_counts["ACCEPT"]:
|
||||
result = "REJECT"
|
||||
else:
|
||||
result = "TIE"
|
||||
|
||||
return {
|
||||
"reached": all_voted and result in ["ACCEPT", "REJECT"],
|
||||
"votes": vote_counts,
|
||||
"missing_votes": missing,
|
||||
"result": result
|
||||
}
|
||||
|
||||
# === Progress Tracking ===
|
||||
|
||||
def update_progress(self, agent: str, phase: str, step: str, details: Dict = None) -> bool:
|
||||
"""Update agent progress"""
|
||||
return self.write(
|
||||
BlackboardSection.PROGRESS.value,
|
||||
agent,
|
||||
{
|
||||
"phase": phase,
|
||||
"step": step,
|
||||
"details": details or {},
|
||||
"updated_at": datetime.utcnow().isoformat()
|
||||
},
|
||||
agent
|
||||
)
|
||||
|
||||
def get_all_progress(self) -> Dict[str, Any]:
|
||||
"""Get progress for all agents"""
|
||||
return self.read(BlackboardSection.PROGRESS.value) or {}
|
||||
|
||||
# === Watchers ===
|
||||
|
||||
def watch(self, section: str, callback: callable):
|
||||
"""Watch a section for changes"""
|
||||
with self._lock:
|
||||
if section not in self._watchers:
|
||||
self._watchers[section] = []
|
||||
self._watchers[section].append(callback)
|
||||
|
||||
def unwatch(self, section: str, callback: callable = None):
|
||||
"""Remove a watcher"""
|
||||
with self._lock:
|
||||
if section not in self._watchers:
|
||||
return
|
||||
if callback:
|
||||
self._watchers[section] = [
|
||||
cb for cb in self._watchers[section] if cb != callback
|
||||
]
|
||||
else:
|
||||
del self._watchers[section]
|
||||
|
||||
def _notify_watchers(self, section: str, key: str, value: Any, author: str):
|
||||
"""Notify watchers of a change (called with lock held)"""
|
||||
callbacks = self._watchers.get(section, []).copy()
|
||||
|
||||
# Release lock before calling callbacks
|
||||
for callback in callbacks:
|
||||
try:
|
||||
callback(section, key, value, author)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# === Test Helpers ===
|
||||
|
||||
def reset(self):
|
||||
"""Reset all state for testing"""
|
||||
with self._lock:
|
||||
self._sections = {
|
||||
section.value: {} for section in BlackboardSection
|
||||
}
|
||||
self._votes.clear()
|
||||
self._history.clear()
|
||||
self._watchers.clear()
|
||||
|
||||
def get_history(self) -> List[Dict[str, Any]]:
|
||||
"""Get operation history for assertions"""
|
||||
with self._lock:
|
||||
return self._history.copy()
|
||||
|
||||
def get_all_state(self) -> Dict[str, Any]:
|
||||
"""Get complete state for test assertions"""
|
||||
with self._lock:
|
||||
return {
|
||||
"sections": {
|
||||
section: {
|
||||
k: {"value": v.value, "author": v.author, "version": v.version}
|
||||
for k, v in entries.items()
|
||||
}
|
||||
for section, entries in self._sections.items()
|
||||
},
|
||||
"votes": [
|
||||
{"proposal": v.proposal_id, "agent": v.agent, "vote": v.vote}
|
||||
for v in self._votes
|
||||
]
|
||||
}
|
||||
|
||||
def inject_state(self, section: str, entries: Dict[str, Any], author: str = "test"):
|
||||
"""Inject state for testing"""
|
||||
with self._lock:
|
||||
if section not in self._sections:
|
||||
self._sections[section] = {}
|
||||
|
||||
for key, value in entries.items():
|
||||
self._sections[section][key] = BlackboardEntry(
|
||||
key=key,
|
||||
value=value,
|
||||
author=author,
|
||||
timestamp=datetime.utcnow()
|
||||
)
|
||||
|
||||
def simulate_conflict(self, section: str, key: str, agents: List[str],
|
||||
values: List[Any]) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Simulate a write conflict for testing.
|
||||
|
||||
Multiple agents write to the same key in rapid succession.
|
||||
Returns the sequence of writes.
|
||||
"""
|
||||
writes = []
|
||||
for agent, value in zip(agents, values):
|
||||
self.write(section, key, value, agent)
|
||||
entry = self._sections[section][key]
|
||||
writes.append({
|
||||
"agent": agent,
|
||||
"value": value,
|
||||
"version": entry.version
|
||||
})
|
||||
return writes
|
||||
@@ -0,0 +1,445 @@
|
||||
"""
|
||||
MockDragonfly - Simulates DragonflyDB/Redis for testing.
|
||||
|
||||
Provides deterministic state management, locks, and pub/sub
|
||||
without requiring a real Redis instance.
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Any, Optional, Callable, Set
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockLock:
|
||||
"""Represents a distributed lock"""
|
||||
key: str
|
||||
owner: str
|
||||
acquired_at: datetime
|
||||
ttl: int # seconds
|
||||
released: bool = False
|
||||
|
||||
def is_valid(self) -> bool:
|
||||
if self.released:
|
||||
return False
|
||||
if datetime.utcnow() > self.acquired_at + timedelta(seconds=self.ttl):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class MockDragonfly:
|
||||
"""
|
||||
Mock DragonflyDB/Redis implementation for testing.
|
||||
|
||||
Simulates:
|
||||
- Key-value storage (strings, hashes, lists, sets)
|
||||
- Key expiration (TTL)
|
||||
- Distributed locks
|
||||
- Pub/Sub messaging
|
||||
- Atomic operations
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._strings: Dict[str, Any] = {}
|
||||
self._hashes: Dict[str, Dict[str, Any]] = {}
|
||||
self._lists: Dict[str, List[Any]] = {}
|
||||
self._sets: Dict[str, Set[str]] = {}
|
||||
self._expiry: Dict[str, datetime] = {}
|
||||
self._locks: Dict[str, MockLock] = {}
|
||||
self._subscribers: Dict[str, List[Callable]] = {}
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def _check_expiry(self, key: str) -> bool:
|
||||
"""Check if key has expired. Returns True if key is valid."""
|
||||
if key in self._expiry:
|
||||
if datetime.utcnow() > self._expiry[key]:
|
||||
# Key expired, clean up
|
||||
self._delete_key(key)
|
||||
return False
|
||||
return True
|
||||
|
||||
def _delete_key(self, key: str):
|
||||
"""Delete key from all stores"""
|
||||
self._strings.pop(key, None)
|
||||
self._hashes.pop(key, None)
|
||||
self._lists.pop(key, None)
|
||||
self._sets.pop(key, None)
|
||||
self._expiry.pop(key, None)
|
||||
|
||||
# === String Operations ===
|
||||
|
||||
def set(self, key: str, value: Any, ex: int = None, nx: bool = False) -> bool:
|
||||
"""
|
||||
Set a string value.
|
||||
|
||||
Args:
|
||||
key: Key name
|
||||
value: Value to store
|
||||
ex: Expiry in seconds
|
||||
nx: Only set if key does not exist
|
||||
"""
|
||||
with self._lock:
|
||||
if nx and key in self._strings and self._check_expiry(key):
|
||||
return False
|
||||
|
||||
self._strings[key] = value
|
||||
if ex:
|
||||
self._expiry[key] = datetime.utcnow() + timedelta(seconds=ex)
|
||||
return True
|
||||
|
||||
def get(self, key: str) -> Optional[Any]:
|
||||
"""Get a string value"""
|
||||
with self._lock:
|
||||
if not self._check_expiry(key):
|
||||
return None
|
||||
return self._strings.get(key)
|
||||
|
||||
def incr(self, key: str, amount: int = 1) -> int:
|
||||
"""Increment a numeric value"""
|
||||
with self._lock:
|
||||
current = int(self._strings.get(key, 0))
|
||||
new_value = current + amount
|
||||
self._strings[key] = new_value
|
||||
return new_value
|
||||
|
||||
def decr(self, key: str, amount: int = 1) -> int:
|
||||
"""Decrement a numeric value"""
|
||||
return self.incr(key, -amount)
|
||||
|
||||
# === Hash Operations ===
|
||||
|
||||
def hset(self, key: str, field: str = None, value: Any = None,
|
||||
mapping: Dict[str, Any] = None) -> int:
|
||||
"""
|
||||
Set hash field(s).
|
||||
|
||||
Can be called as:
|
||||
- hset(key, field, value)
|
||||
- hset(key, mapping={...})
|
||||
"""
|
||||
with self._lock:
|
||||
if key not in self._hashes:
|
||||
self._hashes[key] = {}
|
||||
|
||||
count = 0
|
||||
if mapping:
|
||||
for f, v in mapping.items():
|
||||
if f not in self._hashes[key]:
|
||||
count += 1
|
||||
self._hashes[key][f] = v
|
||||
elif field is not None:
|
||||
if field not in self._hashes[key]:
|
||||
count = 1
|
||||
self._hashes[key][field] = value
|
||||
|
||||
return count
|
||||
|
||||
def hget(self, key: str, field: str) -> Optional[Any]:
|
||||
"""Get a hash field"""
|
||||
with self._lock:
|
||||
if not self._check_expiry(key):
|
||||
return None
|
||||
return self._hashes.get(key, {}).get(field)
|
||||
|
||||
def hgetall(self, key: str) -> Dict[str, Any]:
|
||||
"""Get all hash fields"""
|
||||
with self._lock:
|
||||
if not self._check_expiry(key):
|
||||
return {}
|
||||
return self._hashes.get(key, {}).copy()
|
||||
|
||||
def hincrby(self, key: str, field: str, amount: int = 1) -> int:
|
||||
"""Increment a hash field"""
|
||||
with self._lock:
|
||||
if key not in self._hashes:
|
||||
self._hashes[key] = {}
|
||||
current = int(self._hashes[key].get(field, 0))
|
||||
new_value = current + amount
|
||||
self._hashes[key][field] = new_value
|
||||
return new_value
|
||||
|
||||
def hdel(self, key: str, *fields: str) -> int:
|
||||
"""Delete hash fields"""
|
||||
with self._lock:
|
||||
if key not in self._hashes:
|
||||
return 0
|
||||
count = 0
|
||||
for field in fields:
|
||||
if field in self._hashes[key]:
|
||||
del self._hashes[key][field]
|
||||
count += 1
|
||||
return count
|
||||
|
||||
# === List Operations ===
|
||||
|
||||
def lpush(self, key: str, *values: Any) -> int:
|
||||
"""Push values to left of list"""
|
||||
with self._lock:
|
||||
if key not in self._lists:
|
||||
self._lists[key] = []
|
||||
for value in reversed(values):
|
||||
self._lists[key].insert(0, value)
|
||||
return len(self._lists[key])
|
||||
|
||||
def rpush(self, key: str, *values: Any) -> int:
|
||||
"""Push values to right of list"""
|
||||
with self._lock:
|
||||
if key not in self._lists:
|
||||
self._lists[key] = []
|
||||
self._lists[key].extend(values)
|
||||
return len(self._lists[key])
|
||||
|
||||
def lpop(self, key: str) -> Optional[Any]:
|
||||
"""Pop from left of list"""
|
||||
with self._lock:
|
||||
if key not in self._lists or not self._lists[key]:
|
||||
return None
|
||||
return self._lists[key].pop(0)
|
||||
|
||||
def rpop(self, key: str) -> Optional[Any]:
|
||||
"""Pop from right of list"""
|
||||
with self._lock:
|
||||
if key not in self._lists or not self._lists[key]:
|
||||
return None
|
||||
return self._lists[key].pop()
|
||||
|
||||
def lrange(self, key: str, start: int, stop: int) -> List[Any]:
|
||||
"""Get range of list elements"""
|
||||
with self._lock:
|
||||
if key not in self._lists:
|
||||
return []
|
||||
# Redis uses inclusive stop, Python uses exclusive
|
||||
if stop == -1:
|
||||
return self._lists[key][start:]
|
||||
return self._lists[key][start:stop + 1]
|
||||
|
||||
def llen(self, key: str) -> int:
|
||||
"""Get list length"""
|
||||
with self._lock:
|
||||
return len(self._lists.get(key, []))
|
||||
|
||||
def ltrim(self, key: str, start: int, stop: int) -> bool:
|
||||
"""Trim list to specified range"""
|
||||
with self._lock:
|
||||
if key not in self._lists:
|
||||
return True
|
||||
if stop == -1:
|
||||
self._lists[key] = self._lists[key][start:]
|
||||
else:
|
||||
self._lists[key] = self._lists[key][start:stop + 1]
|
||||
return True
|
||||
|
||||
# === Set Operations ===
|
||||
|
||||
def sadd(self, key: str, *members: str) -> int:
|
||||
"""Add members to set"""
|
||||
with self._lock:
|
||||
if key not in self._sets:
|
||||
self._sets[key] = set()
|
||||
before = len(self._sets[key])
|
||||
self._sets[key].update(members)
|
||||
return len(self._sets[key]) - before
|
||||
|
||||
def srem(self, key: str, *members: str) -> int:
|
||||
"""Remove members from set"""
|
||||
with self._lock:
|
||||
if key not in self._sets:
|
||||
return 0
|
||||
before = len(self._sets[key])
|
||||
self._sets[key] -= set(members)
|
||||
return before - len(self._sets[key])
|
||||
|
||||
def smembers(self, key: str) -> Set[str]:
|
||||
"""Get all set members"""
|
||||
with self._lock:
|
||||
return self._sets.get(key, set()).copy()
|
||||
|
||||
def sismember(self, key: str, member: str) -> bool:
|
||||
"""Check if member is in set"""
|
||||
with self._lock:
|
||||
return member in self._sets.get(key, set())
|
||||
|
||||
# === Key Operations ===
|
||||
|
||||
def exists(self, *keys: str) -> int:
|
||||
"""Check if keys exist"""
|
||||
with self._lock:
|
||||
count = 0
|
||||
for key in keys:
|
||||
if self._check_expiry(key):
|
||||
if key in self._strings or key in self._hashes or \
|
||||
key in self._lists or key in self._sets:
|
||||
count += 1
|
||||
return count
|
||||
|
||||
def delete(self, *keys: str) -> int:
|
||||
"""Delete keys"""
|
||||
with self._lock:
|
||||
count = 0
|
||||
for key in keys:
|
||||
if key in self._strings or key in self._hashes or \
|
||||
key in self._lists or key in self._sets:
|
||||
self._delete_key(key)
|
||||
count += 1
|
||||
return count
|
||||
|
||||
def expire(self, key: str, seconds: int) -> bool:
|
||||
"""Set key expiration"""
|
||||
with self._lock:
|
||||
if key in self._strings or key in self._hashes or \
|
||||
key in self._lists or key in self._sets:
|
||||
self._expiry[key] = datetime.utcnow() + timedelta(seconds=seconds)
|
||||
return True
|
||||
return False
|
||||
|
||||
def ttl(self, key: str) -> int:
|
||||
"""Get TTL in seconds (-1 if no expiry, -2 if key doesn't exist)"""
|
||||
with self._lock:
|
||||
exists = key in self._strings or key in self._hashes or \
|
||||
key in self._lists or key in self._sets
|
||||
if not exists:
|
||||
return -2
|
||||
if key not in self._expiry:
|
||||
return -1
|
||||
remaining = (self._expiry[key] - datetime.utcnow()).total_seconds()
|
||||
return max(0, int(remaining))
|
||||
|
||||
def keys(self, pattern: str = "*") -> List[str]:
|
||||
"""Get keys matching pattern"""
|
||||
with self._lock:
|
||||
all_keys = set()
|
||||
all_keys.update(self._strings.keys())
|
||||
all_keys.update(self._hashes.keys())
|
||||
all_keys.update(self._lists.keys())
|
||||
all_keys.update(self._sets.keys())
|
||||
|
||||
if pattern == "*":
|
||||
return list(all_keys)
|
||||
|
||||
# Simple pattern matching (only supports * and prefix*)
|
||||
import fnmatch
|
||||
return [k for k in all_keys if fnmatch.fnmatch(k, pattern)]
|
||||
|
||||
# === Distributed Locks ===
|
||||
|
||||
def acquire_lock(self, key: str, owner: str, ttl: int = 30) -> bool:
|
||||
"""
|
||||
Acquire a distributed lock.
|
||||
|
||||
Args:
|
||||
key: Lock key
|
||||
owner: Lock owner identifier
|
||||
ttl: Lock TTL in seconds
|
||||
"""
|
||||
with self._lock:
|
||||
if key in self._locks and self._locks[key].is_valid():
|
||||
if self._locks[key].owner != owner:
|
||||
return False
|
||||
|
||||
self._locks[key] = MockLock(
|
||||
key=key,
|
||||
owner=owner,
|
||||
acquired_at=datetime.utcnow(),
|
||||
ttl=ttl
|
||||
)
|
||||
return True
|
||||
|
||||
def release_lock(self, key: str, owner: str) -> bool:
|
||||
"""Release a distributed lock"""
|
||||
with self._lock:
|
||||
if key not in self._locks:
|
||||
return False
|
||||
if self._locks[key].owner != owner:
|
||||
return False
|
||||
self._locks[key].released = True
|
||||
return True
|
||||
|
||||
def refresh_lock(self, key: str, owner: str, ttl: int = 30) -> bool:
|
||||
"""Refresh lock TTL"""
|
||||
with self._lock:
|
||||
if key not in self._locks:
|
||||
return False
|
||||
lock = self._locks[key]
|
||||
if lock.owner != owner or not lock.is_valid():
|
||||
return False
|
||||
lock.acquired_at = datetime.utcnow()
|
||||
lock.ttl = ttl
|
||||
return True
|
||||
|
||||
# === Pub/Sub ===
|
||||
|
||||
def subscribe(self, channel: str, callback: Callable[[str, Any], None]):
|
||||
"""Subscribe to a channel"""
|
||||
with self._lock:
|
||||
if channel not in self._subscribers:
|
||||
self._subscribers[channel] = []
|
||||
self._subscribers[channel].append(callback)
|
||||
|
||||
def unsubscribe(self, channel: str, callback: Callable = None):
|
||||
"""Unsubscribe from a channel"""
|
||||
with self._lock:
|
||||
if channel not in self._subscribers:
|
||||
return
|
||||
if callback:
|
||||
self._subscribers[channel] = [
|
||||
cb for cb in self._subscribers[channel] if cb != callback
|
||||
]
|
||||
else:
|
||||
del self._subscribers[channel]
|
||||
|
||||
def publish(self, channel: str, message: Any) -> int:
|
||||
"""Publish message to channel"""
|
||||
with self._lock:
|
||||
callbacks = self._subscribers.get(channel, []).copy()
|
||||
|
||||
count = 0
|
||||
for callback in callbacks:
|
||||
try:
|
||||
callback(channel, message)
|
||||
count += 1
|
||||
except Exception:
|
||||
pass
|
||||
return count
|
||||
|
||||
# === Test Helpers ===
|
||||
|
||||
def reset(self):
|
||||
"""Reset all state for testing"""
|
||||
with self._lock:
|
||||
self._strings.clear()
|
||||
self._hashes.clear()
|
||||
self._lists.clear()
|
||||
self._sets.clear()
|
||||
self._expiry.clear()
|
||||
self._locks.clear()
|
||||
self._subscribers.clear()
|
||||
|
||||
def get_all_state(self) -> Dict[str, Any]:
|
||||
"""Get complete state for test assertions"""
|
||||
with self._lock:
|
||||
return {
|
||||
"strings": self._strings.copy(),
|
||||
"hashes": {k: v.copy() for k, v in self._hashes.items()},
|
||||
"lists": {k: v.copy() for k, v in self._lists.items()},
|
||||
"sets": {k: v.copy() for k, v in self._sets.items()},
|
||||
"locks": {k: {"owner": v.owner, "valid": v.is_valid()}
|
||||
for k, v in self._locks.items()}
|
||||
}
|
||||
|
||||
def inject_state(self, state: Dict[str, Any]):
|
||||
"""Inject state for testing"""
|
||||
with self._lock:
|
||||
if "strings" in state:
|
||||
self._strings.update(state["strings"])
|
||||
if "hashes" in state:
|
||||
for k, v in state["hashes"].items():
|
||||
self._hashes[k] = v.copy()
|
||||
if "lists" in state:
|
||||
for k, v in state["lists"].items():
|
||||
self._lists[k] = v.copy()
|
||||
if "sets" in state:
|
||||
for k, v in state["sets"].items():
|
||||
self._sets[k] = set(v)
|
||||
@@ -0,0 +1,308 @@
|
||||
"""
|
||||
MockLLM - Simulates LLM responses for deterministic testing.
|
||||
|
||||
Provides pattern-based response matching for testing agent behavior
|
||||
without requiring actual LLM API calls.
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Any, Optional, Callable, Tuple
|
||||
from dataclasses import dataclass, field
|
||||
import re
|
||||
import json
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockResponse:
|
||||
"""A configured mock response"""
|
||||
pattern: str
|
||||
response: str
|
||||
confidence: float = 0.9
|
||||
delay: float = 0.0 # Simulated latency
|
||||
tokens_used: int = 100
|
||||
call_count: int = 0
|
||||
|
||||
|
||||
class MockLLM:
|
||||
"""
|
||||
Mock LLM implementation for testing.
|
||||
|
||||
Features:
|
||||
- Pattern-based response matching
|
||||
- Configurable confidence levels
|
||||
- Token counting simulation
|
||||
- Call tracking for assertions
|
||||
- Error injection for failure testing
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._responses: List[MockResponse] = []
|
||||
self._default_response: str = '{"decision": "EXECUTE", "confidence": 0.5, "steps": []}'
|
||||
self._call_history: List[Dict[str, Any]] = []
|
||||
self._total_tokens: int = 0
|
||||
self._error_mode: Optional[str] = None
|
||||
self._error_after: int = 0
|
||||
self._call_count: int = 0
|
||||
|
||||
# Set up some default patterns
|
||||
self._setup_defaults()
|
||||
|
||||
def _setup_defaults(self):
|
||||
"""Set up default response patterns"""
|
||||
# Plan generation
|
||||
self.add_response(
|
||||
pattern="generate.*plan|create.*plan|plan.*for",
|
||||
response=json.dumps({
|
||||
"decision": "PLAN",
|
||||
"confidence": 0.85,
|
||||
"plan": {
|
||||
"title": "Generated Plan",
|
||||
"steps": [
|
||||
{"action": "analyze", "description": "Analyze requirements"},
|
||||
{"action": "implement", "description": "Implement solution"},
|
||||
{"action": "verify", "description": "Verify results"}
|
||||
]
|
||||
},
|
||||
"assumptions": [],
|
||||
"risks": []
|
||||
}),
|
||||
confidence=0.85
|
||||
)
|
||||
|
||||
# Read operations
|
||||
self.add_response(
|
||||
pattern="read|view|show|list|get",
|
||||
response=json.dumps({
|
||||
"decision": "EXECUTE",
|
||||
"confidence": 0.95,
|
||||
"action": "read",
|
||||
"result": "Operation completed successfully"
|
||||
}),
|
||||
confidence=0.95
|
||||
)
|
||||
|
||||
# Execute operations
|
||||
self.add_response(
|
||||
pattern="execute|run|deploy|apply",
|
||||
response=json.dumps({
|
||||
"decision": "EXECUTE",
|
||||
"confidence": 0.80,
|
||||
"action": "execute",
|
||||
"steps": [
|
||||
{"command": "example_command", "status": "pending"}
|
||||
],
|
||||
"requires_approval": True
|
||||
}),
|
||||
confidence=0.80
|
||||
)
|
||||
|
||||
# Error/unknown
|
||||
self.add_response(
|
||||
pattern="error|fail|invalid",
|
||||
response=json.dumps({
|
||||
"decision": "ERROR",
|
||||
"confidence": 0.1,
|
||||
"error": "Simulated error response",
|
||||
"recommendations": ["Check input parameters", "Verify permissions"]
|
||||
}),
|
||||
confidence=0.1
|
||||
)
|
||||
|
||||
def add_response(self, pattern: str, response: str, confidence: float = 0.9,
|
||||
delay: float = 0.0, tokens: int = 100):
|
||||
"""
|
||||
Add a response pattern.
|
||||
|
||||
Args:
|
||||
pattern: Regex pattern to match against prompt
|
||||
response: Response to return (usually JSON string)
|
||||
confidence: Confidence score for this response
|
||||
delay: Simulated response latency (seconds)
|
||||
tokens: Simulated token usage
|
||||
"""
|
||||
self._responses.insert(0, MockResponse(
|
||||
pattern=pattern,
|
||||
response=response,
|
||||
confidence=confidence,
|
||||
delay=delay,
|
||||
tokens_used=tokens
|
||||
))
|
||||
|
||||
def set_default_response(self, response: str):
|
||||
"""Set the default response when no pattern matches"""
|
||||
self._default_response = response
|
||||
|
||||
def set_error_mode(self, error_type: str, after_calls: int = 0):
|
||||
"""
|
||||
Configure error injection.
|
||||
|
||||
Args:
|
||||
error_type: Type of error ("timeout", "rate_limit", "api_error", None)
|
||||
after_calls: Number of successful calls before error
|
||||
"""
|
||||
self._error_mode = error_type
|
||||
self._error_after = after_calls
|
||||
|
||||
def complete(self, prompt: str, max_tokens: int = 1000,
|
||||
temperature: float = 0.7) -> Tuple[str, Dict[str, Any]]:
|
||||
"""
|
||||
Generate a completion for the prompt.
|
||||
|
||||
Returns: (response_text, metadata)
|
||||
"""
|
||||
self._call_count += 1
|
||||
|
||||
# Record call
|
||||
call_record = {
|
||||
"prompt": prompt[:500], # Truncate for storage
|
||||
"max_tokens": max_tokens,
|
||||
"temperature": temperature,
|
||||
"timestamp": __import__("datetime").datetime.utcnow().isoformat()
|
||||
}
|
||||
|
||||
# Check for error injection
|
||||
if self._error_mode and self._call_count > self._error_after:
|
||||
call_record["error"] = self._error_mode
|
||||
self._call_history.append(call_record)
|
||||
raise self._create_error(self._error_mode)
|
||||
|
||||
# Find matching response
|
||||
prompt_lower = prompt.lower()
|
||||
matched_response = None
|
||||
|
||||
for mock_resp in self._responses:
|
||||
if re.search(mock_resp.pattern, prompt_lower, re.IGNORECASE):
|
||||
matched_response = mock_resp
|
||||
matched_response.call_count += 1
|
||||
break
|
||||
|
||||
if matched_response:
|
||||
response = matched_response.response
|
||||
tokens = matched_response.tokens_used
|
||||
confidence = matched_response.confidence
|
||||
|
||||
# Simulate delay
|
||||
if matched_response.delay > 0:
|
||||
import time
|
||||
time.sleep(matched_response.delay)
|
||||
else:
|
||||
response = self._default_response
|
||||
tokens = 50
|
||||
confidence = 0.5
|
||||
|
||||
self._total_tokens += tokens
|
||||
|
||||
metadata = {
|
||||
"tokens_used": tokens,
|
||||
"confidence": confidence,
|
||||
"pattern_matched": matched_response.pattern if matched_response else None,
|
||||
"total_tokens": self._total_tokens
|
||||
}
|
||||
|
||||
call_record["response"] = response[:500]
|
||||
call_record["metadata"] = metadata
|
||||
self._call_history.append(call_record)
|
||||
|
||||
return response, metadata
|
||||
|
||||
def _create_error(self, error_type: str) -> Exception:
|
||||
"""Create an appropriate error for testing"""
|
||||
if error_type == "timeout":
|
||||
return TimeoutError("LLM request timed out")
|
||||
elif error_type == "rate_limit":
|
||||
return Exception("Rate limit exceeded")
|
||||
elif error_type == "api_error":
|
||||
return Exception("API error: 500 Internal Server Error")
|
||||
else:
|
||||
return Exception(f"Unknown error: {error_type}")
|
||||
|
||||
def chat(self, messages: List[Dict[str, str]], **kwargs) -> Tuple[str, Dict[str, Any]]:
|
||||
"""
|
||||
Chat completion interface.
|
||||
|
||||
Args:
|
||||
messages: List of {"role": "user/assistant/system", "content": "..."}
|
||||
"""
|
||||
# Combine messages into a single prompt
|
||||
prompt = "\n".join([
|
||||
f"{m['role']}: {m['content']}"
|
||||
for m in messages
|
||||
])
|
||||
return self.complete(prompt, **kwargs)
|
||||
|
||||
# === Test Helpers ===
|
||||
|
||||
def reset(self):
|
||||
"""Reset all state for testing"""
|
||||
self._responses.clear()
|
||||
self._call_history.clear()
|
||||
self._total_tokens = 0
|
||||
self._error_mode = None
|
||||
self._error_after = 0
|
||||
self._call_count = 0
|
||||
self._setup_defaults()
|
||||
|
||||
def get_call_history(self) -> List[Dict[str, Any]]:
|
||||
"""Get call history for assertions"""
|
||||
return self._call_history.copy()
|
||||
|
||||
def get_call_count(self) -> int:
|
||||
"""Get total number of calls"""
|
||||
return self._call_count
|
||||
|
||||
def get_total_tokens(self) -> int:
|
||||
"""Get total tokens used"""
|
||||
return self._total_tokens
|
||||
|
||||
def assert_called_with(self, pattern: str) -> bool:
|
||||
"""Check if any call matched a pattern"""
|
||||
for call in self._call_history:
|
||||
if re.search(pattern, call["prompt"], re.IGNORECASE):
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_response_stats(self) -> Dict[str, int]:
|
||||
"""Get call counts per response pattern"""
|
||||
return {
|
||||
resp.pattern: resp.call_count
|
||||
for resp in self._responses
|
||||
if resp.call_count > 0
|
||||
}
|
||||
|
||||
|
||||
class MockLLMBuilder:
|
||||
"""Builder for creating configured MockLLM instances"""
|
||||
|
||||
def __init__(self):
|
||||
self._llm = MockLLM()
|
||||
|
||||
def with_response(self, pattern: str, response: str, **kwargs) -> 'MockLLMBuilder':
|
||||
"""Add a response pattern"""
|
||||
self._llm.add_response(pattern, response, **kwargs)
|
||||
return self
|
||||
|
||||
def with_json_response(self, pattern: str, data: Dict[str, Any],
|
||||
**kwargs) -> 'MockLLMBuilder':
|
||||
"""Add a JSON response pattern"""
|
||||
self._llm.add_response(pattern, json.dumps(data), **kwargs)
|
||||
return self
|
||||
|
||||
def with_error_after(self, calls: int, error_type: str = "api_error") -> 'MockLLMBuilder':
|
||||
"""Configure error injection"""
|
||||
self._llm.set_error_mode(error_type, calls)
|
||||
return self
|
||||
|
||||
def with_high_confidence(self) -> 'MockLLMBuilder':
|
||||
"""Configure all responses for high confidence"""
|
||||
for resp in self._llm._responses:
|
||||
resp.confidence = 0.95
|
||||
return self
|
||||
|
||||
def with_low_confidence(self) -> 'MockLLMBuilder':
|
||||
"""Configure all responses for low confidence"""
|
||||
for resp in self._llm._responses:
|
||||
resp.confidence = 0.3
|
||||
return self
|
||||
|
||||
def build(self) -> MockLLM:
|
||||
"""Build the configured MockLLM"""
|
||||
return self._llm
|
||||
@@ -0,0 +1,360 @@
|
||||
"""
|
||||
MockVault - Simulates HashiCorp Vault for testing.
|
||||
|
||||
Provides deterministic policy evaluation, token management,
|
||||
and secrets access without requiring a real Vault instance.
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Any, Optional, Tuple
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta
|
||||
import json
|
||||
import hashlib
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockToken:
|
||||
"""Represents a Vault token"""
|
||||
accessor: str
|
||||
policies: List[str]
|
||||
ttl: int
|
||||
created_at: datetime
|
||||
renewable: bool = True
|
||||
revoked: bool = False
|
||||
|
||||
def is_valid(self) -> bool:
|
||||
if self.revoked:
|
||||
return False
|
||||
if datetime.utcnow() > self.created_at + timedelta(seconds=self.ttl):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockAppRole:
|
||||
"""Represents an AppRole configuration"""
|
||||
role_id: str
|
||||
policies: List[str]
|
||||
token_ttl: int
|
||||
token_max_ttl: int
|
||||
secret_id_ttl: int
|
||||
|
||||
|
||||
class MockVault:
|
||||
"""
|
||||
Mock Vault implementation for testing.
|
||||
|
||||
Simulates:
|
||||
- AppRole authentication
|
||||
- Token lifecycle (creation, renewal, revocation)
|
||||
- Policy evaluation
|
||||
- KV secrets access
|
||||
- Audit logging
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.kv_store: Dict[str, Dict[str, Any]] = {}
|
||||
self.policies: Dict[str, Dict[str, List[str]]] = {}
|
||||
self.approles: Dict[str, MockAppRole] = {}
|
||||
self.tokens: Dict[str, MockToken] = {}
|
||||
self.secret_ids: Dict[str, Dict[str, Any]] = {}
|
||||
self.audit_log: List[Dict[str, Any]] = []
|
||||
|
||||
# Initialize default tier policies
|
||||
self._setup_default_policies()
|
||||
self._setup_default_approles()
|
||||
|
||||
def _setup_default_policies(self):
|
||||
"""Set up default trust tier policies"""
|
||||
self.policies = {
|
||||
"t0-observer": {
|
||||
"secret/data/docs/*": ["read", "list"],
|
||||
"secret/data/inventory/*": ["read", "list"],
|
||||
"secret/data/logs/*": ["read", "list"]
|
||||
},
|
||||
"t1-operator": {
|
||||
"secret/data/docs/*": ["read", "list"],
|
||||
"secret/data/inventory/*": ["read", "list"],
|
||||
"ssh/creds/sandbox-*": ["read"],
|
||||
"secret/data/ansible/sandbox/*": ["read"]
|
||||
},
|
||||
"t2-builder": {
|
||||
"secret/data/docs/*": ["read", "list"],
|
||||
"secret/data/inventory/*": ["read", "list"],
|
||||
"ssh/creds/sandbox-*": ["read"],
|
||||
"secret/data/frameworks/*": ["create", "read", "update", "list"]
|
||||
},
|
||||
"t3-executor": {
|
||||
"secret/data/*": ["read", "list"],
|
||||
"ssh/creds/*/root-controlled": ["read"],
|
||||
"proxmox/creds/staging": ["read"]
|
||||
},
|
||||
"t4-architect": {
|
||||
"secret/data/*": ["read", "list", "create", "update"],
|
||||
"sys/policies/acl/*": ["read", "list"],
|
||||
"secret/data/governance/*": ["create", "read", "update"]
|
||||
}
|
||||
}
|
||||
|
||||
def _setup_default_approles(self):
|
||||
"""Set up default AppRole configurations"""
|
||||
tier_configs = [
|
||||
("tier0-agent", ["t0-observer"], 3600, 14400, 86400),
|
||||
("tier1-agent", ["t1-operator"], 1800, 7200, 86400),
|
||||
("tier2-agent", ["t2-builder"], 1800, 7200, 86400),
|
||||
("tier3-agent", ["t3-executor"], 900, 3600, 86400),
|
||||
("tier4-agent", ["t4-architect"], 900, 3600, 86400),
|
||||
]
|
||||
|
||||
for role_name, policies, ttl, max_ttl, secret_ttl in tier_configs:
|
||||
role_id = self._generate_id(f"role-{role_name}")
|
||||
self.approles[role_name] = MockAppRole(
|
||||
role_id=role_id,
|
||||
policies=policies,
|
||||
token_ttl=ttl,
|
||||
token_max_ttl=max_ttl,
|
||||
secret_id_ttl=secret_ttl
|
||||
)
|
||||
|
||||
def _generate_id(self, prefix: str) -> str:
|
||||
"""Generate a deterministic ID"""
|
||||
h = hashlib.sha256(prefix.encode()).hexdigest()[:8]
|
||||
return f"{prefix}-{h}"
|
||||
|
||||
def _audit(self, operation: str, path: str, success: bool, details: Dict = None):
|
||||
"""Record audit log entry"""
|
||||
self.audit_log.append({
|
||||
"timestamp": datetime.utcnow().isoformat(),
|
||||
"operation": operation,
|
||||
"path": path,
|
||||
"success": success,
|
||||
"details": details or {}
|
||||
})
|
||||
|
||||
# === AppRole Authentication ===
|
||||
|
||||
def get_role_id(self, role_name: str) -> Optional[str]:
|
||||
"""Get role_id for an AppRole"""
|
||||
if role_name in self.approles:
|
||||
return self.approles[role_name].role_id
|
||||
return None
|
||||
|
||||
def generate_secret_id(self, role_name: str) -> Optional[str]:
|
||||
"""Generate a new secret_id for an AppRole"""
|
||||
if role_name not in self.approles:
|
||||
return None
|
||||
|
||||
secret_id = self._generate_id(f"secret-{role_name}-{datetime.utcnow().timestamp()}")
|
||||
self.secret_ids[secret_id] = {
|
||||
"role_name": role_name,
|
||||
"created_at": datetime.utcnow(),
|
||||
"ttl": self.approles[role_name].secret_id_ttl
|
||||
}
|
||||
return secret_id
|
||||
|
||||
def approle_login(self, role_id: str, secret_id: str) -> Tuple[bool, Optional[str], str]:
|
||||
"""
|
||||
Authenticate with AppRole credentials.
|
||||
|
||||
Returns: (success, token, message)
|
||||
"""
|
||||
# Find role by role_id
|
||||
role_name = None
|
||||
for name, approle in self.approles.items():
|
||||
if approle.role_id == role_id:
|
||||
role_name = name
|
||||
break
|
||||
|
||||
if not role_name:
|
||||
self._audit("approle_login", "auth/approle/login", False, {"error": "invalid_role_id"})
|
||||
return False, None, "Invalid role_id"
|
||||
|
||||
# Validate secret_id
|
||||
if secret_id not in self.secret_ids:
|
||||
self._audit("approle_login", "auth/approle/login", False, {"error": "invalid_secret_id"})
|
||||
return False, None, "Invalid secret_id"
|
||||
|
||||
secret_info = self.secret_ids[secret_id]
|
||||
if secret_info["role_name"] != role_name:
|
||||
self._audit("approle_login", "auth/approle/login", False, {"error": "secret_id_mismatch"})
|
||||
return False, None, "secret_id does not match role"
|
||||
|
||||
# Check secret_id TTL
|
||||
age = (datetime.utcnow() - secret_info["created_at"]).total_seconds()
|
||||
if age > secret_info["ttl"]:
|
||||
self._audit("approle_login", "auth/approle/login", False, {"error": "secret_id_expired"})
|
||||
return False, None, "secret_id expired"
|
||||
|
||||
# Create token
|
||||
approle = self.approles[role_name]
|
||||
token = self._generate_id(f"token-{role_name}-{datetime.utcnow().timestamp()}")
|
||||
accessor = self._generate_id(f"accessor-{token}")
|
||||
|
||||
self.tokens[token] = MockToken(
|
||||
accessor=accessor,
|
||||
policies=approle.policies,
|
||||
ttl=approle.token_ttl,
|
||||
created_at=datetime.utcnow()
|
||||
)
|
||||
|
||||
self._audit("approle_login", "auth/approle/login", True, {
|
||||
"role": role_name,
|
||||
"token_accessor": accessor
|
||||
})
|
||||
|
||||
return True, token, f"Authenticated as {role_name}"
|
||||
|
||||
# === Token Management ===
|
||||
|
||||
def validate_token(self, token: str) -> Tuple[bool, Optional[MockToken]]:
|
||||
"""Validate a token"""
|
||||
if token not in self.tokens:
|
||||
return False, None
|
||||
|
||||
mock_token = self.tokens[token]
|
||||
if not mock_token.is_valid():
|
||||
return False, None
|
||||
|
||||
return True, mock_token
|
||||
|
||||
def revoke_token(self, token: str) -> bool:
|
||||
"""Revoke a token"""
|
||||
if token in self.tokens:
|
||||
self.tokens[token].revoked = True
|
||||
self._audit("revoke", f"auth/token/revoke/{token[:8]}...", True)
|
||||
return True
|
||||
return False
|
||||
|
||||
def renew_token(self, token: str, increment: int = None) -> Tuple[bool, int]:
|
||||
"""
|
||||
Renew a token.
|
||||
|
||||
Returns: (success, new_ttl)
|
||||
"""
|
||||
valid, mock_token = self.validate_token(token)
|
||||
if not valid:
|
||||
return False, 0
|
||||
|
||||
if not mock_token.renewable:
|
||||
return False, 0
|
||||
|
||||
# Reset TTL (up to max)
|
||||
mock_token.created_at = datetime.utcnow()
|
||||
self._audit("renew", f"auth/token/renew", True, {"new_ttl": mock_token.ttl})
|
||||
|
||||
return True, mock_token.ttl
|
||||
|
||||
# === Policy Evaluation ===
|
||||
|
||||
def check_policy(self, token: str, path: str, capability: str) -> bool:
|
||||
"""
|
||||
Check if token has capability for path.
|
||||
|
||||
Supports glob patterns (* at end of path).
|
||||
"""
|
||||
valid, mock_token = self.validate_token(token)
|
||||
if not valid:
|
||||
return False
|
||||
|
||||
for policy_name in mock_token.policies:
|
||||
if policy_name not in self.policies:
|
||||
continue
|
||||
|
||||
policy = self.policies[policy_name]
|
||||
for policy_path, capabilities in policy.items():
|
||||
if self._path_matches(path, policy_path):
|
||||
if capability in capabilities:
|
||||
self._audit("policy_check", path, True, {
|
||||
"policy": policy_name,
|
||||
"capability": capability
|
||||
})
|
||||
return True
|
||||
|
||||
self._audit("policy_check", path, False, {"capability": capability})
|
||||
return False
|
||||
|
||||
def _path_matches(self, actual_path: str, policy_path: str) -> bool:
|
||||
"""Check if actual path matches policy path pattern"""
|
||||
if policy_path.endswith("/*"):
|
||||
prefix = policy_path[:-2]
|
||||
return actual_path.startswith(prefix)
|
||||
elif policy_path.endswith("*"):
|
||||
prefix = policy_path[:-1]
|
||||
return actual_path.startswith(prefix)
|
||||
else:
|
||||
return actual_path == policy_path
|
||||
|
||||
# === KV Secrets ===
|
||||
|
||||
def kv_put(self, path: str, data: Dict[str, Any], token: str = None) -> bool:
|
||||
"""Store a secret"""
|
||||
if token and not self.check_policy(token, path, "create"):
|
||||
return False
|
||||
|
||||
self.kv_store[path] = {
|
||||
"data": data,
|
||||
"metadata": {
|
||||
"created_time": datetime.utcnow().isoformat(),
|
||||
"version": self.kv_store.get(path, {}).get("metadata", {}).get("version", 0) + 1
|
||||
}
|
||||
}
|
||||
self._audit("kv_put", path, True)
|
||||
return True
|
||||
|
||||
def kv_get(self, path: str, token: str = None) -> Optional[Dict[str, Any]]:
|
||||
"""Retrieve a secret"""
|
||||
if token and not self.check_policy(token, path, "read"):
|
||||
return None
|
||||
|
||||
if path in self.kv_store:
|
||||
self._audit("kv_get", path, True)
|
||||
return self.kv_store[path]["data"]
|
||||
|
||||
self._audit("kv_get", path, False, {"error": "not_found"})
|
||||
return None
|
||||
|
||||
def kv_list(self, path: str, token: str = None) -> List[str]:
|
||||
"""List secrets under path"""
|
||||
if token and not self.check_policy(token, path, "list"):
|
||||
return []
|
||||
|
||||
prefix = path.rstrip("/") + "/"
|
||||
keys = []
|
||||
for key in self.kv_store.keys():
|
||||
if key.startswith(prefix):
|
||||
# Get the next path component
|
||||
remainder = key[len(prefix):]
|
||||
next_component = remainder.split("/")[0]
|
||||
if next_component and next_component not in keys:
|
||||
keys.append(next_component)
|
||||
|
||||
return keys
|
||||
|
||||
# === Test Helpers ===
|
||||
|
||||
def reset(self):
|
||||
"""Reset all state for testing"""
|
||||
self.kv_store.clear()
|
||||
self.tokens.clear()
|
||||
self.secret_ids.clear()
|
||||
self.audit_log.clear()
|
||||
self._setup_default_policies()
|
||||
self._setup_default_approles()
|
||||
|
||||
def get_audit_log(self) -> List[Dict[str, Any]]:
|
||||
"""Get audit log for test assertions"""
|
||||
return self.audit_log.copy()
|
||||
|
||||
def inject_policy(self, name: str, rules: Dict[str, List[str]]):
|
||||
"""Inject a custom policy for testing"""
|
||||
self.policies[name] = rules
|
||||
|
||||
def inject_token(self, token: str, policies: List[str], ttl: int = 3600):
|
||||
"""Inject a pre-created token for testing"""
|
||||
accessor = self._generate_id(f"accessor-{token}")
|
||||
self.tokens[token] = MockToken(
|
||||
accessor=accessor,
|
||||
policies=policies,
|
||||
ttl=ttl,
|
||||
created_at=datetime.utcnow()
|
||||
)
|
||||
Reference in New Issue
Block a user