430 lines
17 KiB
Python
430 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import unittest
|
|
import uuid
|
|
from copy import deepcopy
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
|
|
class ExternalCredentialCryptoTests(unittest.TestCase):
|
|
def test_requires_master_key_and_never_falls_back_to_plaintext(self):
|
|
from core.external_systems.crypto import encrypt_secret
|
|
|
|
with patch.dict(os.environ, {}, clear=True):
|
|
with self.assertRaisesRegex(RuntimeError, "ZCBOT_CREDENTIAL_MASTER_KEY"):
|
|
encrypt_secret("secret")
|
|
|
|
def test_roundtrip_uses_ciphertext(self):
|
|
from core.external_systems.crypto import decrypt_secret, encrypt_secret
|
|
|
|
with patch.dict(os.environ, {"ZCBOT_CREDENTIAL_MASTER_KEY": "unit-test-key-at-least-32-characters"}, clear=False):
|
|
stored = encrypt_secret("mes-password")
|
|
self.assertTrue(stored.startswith("v1:"))
|
|
self.assertNotIn("mes-password", stored)
|
|
self.assertEqual(decrypt_secret(stored), "mes-password")
|
|
|
|
def test_rejects_short_master_key(self):
|
|
from core.external_systems.crypto import configured, encrypt_secret
|
|
|
|
with patch.dict(os.environ, {"ZCBOT_CREDENTIAL_MASTER_KEY": "too-short"}, clear=False):
|
|
self.assertFalse(configured())
|
|
with self.assertRaisesRegex(RuntimeError, "至少需要 32"):
|
|
encrypt_secret("mes-password")
|
|
|
|
|
|
def _cfg(*, allowed=frozenset(), recommended=()):
|
|
from core.external_systems.factory import FactoryMesConfig
|
|
|
|
return FactoryMesConfig(
|
|
base_url="https://factory.invalid",
|
|
openapi_url="https://factory.invalid/swagger.json",
|
|
login_path="/api/auth/token/",
|
|
allowed_post_operations=frozenset(allowed),
|
|
timeout_seconds=5,
|
|
max_result_bytes=65536,
|
|
max_total_result_bytes=262144,
|
|
max_page_size=200,
|
|
verify_tls=True,
|
|
query_guidance="先查数据集目录",
|
|
recommended_operation_ids=tuple(recommended),
|
|
)
|
|
|
|
|
|
_SPEC = {
|
|
"swagger": "2.0",
|
|
"paths": {
|
|
"/api/bi/dataset/": {
|
|
"get": {
|
|
"operationId": "bi_dataset_list",
|
|
"summary": "复杂统计查询的数据集目录",
|
|
"tags": ["BI", "数据集", "报表"],
|
|
"parameters": [],
|
|
}
|
|
},
|
|
"/api/qm/ftestwork/{batch}/": {
|
|
"get": {
|
|
"operationId": "qm_ftestwork_read",
|
|
"summary": "查询成品检验批次",
|
|
"tags": ["quality"],
|
|
"parameters": [
|
|
{"name": "batch", "in": "path", "required": True, "type": "string"},
|
|
{"name": "page_size", "in": "query", "required": False, "type": "integer"},
|
|
],
|
|
}
|
|
},
|
|
"/api/bi/dataset/{code}/exec/": {
|
|
"post": {
|
|
"operationId": "bi_dataset_exec",
|
|
"summary": "执行只读数据集",
|
|
"parameters": [
|
|
{"name": "code", "in": "path", "required": True, "type": "string"},
|
|
{"name": "payload", "in": "body", "required": True, "schema": {"type": "object"}},
|
|
],
|
|
}
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
class _Response:
|
|
def __init__(self, status_code=200, payload=None):
|
|
self.status_code = status_code
|
|
self._payload = payload
|
|
self.headers = {"content-type": "application/json"}
|
|
self.text = json.dumps(payload, ensure_ascii=False)
|
|
|
|
def json(self):
|
|
return self._payload
|
|
|
|
|
|
class _Http:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
return False
|
|
|
|
def post(self, url, **kwargs):
|
|
self.calls.append(("POST", url, kwargs))
|
|
return _Response(payload={"access": "remote-jwt"})
|
|
|
|
def get(self, url, **kwargs):
|
|
self.calls.append(("GET", url, kwargs))
|
|
return _Response(payload=_SPEC)
|
|
|
|
def request(self, method, url, **kwargs):
|
|
self.calls.append((method, url, kwargs))
|
|
return _Response(payload={"count": 1, "results": [{"batch": "B/1"}]})
|
|
|
|
|
|
class FactoryOpenApiConnectorTests(unittest.TestCase):
|
|
def setUp(self):
|
|
from core.external_systems import factory
|
|
factory._SPEC_CACHE.clear()
|
|
|
|
def test_admin_mapping_builds_bounded_runtime_config(self):
|
|
from core.external_systems.factory import FactoryMesConfig
|
|
|
|
cfg = FactoryMesConfig.from_mapping({
|
|
"base_url": "https://factory.invalid/",
|
|
"openapi_url": "https://factory.invalid/swagger.json",
|
|
"allowed_post_operations": "bi_dataset_exec, report_preview",
|
|
"timeout_seconds": 999,
|
|
"max_result_bytes": 1,
|
|
"verify_tls": True,
|
|
})
|
|
self.assertEqual(cfg.base_url, "https://factory.invalid")
|
|
self.assertEqual(cfg.timeout_seconds, 60)
|
|
self.assertEqual(cfg.max_result_bytes, 4096)
|
|
self.assertEqual(cfg.max_total_result_bytes, 262144)
|
|
self.assertEqual(cfg.max_page_size, 200)
|
|
self.assertEqual(cfg.allowed_post_operations, {"bi_dataset_exec", "report_preview"})
|
|
self.assertIn("dataset list", cfg.query_guidance)
|
|
self.assertEqual(
|
|
cfg.recommended_operation_ids,
|
|
("bi_dataset_list", "bi_dataset_exec"),
|
|
)
|
|
|
|
def test_admin_mapping_rejects_embedded_url_credentials(self):
|
|
from core.external_systems.factory import FactoryMesConfig, FactoryMesError
|
|
|
|
with self.assertRaisesRegex(FactoryMesError, "不能内嵌凭据"):
|
|
FactoryMesConfig.from_mapping({
|
|
"base_url": "https://user:secret@factory.invalid",
|
|
"openapi_url": "https://factory.invalid/swagger.json",
|
|
})
|
|
|
|
def test_search_discovers_operation_without_exposing_credentials(self):
|
|
from core.external_systems.factory import FactoryMesClient
|
|
|
|
http = _Http()
|
|
client = FactoryMesClient("mes-user", "mes-password", _cfg())
|
|
with patch.object(client, "_client", return_value=http):
|
|
result = client.search("成品检验")
|
|
self.assertEqual(result[0]["operation_id"], "qm_ftestwork_read")
|
|
rendered = json.dumps(result, ensure_ascii=False)
|
|
self.assertNotIn("mes-password", rendered)
|
|
self.assertNotIn("remote-jwt", rendered)
|
|
|
|
def test_search_pins_callable_admin_recommendations_without_keyword_match(self):
|
|
from core.external_systems.factory import FactoryMesClient
|
|
|
|
http = _Http()
|
|
client = FactoryMesClient(
|
|
"mes-user",
|
|
"mes-password",
|
|
_cfg(
|
|
allowed={"bi_dataset_exec"},
|
|
recommended=("bi_dataset_list", "bi_dataset_exec"),
|
|
),
|
|
)
|
|
with patch.object(client, "_client", return_value=http):
|
|
result = client.search("某工段上月产量")
|
|
self.assertEqual(
|
|
[item["operation_id"] for item in result[:2]],
|
|
["bi_dataset_list", "bi_dataset_exec"],
|
|
)
|
|
self.assertTrue(all(item["recommended"] for item in result[:2]))
|
|
|
|
def test_get_call_resolves_encoded_path_and_query(self):
|
|
from core.external_systems.factory import FactoryMesClient
|
|
|
|
http = _Http()
|
|
client = FactoryMesClient("mes-user", "mes-password", _cfg())
|
|
with patch.object(client, "_client", return_value=http):
|
|
result = client.call(
|
|
"qm_ftestwork_read",
|
|
arguments={"batch": "B/1", "page_size": 50},
|
|
)
|
|
method, url, kwargs = [call for call in http.calls if call[0] == "GET" and "/api/" in call[1]][0]
|
|
self.assertEqual(method, "GET")
|
|
self.assertIn("B%2F1", url)
|
|
self.assertEqual(kwargs["params"], {"page_size": 50})
|
|
self.assertEqual(result["data"]["count"], 1)
|
|
|
|
def test_get_call_bounds_pagination_for_agent_queries(self):
|
|
from core.external_systems.factory import FactoryMesClient, FactoryMesError
|
|
|
|
spec = deepcopy(_SPEC)
|
|
spec["paths"]["/api/qm/ftestwork/{batch}/"]["get"]["parameters"].extend([
|
|
{"name": "page", "in": "query", "required": False, "type": "integer"},
|
|
{"name": "pageoff", "in": "query", "required": False, "type": "boolean"},
|
|
])
|
|
http = _Http()
|
|
client = FactoryMesClient("u", "p", _cfg())
|
|
with patch.object(client, "_client", return_value=http), patch.object(
|
|
client, "_fetch_spec", return_value=spec
|
|
):
|
|
client.call(
|
|
"qm_ftestwork_read",
|
|
arguments={"batch": "B1", "page": 1, "page_size": 99999},
|
|
)
|
|
request = next(call for call in http.calls if call[0] == "GET")
|
|
self.assertEqual(request[2]["params"]["page_size"], 200)
|
|
|
|
with patch.object(client, "authenticate", return_value="jwt"), patch.object(
|
|
client, "_fetch_spec", return_value=spec
|
|
), self.assertRaisesRegex(FactoryMesError, "不允许 page=0"):
|
|
client.call(
|
|
"qm_ftestwork_read",
|
|
arguments={"batch": "B1", "page": 0},
|
|
)
|
|
|
|
def test_swagger_base_path_is_added_to_operation_url(self):
|
|
from core.external_systems.factory import FactoryMesClient
|
|
|
|
spec = deepcopy(_SPEC)
|
|
spec["basePath"] = "/api"
|
|
spec["paths"] = {
|
|
path.removeprefix("/api"): value
|
|
for path, value in spec["paths"].items()
|
|
}
|
|
http = _Http()
|
|
client = FactoryMesClient("u", "p", _cfg())
|
|
with patch.object(client, "_client", return_value=http), patch.object(
|
|
client, "_fetch_spec", return_value=spec
|
|
):
|
|
client.call("qm_ftestwork_read", arguments={"batch": "B1"})
|
|
request = next(call for call in http.calls if call[0] == "GET")
|
|
self.assertEqual(request[1], "https://factory.invalid/api/qm/ftestwork/B1/")
|
|
|
|
def test_api_base_path_does_not_change_login_url(self):
|
|
from core.external_systems.factory import FactoryMesClient
|
|
|
|
http = _Http()
|
|
client = FactoryMesClient("u", "p", _cfg())
|
|
with patch.object(client, "_client", return_value=http):
|
|
client.authenticate()
|
|
request = next(call for call in http.calls if call[0] == "POST")
|
|
self.assertEqual(request[1], "https://factory.invalid/api/auth/token/")
|
|
|
|
def test_base_path_is_not_duplicated_when_operation_already_contains_it(self):
|
|
from core.external_systems.factory import FactoryMesClient, FactoryMesConfig
|
|
|
|
spec = {**deepcopy(_SPEC), "basePath": "/api"}
|
|
http = _Http()
|
|
client = FactoryMesClient("u", "p", _cfg())
|
|
with patch.object(client, "_client", return_value=http), patch.object(
|
|
client, "_fetch_spec", return_value=spec
|
|
):
|
|
client.call("qm_ftestwork_read", arguments={"batch": "B1"})
|
|
request = next(call for call in http.calls if call[0] == "GET")
|
|
self.assertNotIn("/api/api/", request[1])
|
|
|
|
configured_prefix = FactoryMesClient(
|
|
"u",
|
|
"p",
|
|
FactoryMesConfig(
|
|
**{**_cfg().__dict__, "base_url": "https://factory.invalid/api"}
|
|
),
|
|
)
|
|
self.assertEqual(
|
|
configured_prefix._operation_url(spec, "/api/qm/ftestwork/B1/"),
|
|
"https://factory.invalid/api/qm/ftestwork/B1/",
|
|
)
|
|
|
|
def test_openapi_server_path_is_used_but_cross_origin_server_is_rejected(self):
|
|
from core.external_systems.factory import FactoryMesClient, FactoryMesError
|
|
|
|
client = FactoryMesClient("u", "p", _cfg())
|
|
same_origin = {"openapi": "3.0.0", "servers": [{"url": "/v1"}], "paths": {}}
|
|
self.assertEqual(
|
|
client._operation_url(same_origin, "/quality/results/"),
|
|
"https://factory.invalid/v1/quality/results/",
|
|
)
|
|
cross_origin = {
|
|
"openapi": "3.0.0",
|
|
"servers": [{"url": "https://attacker.invalid/v1"}],
|
|
"paths": {},
|
|
}
|
|
with self.assertRaisesRegex(FactoryMesError, "越出 Factory MES 主机"):
|
|
client._operation_url(cross_origin, "/quality/results/")
|
|
|
|
def test_no_declared_base_path_keeps_existing_url_behavior(self):
|
|
from core.external_systems.factory import FactoryMesClient
|
|
|
|
client = FactoryMesClient("u", "p", _cfg())
|
|
self.assertEqual(
|
|
client._operation_url(_SPEC, "/api/qm/ftestwork/B1/"),
|
|
"https://factory.invalid/api/qm/ftestwork/B1/",
|
|
)
|
|
|
|
def test_post_is_denied_unless_admin_allowlists_operation(self):
|
|
from core.external_systems.factory import FactoryMesClient, FactoryMesError
|
|
|
|
denied = FactoryMesClient("u", "p", _cfg())
|
|
with patch.object(denied, "authenticate", return_value="jwt"), patch.object(
|
|
denied, "_fetch_spec", return_value=_SPEC
|
|
):
|
|
with self.assertRaisesRegex(FactoryMesError, "只读调用范围"):
|
|
denied.call("bi_dataset_exec", arguments={"code": "x", "payload": {}})
|
|
|
|
http = _Http()
|
|
allowed = FactoryMesClient("u", "p", _cfg(allowed={"bi_dataset_exec"}))
|
|
with patch.object(allowed, "_client", return_value=http):
|
|
result = allowed.call(
|
|
"bi_dataset_exec",
|
|
arguments={"code": "yield", "payload": {"query": {"month": "2026-08"}}},
|
|
)
|
|
request = [call for call in http.calls if call[0] == "POST" and "/dataset/" in call[1]][0]
|
|
self.assertEqual(request[2]["json"], {"query": {"month": "2026-08"}})
|
|
self.assertFalse(result["truncated"])
|
|
|
|
def test_allowlisted_post_accepts_separate_body_field(self):
|
|
from core.external_systems.factory import FactoryMesClient
|
|
|
|
http = _Http()
|
|
allowed = FactoryMesClient("u", "p", _cfg(allowed={"bi_dataset_exec"}))
|
|
with patch.object(allowed, "_client", return_value=http):
|
|
result = allowed.call(
|
|
"bi_dataset_exec",
|
|
arguments={"code": "quality"},
|
|
body={"batch": "B-1"},
|
|
)
|
|
request = [call for call in http.calls if call[0] == "POST" and "/dataset/" in call[1]][0]
|
|
self.assertEqual(request[2]["json"], {"batch": "B-1"})
|
|
self.assertFalse(result["truncated"])
|
|
|
|
def test_rejects_unknown_arguments(self):
|
|
from core.external_systems.factory import FactoryMesClient, FactoryMesError
|
|
|
|
client = FactoryMesClient("u", "p", _cfg())
|
|
with patch.object(client, "authenticate", return_value="jwt"), patch.object(
|
|
client, "_fetch_spec", return_value=_SPEC
|
|
):
|
|
with self.assertRaisesRegex(FactoryMesError, "接口定义之外"):
|
|
client.call(
|
|
"qm_ftestwork_read",
|
|
arguments={"batch": "B1", "unexpected": "x"},
|
|
)
|
|
|
|
|
|
class ExternalSystemToolSafetyTests(unittest.TestCase):
|
|
def test_tools_are_scoped_to_constructor_user(self):
|
|
from tools.external_systems import ExternalSystemListTool
|
|
|
|
uid = uuid.uuid4()
|
|
with patch(
|
|
"tools.external_systems.list_external_systems",
|
|
return_value=[{
|
|
"external_system_id": str(uuid.uuid4()),
|
|
"status": "active",
|
|
"username_masked": "me***r",
|
|
}],
|
|
) as listed:
|
|
output = ExternalSystemListTool(uid).execute()
|
|
listed.assert_called_once_with(uid)
|
|
self.assertNotIn("password", output.lower())
|
|
|
|
def test_search_returns_admin_guidance_with_recommended_operations(self):
|
|
from tools.external_systems import ExternalSystemSearchTool
|
|
|
|
uid = uuid.uuid4()
|
|
client = SimpleNamespace(
|
|
cfg=SimpleNamespace(
|
|
query_guidance="统计查询先查看数据集目录",
|
|
recommended_operation_ids=("bi_dataset_list",),
|
|
),
|
|
search=lambda query, limit: [{"operation_id": "bi_dataset_list"}],
|
|
)
|
|
with patch(
|
|
"tools.external_systems._row_and_client",
|
|
return_value=(SimpleNamespace(), client),
|
|
):
|
|
output = ExternalSystemSearchTool(uid).execute(str(uuid.uuid4()), "产量")
|
|
payload = json.loads(output)
|
|
self.assertEqual(payload["query_guidance"], "统计查询先查看数据集目录")
|
|
self.assertEqual(payload["recommended_operation_ids"], ["bi_dataset_list"])
|
|
|
|
def test_call_discards_result_over_per_run_external_budget(self):
|
|
from tools.external_systems import ExternalSystemCallTool
|
|
|
|
uid = uuid.uuid4()
|
|
client = SimpleNamespace(
|
|
cfg=SimpleNamespace(max_total_result_bytes=32),
|
|
call=lambda *args, **kwargs: {"data": "x" * 100},
|
|
)
|
|
with patch(
|
|
"tools.external_systems._row_and_client",
|
|
return_value=(SimpleNamespace(), client),
|
|
):
|
|
output = ExternalSystemCallTool(uid).execute(
|
|
str(uuid.uuid4()), "detail_list"
|
|
)
|
|
self.assertIn("累计返回量超过上限", output)
|
|
self.assertNotIn("x" * 20, output)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|