rex/audit.py
Rama master · réplica de lectura del repositorio original.
"""Audit logging for Rex tool execution.
This module provides audit logging functionality for tracking all tool
invocations, policy decisions, and execution results. It supports:
- Persistent storage in JSON Lines format
- Automatic redaction of sensitive fields
- Time-range queries for log retrieval
- Export functionality for archival
All log entries are automatically redacted using redact_sensitive_keys()
before being written to disk to prevent sensitive data exposure.
"""
from __future__ import annotations
import logging
from datetime import UTC, datetime
from pathlib import Path
from threading import Lock
from typing import Any, Literal
from pydantic import BaseModel, Field
from rex.contracts.core import redact_sensitive_keys
logger = logging.getLogger(__name__)
# Default log directory path
DEFAULT_LOG_DIR = Path("data/logs")
DEFAULT_LOG_FILE = "audit.log"
def _utc_now() -> datetime:
"""Return the current UTC datetime."""
return datetime.now(UTC)
class LogEntry(BaseModel):
"""A single audit log entry for tool execution.
Each entry captures the full context of a tool invocation including
the tool name, arguments, policy decision, and result. Sensitive
fields are redacted before persistence.
"""
timestamp: datetime = Field(
default_factory=_utc_now,
description="When this log entry was created (UTC)",
)
action_id: str = Field(
...,
description="Unique identifier for this action",
)
task_id: str | None = Field(
default=None,
description="ID of the parent task, if any",
)
tool: str = Field(
...,
description="Name of the tool that was invoked",
)
tool_call_args: dict[str, Any] = Field(
default_factory=dict,
description="Arguments passed to the tool (redacted before storage)",
)
policy_decision: Literal["allowed", "denied", "requires_approval"] = Field(
...,
description="Policy engine decision for this tool call",
)
tool_result: dict[str, Any] | None = Field(
default=None,
description="Result from tool execution (redacted before storage)",
)
error: str | None = Field(
default=None,
description="Error message if the tool execution failed",
)
redacted: bool = Field(
default=False,
description="Whether sensitive fields have been redacted",
)
requested_by: str | None = Field(
default=None,
description="Identifier of who/what requested this action",
)
duration_ms: int | None = Field(
default=None,
description="Execution duration in milliseconds",
)
model_config = {
"json_schema_extra": {
"examples": [
{
"timestamp": "2024-01-15T10:30:00Z",
"action_id": "act_001",
"task_id": "task_001",
"tool": "time_now",
"tool_call_args": {"location": "Dallas, TX"},
"policy_decision": "allowed",
"tool_result": {"local_time": "2024-01-15 10:30"},
"redacted": True,
"duration_ms": 15,
}
]
}
}
class AuditLogger:
"""Persistent audit logger for tool executions.
Stores audit entries in JSON Lines format with automatic redaction
of sensitive fields. Thread-safe for concurrent writes.
Usage:
logger = AuditLogger()
entry = LogEntry(
action_id="act_001",
tool="time_now",
tool_call_args={"location": "Dallas"},
policy_decision="allowed",
)
logger.log(entry)
"""
def __init__(
self,
log_dir: Path | str | None = None,
log_file: str = DEFAULT_LOG_FILE,
log_path: Path | str | None = None,
) -> None:
"""Initialize the audit logger.
Args:
log_dir: Directory to store log files. Defaults to data/logs.
log_file: Name of the log file. Defaults to audit.log.
log_path: Optional full path to the log file. Overrides log_dir/log_file.
"""
if log_path is not None:
self._log_path = Path(log_path)
self._log_dir = self._log_path.parent
self._log_file = self._log_path.name
else:
if log_dir is None:
log_dir = DEFAULT_LOG_DIR
self._log_dir = Path(log_dir)
self._log_file = log_file
self._log_path = self._log_dir / self._log_file
self._lock = Lock()
self._ensure_log_dir()
def _ensure_log_dir(self) -> None:
"""Create the log directory if it doesn't exist."""
try:
self._log_dir.mkdir(parents=True, exist_ok=True)
except OSError as e:
logger.warning("Could not create log directory %s: %s", self._log_dir, e)
@property
def log_path(self) -> Path:
"""Return the path to the audit log file."""
return self._log_path
def _redact_entry(self, entry: LogEntry) -> LogEntry:
"""Apply redaction to sensitive fields in a log entry.
Args:
entry: The log entry to redact.
Returns:
A new LogEntry with sensitive fields redacted.
"""
redacted_args = redact_sensitive_keys(entry.tool_call_args)
redacted_result = (
redact_sensitive_keys(entry.tool_result) if entry.tool_result is not None else None
)
return LogEntry(
timestamp=entry.timestamp,
action_id=entry.action_id,
task_id=entry.task_id,
tool=entry.tool,
tool_call_args=redacted_args, # type: ignore[arg-type]
policy_decision=entry.policy_decision,
tool_result=redacted_result, # type: ignore[arg-type]
error=entry.error,
redacted=True,
requested_by=entry.requested_by,
duration_ms=entry.duration_ms,
)
def log(self, entry: LogEntry) -> None:
"""Log an audit entry to the persistent store.
Redacts sensitive fields before writing to disk.
Args:
entry: The log entry to persist.
"""
redacted_entry = self._redact_entry(entry)
# Serialize with ISO format for datetime
json_line = redacted_entry.model_dump_json() + "\n"
with self._lock:
try:
with open(self._log_path, "a", encoding="utf-8") as f:
f.write(json_line)
except OSError as e:
logger.error("Failed to write audit log: %s", e)
def read(
self,
start_time: datetime | None = None,
end_time: datetime | None = None,
) -> list[LogEntry]:
"""Read log entries from the persistent store.
Args:
start_time: Optional start of time range (inclusive).
end_time: Optional end of time range (inclusive).
Returns:
List of log entries matching the time range criteria.
"""
entries: list[LogEntry] = []
if not self._log_path.exists():
return entries
try:
with open(self._log_path, encoding="utf-8") as f:
for line in f:
line = line.strip()
if not line:
continue
try:
entry = LogEntry.model_validate_json(line)
# Apply time range filter
if start_time is not None and entry.timestamp < start_time:
continue
if end_time is not None and entry.timestamp > end_time:
continue
entries.append(entry)
except Exception as e:
logger.warning("Skipping malformed log entry: %s", e)
except OSError as e:
logger.error("Failed to read audit log: %s", e)
return entries
def read_by_action_id(self, action_id: str) -> LogEntry | None:
"""Read a specific log entry by action ID.
Args:
action_id: The action ID to look up.
Returns:
The matching log entry or None if not found.
"""
if not self._log_path.exists():
return None
try:
with open(self._log_path, encoding="utf-8") as f:
for line in f:
line = line.strip()
if not line:
continue
try:
entry = LogEntry.model_validate_json(line)
if entry.action_id == action_id:
return entry
except Exception:
continue
except OSError as e:
logger.error("Failed to read audit log: %s", e)
return None
def export(self, path: str | Path) -> int:
"""Export all log entries to a specified path.
Args:
path: Destination path for the export file.
Returns:
Number of entries exported.
"""
entries = self.read()
export_path = Path(path)
# Ensure parent directory exists
export_path.parent.mkdir(parents=True, exist_ok=True)
try:
with open(export_path, "w", encoding="utf-8") as f:
for entry in entries:
f.write(entry.model_dump_json() + "\n")
logger.info("Exported %d audit entries to %s", len(entries), export_path)
return len(entries)
except OSError as e:
logger.error("Failed to export audit log: %s", e)
return 0
def clear(self) -> None:
"""Clear all entries from the audit log.
WARNING: This permanently deletes all audit history.
Use only for testing or with explicit user confirmation.
"""
with self._lock:
try:
if self._log_path.exists():
self._log_path.unlink()
except OSError as e:
logger.error("Failed to clear audit log: %s", e)
# Module-level singleton
_default_audit_logger: AuditLogger | None = None
_singleton_lock = Lock()
def get_audit_logger() -> AuditLogger:
"""Get the default audit logger singleton.
Returns:
The shared AuditLogger instance.
"""
global _default_audit_logger
with _singleton_lock:
if _default_audit_logger is None:
_default_audit_logger = AuditLogger()
return _default_audit_logger
def reset_audit_logger() -> None:
"""Reset the default audit logger singleton.
Primarily used for testing to ensure a clean state.
"""
global _default_audit_logger
with _singleton_lock:
_default_audit_logger = None
def replay(entry: LogEntry, *, dry_run: bool = True) -> Any:
"""Replay a tool execution from an audit log entry.
This is a convenience wrapper for rex.replay.replay.
"""
from rex.replay import replay as _replay
return _replay(entry, dry_run=dry_run)
__all__ = [
"LogEntry",
"AuditLogger",
"get_audit_logger",
"reset_audit_logger",
"replay",
"DEFAULT_LOG_DIR",
"DEFAULT_LOG_FILE",
]