Explorar o código

安全:移除旧工具中的硬编码数据库与飞书凭据

SamLee hai 2 días
pai
achega
ce9a2e93a5

+ 53 - 27
agent/agent/tools/builtin/browser/sync_mysql_help.py

@@ -1,14 +1,24 @@
-import pymysql
+import os
+from typing import Any, Dict, Optional, Tuple
 
 
-from typing import Tuple, Any, Dict, Literal, Optional
-from dbutils.pooled_db import PooledDB, PooledDedicatedDBConnection
-from dbutils.steady_db import SteadyDBCursor
-from pymysql.cursors import DictCursor
+def _database_config() -> dict[str, str]:
+    values = {
+        "host": os.environ.get("AGENT_BROWSER_DATABASE_HOST"),
+        "user": os.environ.get("AGENT_BROWSER_DATABASE_USER"),
+        "password": os.environ.get("AGENT_BROWSER_DATABASE_PASSWORD"),
+        "database": os.environ.get("AGENT_BROWSER_DATABASE_NAME"),
+    }
+    missing = [name for name, value in values.items() if not value]
+    if missing:
+        raise RuntimeError(
+            "legacy browser database configuration is missing: " + ", ".join(missing)
+        )
+    return {name: value for name, value in values.items() if value is not None}
 
 
 class SyncMySQLHelper(object):
-    _pool: PooledDB = None
+    _pool: Any = None
     _instance = None
 
     def __new__(cls, *args, **kwargs):
@@ -19,50 +29,69 @@ class SyncMySQLHelper(object):
 
     def get_pool(self):
         if self._pool is None:
+            required = _database_config()
+            import pymysql
+            from dbutils.pooled_db import PooledDB
+
             self._pool = PooledDB(
                 creator=pymysql,
                 mincached=10,
                 maxconnections=20,
                 blocking=True,
-                host='rm-t4na9qj85v7790tf84o.mysql.singapore.rds.aliyuncs.com',
+                host=required["host"],
                 port=3306,
-                user='crawler_admin',
-                password='cyber#crawler_2023',
-                database='aigc-admin-prod')
+                user=required["user"],
+                password=required["password"],
+                database=required["database"],
+            )
 
         return self._pool
 
-    def fetchone(self, sql: str, data: Optional[Tuple[Any, ...]] = None) -> Dict[str, Any]:
+    @staticmethod
+    def _dict_cursor() -> Any:
+        from pymysql.cursors import DictCursor
+
+        return DictCursor
+
+    def fetchone(
+        self, sql: str, data: Optional[Tuple[Any, ...]] = None
+    ) -> Dict[str, Any]:
         pool = self.get_pool()
-        with pool.connection() as conn:  
-            with conn.cursor(DictCursor) as cursor: 
+        with pool.connection() as conn:
+            with conn.cursor(self._dict_cursor()) as cursor:
                 cursor.execute(sql, data)
                 result = cursor.fetchone()
                 return result
 
-    def fetchall(self, sql: str, data: Optional[Tuple[Any, ...]] = None) -> Tuple[Dict[str, Any]]:
+    def fetchall(
+        self, sql: str, data: Optional[Tuple[Any, ...]] = None
+    ) -> Tuple[Dict[str, Any]]:
         pool = self.get_pool()
-        with pool.connection() as conn: 
-            with conn.cursor(DictCursor) as cursor: 
+        with pool.connection() as conn:
+            with conn.cursor(self._dict_cursor()) as cursor:
                 cursor.execute(sql, data)
                 result = cursor.fetchall()
                 return result
 
-    def fetchmany(self,
-                  sql: str,
-                  data: Optional[Tuple[Any, ...]] = None,
-                  size: Optional[int] = None) -> Tuple[Dict[str, Any]]:
+    def fetchmany(
+        self,
+        sql: str,
+        data: Optional[Tuple[Any, ...]] = None,
+        size: Optional[int] = None,
+    ) -> Tuple[Dict[str, Any]]:
         pool = self.get_pool()
