DESKTOP-72TV0V4\caoxiaozhu
b2bc9988a9
- 新增 Go 语言后端服务(server/)，包含用户认证、Agent管理、数据库连接等API - 新增 Python Agent 服务(agent/)，实现Agent核心逻辑和工具集 - 前端从原生HTML迁移到Vue.js框架(web/src/) - 添加 Docker Compose 支持(docker-compose.yml) - 添加项目架构文档(docs/ARCHITECTURE.md) - 添加环境变量示例(.env.example)和本地启动脚本(start-local.ps1) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
97 lines
2.5 KiB
Python
97 lines
2.5 KiB
Python
"""
|
||
数据库查询工具 - 安全的数据查询接口
|
||
"""
|
||
from typing import Dict, Any, List, Optional
|
||
import os
|
||
|
||
|
||
# 只读查询白名单 - 只允许 SELECT 语句
|
||
ALLOWED_TABLES = ["users", "agents", "sessions", "audit_logs"]
|
||
|
||
|
||
class DatabaseQueryTool:
|
||
"""
|
||
数据库查询工具
|
||
|
||
安全特性：
|
||
- 只允许 SELECT 查询
|
||
- 表名白名单
|
||
- 结果数量限制
|
||
"""
|
||
|
||
def __init__(self, connection_string: str = ""):
|
||
self.connection_string = connection_string or os.getenv(
|
||
"DATABASE_URL",
|
||
"postgresql://postgres:postgres@localhost:5432/x_agents"
|
||
)
|
||
self.max_rows = 100 # 最多返回100行
|
||
|
||
def query(self, sql: str, params: List[Any] = None) -> Dict[str, Any]:
|
||
"""
|
||
执行查询
|
||
|
||
Args:
|
||
sql: SQL 查询语句（必须是 SELECT）
|
||
params: 查询参数
|
||
|
||
Returns:
|
||
查询结果
|
||
"""
|
||
# 安全检查1: 必须是 SELECT 语句
|
||
sql_upper = sql.strip().upper()
|
||
if not sql_upper.startswith("SELECT"):
|
||
return {
|
||
"success": False,
|
||
"error": "Only SELECT queries are allowed"
|
||
}
|
||
|
||
# 安全检查2: 禁止危险关键字
|
||
dangerous_keywords = [
|
||
"DROP", "DELETE", "INSERT", "UPDATE", "ALTER",
|
||
"CREATE", "TRUNCATE", "EXEC", "EXECUTE"
|
||
]
|
||
for keyword in dangerous_keywords:
|
||
if keyword in sql_upper:
|
||
return {
|
||
"success": False,
|
||
"error": f"Keyword '{keyword}' is not allowed"
|
||
}
|
||
|
||
# 安全检查3: 表名白名单
|
||
for table in ALLOWED_TABLES:
|
||
if f"FROM {table}" in sql_upper or f"JOIN {table}" in sql_upper:
|
||
# 表名在白名单中，允许
|
||
break
|
||
else:
|
||
# 没有找到白名单表
|
||
return {
|
||
"success": False,
|
||
"error": f"Table not in whitelist. Allowed: {ALLOWED_TABLES}"
|
||
}
|
||
|
||
# TODO: 实际执行查询（需要数据库连接）
|
||
# 这里返回模拟数据
|
||
return {
|
||
"success": True,
|
||
"message": "Query executed (mock mode - database not connected)",
|
||
"rows": [],
|
||
"columns": []
|
||
}
|
||
|
||
|
||
# 全局实例
|
||
db_tool = DatabaseQueryTool()
|
||
|
||
|
||
def query_data(sql: str) -> Dict[str, Any]:
|
||
"""
|
||
查询数据工具
|
||
|
||
Args:
|
||
sql: SELECT 查询语句
|
||
|
||
Returns:
|
||
查询结果
|
||
"""
|
||
return db_tool.query(sql)
|