"""审计日志系统
记录所有数据库写操作，提供审计和追踪功能。
"""
import json
import logging
import os
import time
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Set, Union
# 设置审计日志记录器
audit_logger = logging.getLogger("mcp_dbutils.audit")
audit_logger.setLevel(logging.INFO)
# 内存缓冲区，保存最近的日志记录
_memory_buffer: List[Dict[str, Any]] = []
_memory_buffer_size = 1000 # 默认缓冲区大小
# 审计日志配置
_audit_config = {
"enabled": True,
"file_storage": {
"enabled": True,
"path": "logs/audit",
"max_file_size": 10 * 1024 * 1024, # 10MB
"backup_count": 10
},
"content": {
"sanitize_sql": True,
"include_user_context": True
}
}
def configure_audit_logging(config: Dict[str, Any]) -> None:
"""配置审计日志系统
Args:
config: 审计日志配置
"""
global _audit_config, _memory_buffer_size
if config:
_audit_config.update(config)
# 更新内存缓冲区大小
if "memory_buffer" in config and "size" in config["memory_buffer"]:
_memory_buffer_size = config["memory_buffer"]["size"]
# 配置文件处理器
if _audit_config["enabled"] and _audit_config["file_storage"]["enabled"]:
_setup_file_handler()
def _setup_file_handler() -> None:
"""设置文件日志处理器"""
log_dir = _audit_config["file_storage"]["path"]
os.makedirs(log_dir, exist_ok=True)
# 创建文件处理器
file_handler = logging.FileHandler(
f"{log_dir}/dbutils-audit.log",
mode="a",
encoding="utf-8"
)
# 设置格式
formatter = logging.Formatter("%(asctime)s - %(message)s")
file_handler.setFormatter(formatter)
# 添加到记录器
audit_logger.addHandler(file_handler)
def _sanitize_sql(sql: str) -> str:
"""对SQL语句进行脱敏处理
Args:
sql: SQL语句
Returns:
脱敏后的SQL语句
"""
if not _audit_config["content"]["sanitize_sql"]:
return sql
# 简单的脱敏处理，替换VALUES子句中的值
import re
# 替换INSERT语句中的VALUES
sanitized = re.sub(
r"VALUES\s*\((.*?)\)",
"VALUES (?)",
sql,
flags=re.IGNORECASE | re.DOTALL
)
# 替换WHERE子句中的值
sanitized = re.sub(
r"(WHERE\s+\w+\s*=\s*)('[^']*'|\d+)",
r"\1?",
sanitized,
flags=re.IGNORECASE
)
return sanitized
def _get_user_context() -> Optional[str]:
"""获取用户上下文
Returns:
用户上下文信息
"""
if not _audit_config["content"]["include_user_context"]:
return None
# 这里可以添加获取用户上下文的逻辑
# 例如从环境变量、请求头等获取
return None
def log_write_operation(
connection_name: str,
table_name: str,
operation_type: str,
sql: str,
affected_rows: int,
execution_time: float,
status: str = "SUCCESS",
error_message: Optional[str] = None
) -> None:
"""记录写操作到审计日志
Args:
connection_name: 数据库连接名称
table_name: 表名
operation_type: 操作类型（INSERT、UPDATE、DELETE）
sql: SQL语句
affected_rows: 影响的行数
execution_time: 执行时间（毫秒）
status: 操作结果（SUCCESS、FAILED）
error_message: 错误信息（如果失败）
"""
if not _audit_config["enabled"]:
return
# 创建日志记录
log_entry = {
"timestamp": datetime.now().isoformat(),
"connection_name": connection_name,
"table_name": table_name,
"operation_type": operation_type,
"sql_statement": _sanitize_sql(sql),
"affected_rows": affected_rows,
"status": status,
"execution_time": execution_time,
"user_context": _get_user_context()
}
if error_message:
log_entry["error_message"] = error_message
# 添加到内存缓冲区
_memory_buffer.append(log_entry)
if len(_memory_buffer) > _memory_buffer_size:
_memory_buffer.pop(0) # 移除最旧的记录
# 写入日志文件
if _audit_config["file_storage"]["enabled"]:
audit_logger.info(json.dumps(log_entry))
def get_logs(
connection_name: Optional[str] = None,
table_name: Optional[str] = None,
start_time: Optional[str] = None,
end_time: Optional[str] = None,
operation_type: Optional[str] = None,
status: Optional[str] = None,
limit: int = 100
) -> List[Dict[str, Any]]:
"""获取审计日志
Args:
connection_name: 数据库连接名称
table_name: 表名
start_time: 开始时间（ISO 8601格式）
end_time: 结束时间（ISO 8601格式）
operation_type: 操作类型（INSERT、UPDATE、DELETE）
status: 操作状态（SUCCESS、FAILED）
limit: 返回记录数量限制
Returns:
审计日志记录列表
"""
# 从内存缓冲区获取日志
logs = _memory_buffer.copy()
# 应用过滤条件
filtered_logs = []
for log in logs:
# 连接名称过滤
if connection_name and log["connection_name"] != connection_name:
continue
# 表名过滤
if table_name and log["table_name"] != table_name:
continue
# 时间范围过滤
if start_time and log["timestamp"] < start_time:
continue
if end_time and log["timestamp"] > end_time:
continue
# 操作类型过滤
if operation_type and log["operation_type"] != operation_type:
continue
# 状态过滤
if status and log["status"] != status:
continue
filtered_logs.append(log)
# 应用限制
if len(filtered_logs) >= limit:
break
return filtered_logs
def format_logs(logs: List[Dict[str, Any]]) -> str:
"""格式化审计日志
Args:
logs: 审计日志记录列表
Returns:
格式化后的审计日志
"""
if not logs:
return "No audit logs found."
formatted = ["Audit Logs:", "==========="]
for log in logs:
formatted.extend([
f"\nTimestamp: {log['timestamp']}",
f"Connection: {log['connection_name']}",
f"Table: {log['table_name']}",
f"Operation: {log['operation_type']}",
f"Status: {log['status']}",
f"Affected Rows: {log['affected_rows']}",
f"Execution Time: {log['execution_time']:.2f}ms",
f"SQL: {log['sql_statement']}"
])
if "error_message" in log:
formatted.append(f"Error: {log['error_message']}")
if log.get("user_context"):
formatted.append(f"User Context: {log['user_context']}")
return "\n".join(formatted)