-        with pool.connection() as conn:  
-            with conn.cursor(DictCursor) as cursor: 
+        with pool.connection() as conn:
+            with conn.cursor(self._dict_cursor()) as cursor:
                 cursor.execute(sql, data)
                 result = cursor.fetchmany(size=size)
                 return result
 
     def execute(self, sql: str, data: Optional[Tuple[Any, ...]] = None):
+        import pymysql
+
         pool = self.get_pool()
-        with pool.connection() as conn:  
-            with conn.cursor(DictCursor) as cursor:  
+        with pool.connection() as conn:
+            with conn.cursor(self._dict_cursor()) as cursor:
                 try:
                     cursor.execute(sql, data)
                     result = conn.commit()
@@ -81,6 +110,3 @@ class SyncMySQLHelper(object):
 
 
 mysql = SyncMySQLHelper()
-
-
-

+ 5 - 6
agent/agent/tools/builtin/feishu/chat.py

@@ -3,15 +3,14 @@ import os
 import base64
 import httpx
 import asyncio
-from typing import Optional, List, Dict, Any, Union
-from .feishu_client import FeishuClient, FeishuDomain
+from typing import Optional, List, Dict, Any
+from .feishu_client import FeishuClient
 from agent.tools import tool, ToolResult, ToolContext
 from agent.trace.models import MessageContent
 
-# 从环境变量获取飞书配置
-# 也可以在此设置硬编码的默认值,但推荐使用环境变量
-FEISHU_APP_ID = os.getenv("FEISHU_APP_ID", "cli_a90fe317987a9cc9")
-FEISHU_APP_SECRET = os.getenv("FEISHU_APP_SECRET", "nn2dWuXTiRA2N6xodbm4g0qz1AfM2ayi")
+# 凭据只能由运行环境注入;缺失时 FeishuClient 会 fail-closed。
+FEISHU_APP_ID = os.getenv("FEISHU_APP_ID", "")
+FEISHU_APP_SECRET = os.getenv("FEISHU_APP_SECRET", "")
 
 CONTACTS_FILE = os.path.join(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "..")), "config", "feishu_contacts.json")
 CHAT_HISTORY_DIR = os.path.join(os.path.dirname(__file__), "chat_history")

+ 17 - 17
agent/agent/tools/builtin/feishu/chat_test.py

@@ -1,19 +1,8 @@
 import asyncio
-import json
 import os
-import sys
 
-# 将项目根目录添加到 python 路径,确保可以正确导入 agent 包
-PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", ".."))
-if PROJECT_ROOT not in sys.path:
-    sys.path.append(PROJECT_ROOT)
+from agent.tools.builtin.feishu.chat import feishu_send_message_to_contact
 
-from agent.tools.builtin.feishu.chat import (
-    feishu_get_contact_list,
-    feishu_send_message_to_contact,
-    feishu_get_contact_replies,
-    feishu_get_chat_history
-)
 
 async def feishu_tools():
     print("开始测试飞书工具...\n")
@@ -61,16 +50,27 @@ async def feishu_tools():
 
     # # 4. 测试获取历史记录
     # print(f"--- 测试: feishu_get_chat_history (对象: {contact_name}) ---")
-    # result_history = await feishu_get_chat_history(contact_name, page_size=5, page_token="4cXSlmN7uFAnWWU5yfIGMNvUNrBPLlXZREzLcnvUtOcmK2QFKfwEqfbui_UDsR-y8ne0BkzXABiYTAQASh-n7my_3zQp6o3ERRz0bZ4LB5zMvahf8x7OQoso1rjrMaKM")
+    # result_history = await feishu_get_chat_history(
+    #     contact_name,
+    #     page_size=5,
+    #     page_token=os.environ.get("FEISHU_TEST_PAGE_TOKEN"),
+    # )
     # print(f"标题: {result_history.title}")
     # print(f"历史记录输出: {result_history.output}")
     # print("-" * 30 + "\n")
 
+
 if __name__ == "__main__":
-    # 模拟环境变量 (如果在系统环境变量中已设置,此处可省略)
-    os.environ["FEISHU_APP_ID"] = "cli_a90fe317987a9cc9"
-    os.environ["FEISHU_APP_SECRET"] = "nn2dWuXTiRA2N6xodbm4g0qz1AfM2ayi"
-    
+    missing = [
+        name
+        for name in ("FEISHU_APP_ID", "FEISHU_APP_SECRET")
+        if not os.environ.get(name)
+    ]
+    if missing:
+        raise RuntimeError(
+            f"missing required environment variables: {', '.join(missing)}"
+        )
+
     try:
         asyncio.run(feishu_tools())
     except KeyboardInterrupt:

+ 11 - 14
agent/agent/tools/builtin/feishu/feishu_client.py

@@ -26,16 +26,15 @@ from enum import Enum
 from typing import Any, Callable, Dict, List, Optional, Union
 
 import lark_oapi as lark
-from lark_oapi.api.contact.v3 import GetUserRequest, GetUserResponse
+from lark_oapi.api.contact.v3 import GetUserRequest
 from lark_oapi.api.im.v1 import (
     CreateMessageRequest, CreateMessageRequestBody,
     ReplyMessageRequest, ReplyMessageRequestBody,
-    GetMessageRequest, GetMessageResponse,
-    PatchMessageRequest, PatchMessageRequestBody,
+    GetMessageRequest,
     CreateImageRequest, CreateImageRequestBody,
     GetImageRequest,
     CreateFileRequest, CreateFileRequestBody,
-    GetMessageResourceRequest, GetMessageResourceResponse,
+    GetMessageResourceRequest,
     ListMessageRequest, ListMessageResponse
 )
 
@@ -116,6 +115,8 @@ class FeishuClient:
             encrypt_key: 事件加密密钥 (可选)
             verification_token: 事件验证令牌 (可选)
         """
+        if not app_id or not app_secret:
+            raise RuntimeError("FEISHU_APP_ID and FEISHU_APP_SECRET are required")
         self.app_id = app_id
         self.app_secret = app_secret
         self.domain = domain
@@ -557,7 +558,7 @@ class FeishuClient:
             parsed = json.loads(content)
             if item.msg_type == "text" and "text" in parsed:
                 content = parsed["text"]
-        except:
+        except (json.JSONDecodeError, TypeError):
             pass
 
         return {
@@ -618,7 +619,7 @@ class FeishuClient:
                     parsed = json.loads(content)
                     if item.msg_type == "text" and "text" in parsed:
                         content = parsed["text"]
-                except:
+                except (json.JSONDecodeError, TypeError):
                     pass
 
                 messages.append({
@@ -813,7 +814,7 @@ class FeishuClient:
             elif message_type == "post":
                 return self._parse_post_content(parsed)
             return content
-        except:
+        except (json.JSONDecodeError, TypeError, AttributeError):
             return content
 
     def _parse_post_content(self, parsed: Dict) -> str:
@@ -876,8 +877,8 @@ class FeishuClient:
 
 if __name__ == "__main__":
     # 从环境变量获取配置
-    APP_ID = os.getenv("FEISHU_APP_ID", "cli_a90fe317987a9cc9")
-    APP_SECRET = os.getenv("FEISHU_APP_SECRET", "nn2dWuXTiRA2N6xodbm4g0qz1AfM2ayi")
+    APP_ID = os.getenv("FEISHU_APP_ID")
+    APP_SECRET = os.getenv("FEISHU_APP_SECRET")
 
     if not APP_ID or not APP_SECRET:
         print("请设置环境变量 FEISHU_APP_ID 和 FEISHU_APP_SECRET")
@@ -892,7 +893,7 @@ if __name__ == "__main__":
 
     # 消息处理回调
     def handle_message(event: FeishuMessageEvent):
-        print(f"\n收到消息:")
+        print("\n收到消息:")
         print(f"  发送者: {event.sender_name or event.sender_open_id}")
         print(f"  类型: {event.chat_type.value}")
         print(f"  内容: {event.content}")
@@ -902,10 +903,6 @@ if __name__ == "__main__":
         if event.chat_type == ChatType.P2P or event.mentioned_bot:
             # 先回复文字
             reply_text = f"收到你的消息: {event.content}"
-            chat_id = event.chat_id
-            content = event.content
-            content_type = event.content_type # image、text等
-            open_id = event.sender_open_id
             client.send_message(
                 to=event.chat_id,
                 text=reply_text,

+ 68 - 0
agent/tests/test_legacy_browser_database_config.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+from types import ModuleType
+
+import pytest
+
+
+_VARIABLES = (
+    "AGENT_BROWSER_DATABASE_HOST",
+    "AGENT_BROWSER_DATABASE_USER",
+    "AGENT_BROWSER_DATABASE_PASSWORD",
+    "AGENT_BROWSER_DATABASE_NAME",
+)
+
+
+def _load_helper_module() -> ModuleType:
+    path = (
+        Path(__file__).parents[1]
+        / "agent"
+        / "tools"
+        / "builtin"
+        / "browser"
+        / "sync_mysql_help.py"
+    )
+    spec = importlib.util.spec_from_file_location(
+        "agent_browser_database_helper_test", path
+    )
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_legacy_browser_database_config_fails_closed(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    for name in _VARIABLES:
+        monkeypatch.delenv(name, raising=False)
+    module = _load_helper_module()
+    helper = module.SyncMySQLHelper()
+    helper._pool = None
+
+    with pytest.raises(
+        RuntimeError, match="legacy browser database configuration is missing"
+    ):
+        helper.get_pool()
+
+
+def test_legacy_browser_database_config_uses_environment(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    configured = {
+        "AGENT_BROWSER_DATABASE_HOST": "db.example.invalid",
+        "AGENT_BROWSER_DATABASE_USER": "unit-user",
+        "AGENT_BROWSER_DATABASE_PASSWORD": "unit-test-only",
+        "AGENT_BROWSER_DATABASE_NAME": "unit-database",
+    }
+    for name, value in configured.items():
+        monkeypatch.setenv(name, value)
+    module = _load_helper_module()
+    resolved = module._database_config()
+
+    assert resolved["host"] == configured["AGENT_BROWSER_DATABASE_HOST"]
+    assert resolved["user"] == configured["AGENT_BROWSER_DATABASE_USER"]
+    assert resolved["password"] == configured["AGENT_BROWSER_DATABASE_PASSWORD"]
+    assert resolved["database"] == configured["AGENT_BROWSER_DATABASE_NAME"]

+ 45 - 0
agent/tests/test_legacy_feishu_credentials.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+
+import pytest
+
+
+_FEISHU_FILES = (
+    "chat.py",
+    "chat_test.py",
+    "feishu_client.py",
+)
+_CREDENTIAL_NAMES = {"FEISHU_APP_ID", "FEISHU_APP_SECRET"}
+
+
+@pytest.mark.parametrize("filename", _FEISHU_FILES)
+def test_legacy_feishu_credentials_have_no_nonempty_fallback(filename: str) -> None:
+    path = (
+        Path(__file__).parents[1] / "agent" / "tools" / "builtin" / "feishu" / filename
+    )
+    tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
+
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Attribute):
+            continue
+        if node.func.attr != "getenv" or not node.args:
+            continue
+        name = node.args[0]
+        if not isinstance(name, ast.Constant) or name.value not in _CREDENTIAL_NAMES:
+            continue
+        if len(node.args) > 1:
+            fallback = node.args[1]
+            assert isinstance(fallback, ast.Constant)
+            assert fallback.value in {None, ""}
+
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.Assign) or not isinstance(node.value, ast.Constant):
+            continue
+        for target in node.targets:
+            if not isinstance(target, ast.Subscript):
+                continue
+            key = target.slice
+            if isinstance(key, ast.Constant) and key.value in _CREDENTIAL_NAMES:
+                assert node.value.value in {None, ""}