Codex
7eb9b987ab
Gate source-file downloads and verify registry paths and content hashes. Stream multipart uploads with bounded limits, supervise OpenCode jobs with explicit result markers and timeouts, and preserve managed endpoint security settings. Document the trusted-client boundary and add focused regression coverage. Validation: PYTHONPATH=<staged-tree> python -m unittest discover -s tests (182 tests)
446 lines
18 KiB
Python
446 lines
18 KiB
Python
import hashlib
|
|
import json
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from obsidian_wiki.access_control import source_decision as actual_source_decision
|
|
from obsidian_wiki.file_assets import connect, init_db, upsert_asset
|
|
from obsidian_wiki.mcp_assets import (
|
|
AssetAccessError,
|
|
AssetTooLargeError,
|
|
asset_uri_for_file_id,
|
|
file_id_from_asset_uri,
|
|
get_file_asset_metadata,
|
|
mcp_runtime_options,
|
|
read_file_asset_resource,
|
|
record_allowed_for_mcp,
|
|
search_file_assets,
|
|
)
|
|
|
|
|
|
class McpAssetTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tempdir = tempfile.TemporaryDirectory()
|
|
self.vault = Path(self.tempdir.name)
|
|
(self.vault / "_meta").mkdir()
|
|
self.db_path = self.vault / "_meta" / "file-assets.sqlite"
|
|
self.registry_path = self.vault / "_meta" / "source-registry.jsonl"
|
|
init_db(self.db_path)
|
|
self._write_profile_file()
|
|
|
|
self.drawing_id = self._write_asset_file("drawings/a.pdf", b"%PDF allowed drawing")
|
|
self.order_id = self._write_asset_file("orders/order.xlsx", b"customer price data")
|
|
self.complaint_id = self._write_asset_file("quality/photo.png", b"complaint evidence")
|
|
|
|
self._insert_asset(
|
|
self.drawing_id,
|
|
archive_path="_sources/archive/drawings/a.pdf",
|
|
original_name="B3501_11101c_X2.pdf",
|
|
doc_type="drawing",
|
|
source_type="pdf",
|
|
product_lines=["BTB"],
|
|
business_categories=["products/btb", "drawings/customer"],
|
|
data_classification="technical-confidential",
|
|
visibility="internal",
|
|
summary="B3501 customer drawing.",
|
|
bjd_part_numbers=["101049-X0XXX0"],
|
|
series=["0.5-18"],
|
|
pitch=["0.5"],
|
|
pin_count=["18"],
|
|
allowed_bots=["bjd-general-bot", "missing-bot"],
|
|
)
|
|
self._insert_asset(
|
|
self.order_id,
|
|
archive_path="_sources/archive/orders/order.xlsx",
|
|
original_name="订单项.xlsx",
|
|
doc_type="order_data",
|
|
source_type="spreadsheet",
|
|
business_categories=["orders/raw-excel"],
|
|
data_classification="commercial-confidential",
|
|
visibility="pii",
|
|
summary="Raw customer order lines.",
|
|
)
|
|
self._insert_asset(
|
|
self.complaint_id,
|
|
archive_path="_sources/archive/quality/photo.png",
|
|
original_name="客户反馈来料标签错误图片.png",
|
|
doc_type="evidence_media",
|
|
source_type="image",
|
|
business_categories=["quality/customer-complaints"],
|
|
data_classification="quality-confidential",
|
|
visibility="internal",
|
|
summary="Raw complaint evidence image.",
|
|
)
|
|
|
|
def tearDown(self):
|
|
self.tempdir.cleanup()
|
|
|
|
def _write_asset_file(self, rel_path, data):
|
|
path = self.vault / "_sources" / "archive" / rel_path
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_bytes(data)
|
|
return "sha256:" + hashlib.sha256(data).hexdigest()
|
|
|
|
def _insert_asset(self, file_id, **fields):
|
|
record = {
|
|
"record_type": "source_archive",
|
|
"source_id": file_id,
|
|
"project": "bjd",
|
|
"title": Path(fields["original_name"]).stem,
|
|
"size_bytes": len((self.vault / fields["archive_path"]).read_bytes()),
|
|
"related_pages": ["projects/bjd/references/example.md"],
|
|
"allowed_bots": ["bjd-general-bot"],
|
|
**fields,
|
|
}
|
|
with self.registry_path.open("a", encoding="utf-8") as fh:
|
|
fh.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
asset = {
|
|
"file_id": file_id,
|
|
"sha256": file_id.removeprefix("sha256:"),
|
|
"archive_path": fields["archive_path"],
|
|
"original_name": fields["original_name"],
|
|
"file_ext": Path(fields["original_name"]).suffix.lower(),
|
|
"mime_type": "application/pdf",
|
|
"source_type": fields["source_type"],
|
|
"doc_type": fields["doc_type"],
|
|
"project": "bjd",
|
|
"title": record["title"],
|
|
"summary": fields["summary"],
|
|
"visibility": fields["visibility"],
|
|
"business_categories": json.dumps(fields.get("business_categories", []), ensure_ascii=False),
|
|
"product_lines": json.dumps(fields.get("product_lines", []), ensure_ascii=False),
|
|
"customers": "[]",
|
|
"competitors": "[]",
|
|
"standards": "[]",
|
|
"market_tags": "[]",
|
|
"bjd_part_numbers": json.dumps(fields.get("bjd_part_numbers", []), ensure_ascii=False),
|
|
"competitor_part_numbers": json.dumps(fields.get("competitor_part_numbers", []), ensure_ascii=False),
|
|
"series": json.dumps(fields.get("series", []), ensure_ascii=False),
|
|
"pitch": json.dumps(fields.get("pitch", []), ensure_ascii=False),
|
|
"pin_count": json.dumps(fields.get("pin_count", []), ensure_ascii=False),
|
|
"related_pages": json.dumps(record["related_pages"], ensure_ascii=False),
|
|
"catalog_paths": "[]",
|
|
"extracted_text": None,
|
|
"classification_confidence": 0.9,
|
|
"classification_evidence": "[]",
|
|
"needs_review": 0,
|
|
"size_bytes": record["size_bytes"],
|
|
"modified_at": None,
|
|
"archived_at": "2026-07-05T00:00:00Z",
|
|
}
|
|
with connect(self.db_path) as conn:
|
|
upsert_asset(conn, asset, status="active")
|
|
|
|
def _write_profile_file(self):
|
|
data = {
|
|
"version": 2,
|
|
"profiles": {
|
|
"bjd-general-bot": {
|
|
"allowed_visibility": ["public", "internal"],
|
|
"allowed_data_classification": [
|
|
"public",
|
|
"internal",
|
|
"technical-confidential",
|
|
],
|
|
"allowed_business_domains": ["products", "engineering"],
|
|
"blocked_doc_types": ["order_data", "evidence_media"],
|
|
"include_source_files": False,
|
|
},
|
|
"public-product-bot": {
|
|
"allowed_visibility": ["public", "internal"],
|
|
"allowed_data_classification": [
|
|
"public",
|
|
"internal",
|
|
"technical-confidential",
|
|
],
|
|
"allowed_business_domains": ["products", "engineering"],
|
|
"include_source_files": False,
|
|
},
|
|
},
|
|
}
|
|
(self.vault / "_meta" / "bot-profiles.json").write_text(json.dumps(data), encoding="utf-8")
|
|
|
|
def test_search_returns_allowed_engineering_assets_only(self):
|
|
rows = search_file_assets(self.vault, "B3501 drawing", db_path=self.db_path)
|
|
|
|
self.assertEqual([row["file_id"] for row in rows], [self.drawing_id])
|
|
self.assertEqual(rows[0]["resource_uri"], asset_uri_for_file_id(self.drawing_id))
|
|
self.assertNotIn("absolute_path", rows[0])
|
|
|
|
def test_search_exposes_safe_identifier_fields(self):
|
|
rows = search_file_assets(self.vault, "101049 图纸", db_path=self.db_path)
|
|
|
|
self.assertEqual([row["file_id"] for row in rows], [self.drawing_id])
|
|
self.assertEqual(rows[0]["bjd_part_numbers"], ["101049-X0XXX0"])
|
|
self.assertEqual(rows[0]["series"], ["0.5-18"])
|
|
self.assertEqual(rows[0]["pitch"], ["0.5"])
|
|
self.assertEqual(rows[0]["pin_count"], ["18"])
|
|
|
|
def test_customer_shareable_engineering_files_work_with_source_file_profile(self):
|
|
shareable_id = self._write_asset_file(
|
|
"drawings/customer-shareable.pdf",
|
|
b"%PDF customer drawing",
|
|
)
|
|
self._insert_asset(
|
|
shareable_id,
|
|
archive_path="_sources/archive/drawings/customer-shareable.pdf",
|
|
original_name="北极大-0.5-18客户图.pdf",
|
|
doc_type="drawing",
|
|
source_type="pdf",
|
|
product_lines=["FPC"],
|
|
business_categories=["products/fpc", "drawings/customer-drawings"],
|
|
data_classification="customer-shareable",
|
|
visibility="public",
|
|
summary="FPC 0.5-18 customer-shareable drawing.",
|
|
bjd_part_numbers=["101049-X0XXX0"],
|
|
series=["0.5-18"],
|
|
)
|
|
|
|
rows = search_file_assets(self.vault, "101049 图纸", db_path=self.db_path, profile_name="bjd-general-bot")
|
|
|
|
self.assertEqual({row["file_id"] for row in rows}, {shareable_id, self.drawing_id})
|
|
|
|
def test_search_blocks_sensitive_order_and_complaint_sources(self):
|
|
self.assertEqual(search_file_assets(self.vault, "订单", db_path=self.db_path), [])
|
|
self.assertEqual(search_file_assets(self.vault, "来料标签错误", db_path=self.db_path), [])
|
|
|
|
def test_asset_uri_round_trips_file_id(self):
|
|
uri = asset_uri_for_file_id(self.drawing_id)
|
|
|
|
self.assertEqual(
|
|
uri,
|
|
"asset://source/sha256/" + self.drawing_id.removeprefix("sha256:"),
|
|
)
|
|
self.assertEqual(file_id_from_asset_uri(uri), self.drawing_id)
|
|
|
|
def test_resource_read_rechecks_acl_and_returns_file_bytes(self):
|
|
data = read_file_asset_resource(
|
|
self.vault,
|
|
asset_uri_for_file_id(self.drawing_id),
|
|
db_path=self.db_path,
|
|
allow_source_files=True,
|
|
)
|
|
|
|
self.assertEqual(data["name"], "B3501_11101c_X2.pdf")
|
|
self.assertEqual(data["mime_type"], "application/pdf")
|
|
self.assertEqual(data["bytes"], b"%PDF allowed drawing")
|
|
|
|
def test_resource_read_rejects_blocked_asset_even_by_direct_uri(self):
|
|
with self.assertRaises(AssetAccessError):
|
|
read_file_asset_resource(
|
|
self.vault,
|
|
asset_uri_for_file_id(self.order_id),
|
|
db_path=self.db_path,
|
|
allow_source_files=True,
|
|
)
|
|
|
|
def test_resource_read_enforces_max_size(self):
|
|
with self.assertRaises(AssetTooLargeError):
|
|
read_file_asset_resource(
|
|
self.vault,
|
|
asset_uri_for_file_id(self.drawing_id),
|
|
db_path=self.db_path,
|
|
max_file_bytes=5,
|
|
allow_source_files=True,
|
|
)
|
|
|
|
def test_resource_read_requires_explicit_server_opt_in(self):
|
|
with self.assertRaisesRegex(AssetAccessError, "source file reads are disabled"):
|
|
read_file_asset_resource(
|
|
self.vault,
|
|
asset_uri_for_file_id(self.drawing_id),
|
|
db_path=self.db_path,
|
|
)
|
|
|
|
def test_runtime_options_source_file_gate_defaults_closed_and_reads_env(self):
|
|
with patch.dict("os.environ", {}, clear=True):
|
|
self.assertFalse(mcp_runtime_options()["allow_source_files"])
|
|
with patch.dict("os.environ", {"OBSIDIAN_MCP_ALLOW_SOURCE_FILES": "1"}, clear=True):
|
|
self.assertTrue(mcp_runtime_options()["allow_source_files"])
|
|
|
|
def test_resource_read_rejects_db_registry_archive_path_mismatch(self):
|
|
other_path = "_sources/archive/drawings/other.pdf"
|
|
self._write_asset_file("drawings/other.pdf", b"%PDF other drawing")
|
|
with connect(self.db_path) as conn:
|
|
conn.execute(
|
|
"UPDATE file_assets SET archive_path = ? WHERE file_id = ?",
|
|
(other_path, self.drawing_id),
|
|
)
|
|
|
|
with self.assertRaisesRegex(AssetAccessError, "archive_path mismatch"):
|
|
read_file_asset_resource(
|
|
self.vault,
|
|
asset_uri_for_file_id(self.drawing_id),
|
|
db_path=self.db_path,
|
|
allow_source_files=True,
|
|
)
|
|
|
|
def test_resource_read_accepts_normalization_equivalent_archive_paths(self):
|
|
with connect(self.db_path) as conn:
|
|
conn.execute(
|
|
"UPDATE file_assets SET archive_path = ? WHERE file_id = ?",
|
|
("_sources/archive/drawings/../drawings/a.pdf", self.drawing_id),
|
|
)
|
|
|
|
data = read_file_asset_resource(
|
|
self.vault,
|
|
asset_uri_for_file_id(self.drawing_id),
|
|
db_path=self.db_path,
|
|
allow_source_files=True,
|
|
)
|
|
|
|
self.assertEqual(data["bytes"], b"%PDF allowed drawing")
|
|
|
|
def test_resource_read_rejects_content_that_no_longer_matches_file_id(self):
|
|
(self.vault / "_sources" / "archive" / "drawings" / "a.pdf").write_bytes(
|
|
b"%PDF tampered drawing"
|
|
)
|
|
|
|
with self.assertRaisesRegex(AssetAccessError, "content hash mismatch"):
|
|
read_file_asset_resource(
|
|
self.vault,
|
|
asset_uri_for_file_id(self.drawing_id),
|
|
db_path=self.db_path,
|
|
allow_source_files=True,
|
|
)
|
|
|
|
def test_metadata_uses_safe_fields(self):
|
|
metadata = get_file_asset_metadata(self.vault, self.drawing_id, db_path=self.db_path)
|
|
|
|
self.assertEqual(metadata["file_id"], self.drawing_id)
|
|
self.assertEqual(metadata["doc_type"], "drawing")
|
|
self.assertEqual(metadata["resource_uri"], asset_uri_for_file_id(self.drawing_id))
|
|
self.assertNotIn("absolute_path", metadata)
|
|
|
|
def test_registry_allowed_bots_are_enforced_with_valid_profile(self):
|
|
with self.assertRaises(AssetAccessError):
|
|
get_file_asset_metadata(
|
|
self.vault,
|
|
self.drawing_id,
|
|
db_path=self.db_path,
|
|
profile_name="public-product-bot",
|
|
)
|
|
|
|
def test_missing_requested_profile_fails_closed(self):
|
|
with self.assertRaisesRegex(AssetAccessError, "bot profile not found"):
|
|
search_file_assets(
|
|
self.vault,
|
|
"B3501",
|
|
db_path=self.db_path,
|
|
profile_name="missing-bot",
|
|
)
|
|
|
|
with self.assertRaisesRegex(AssetAccessError, "bot profile not found"):
|
|
get_file_asset_metadata(
|
|
self.vault,
|
|
self.drawing_id,
|
|
db_path=self.db_path,
|
|
profile_name="missing-bot",
|
|
)
|
|
|
|
def test_metadata_acl_does_not_rewrite_profile_source_file_setting(self):
|
|
profile = {
|
|
"allowed_visibility": ["internal"],
|
|
"allowed_data_classification": ["technical-confidential"],
|
|
"allowed_business_domains": ["products", "engineering"],
|
|
"include_source_files": False,
|
|
}
|
|
record = {
|
|
"archive_path": "_sources/archive/drawings/a.pdf",
|
|
"doc_type": "drawing",
|
|
"visibility": "internal",
|
|
"data_classification": "customer-shareable",
|
|
"business_categories": ["products/btb", "drawings/customer"],
|
|
"allowed_bots": ["bjd-general-bot"],
|
|
}
|
|
|
|
with patch(
|
|
"obsidian_wiki.mcp_assets.source_decision",
|
|
wraps=actual_source_decision,
|
|
) as mocked_decision:
|
|
allowed, reason = record_allowed_for_mcp(
|
|
record,
|
|
profile_name="bjd-general-bot",
|
|
profile=profile,
|
|
)
|
|
|
|
self.assertTrue(allowed, reason)
|
|
runtime_profile = mocked_decision.call_args.args[1]
|
|
self.assertIsNot(runtime_profile, profile)
|
|
self.assertTrue(runtime_profile["include_source_files"])
|
|
self.assertIn("customer-shareable", runtime_profile["allowed_data_classification"])
|
|
self.assertFalse(profile["include_source_files"])
|
|
self.assertNotIn("customer-shareable", profile["allowed_data_classification"])
|
|
|
|
def test_search_filters_acl_before_applying_limit(self):
|
|
allowed_id = self._write_asset_file(
|
|
"drawings/allowed-limit.pdf",
|
|
b"%PDF allowed lower-ranked limit result",
|
|
)
|
|
self._insert_asset(
|
|
allowed_id,
|
|
archive_path="_sources/archive/drawings/allowed-limit.pdf",
|
|
original_name="Allowed limit drawing.pdf",
|
|
doc_type="drawing",
|
|
source_type="pdf",
|
|
product_lines=["BTB"],
|
|
business_categories=["products/btb"],
|
|
data_classification="technical-confidential",
|
|
visibility="internal",
|
|
summary="BLOCKME appears only as unstructured text.",
|
|
)
|
|
for index in range(4):
|
|
blocked_id = self._write_asset_file(
|
|
f"drawings/blocked-limit-{index}.pdf",
|
|
f"%PDF blocked limit result {index}".encode(),
|
|
)
|
|
self._insert_asset(
|
|
blocked_id,
|
|
archive_path=f"_sources/archive/drawings/blocked-limit-{index}.pdf",
|
|
original_name=f"Blocked limit drawing {index}.pdf",
|
|
doc_type="drawing",
|
|
source_type="pdf",
|
|
product_lines=["BTB"],
|
|
business_categories=["products/btb"],
|
|
data_classification="technical-confidential",
|
|
visibility="internal",
|
|
summary="BLOCKME structured identifier result.",
|
|
bjd_part_numbers=["BLOCKME"],
|
|
allowed_bots=["other-bot"],
|
|
)
|
|
|
|
rows = search_file_assets(self.vault, "BLOCKME", db_path=self.db_path, limit=1)
|
|
|
|
self.assertEqual([row["file_id"] for row in rows], [allowed_id])
|
|
|
|
def test_search_clamps_limit_to_public_range(self):
|
|
for index in range(101):
|
|
file_id = self._write_asset_file(
|
|
f"drawings/bulk-{index}.pdf",
|
|
f"%PDF BULK-LIMIT {index}".encode(),
|
|
)
|
|
self._insert_asset(
|
|
file_id,
|
|
archive_path=f"_sources/archive/drawings/bulk-{index}.pdf",
|
|
original_name=f"Bulk limit drawing {index}.pdf",
|
|
doc_type="drawing",
|
|
source_type="pdf",
|
|
product_lines=["BTB"],
|
|
business_categories=["products/btb"],
|
|
data_classification="technical-confidential",
|
|
visibility="internal",
|
|
summary="BULK-LIMIT clamp candidate.",
|
|
)
|
|
|
|
lower = search_file_assets(self.vault, "B3501", db_path=self.db_path, limit=0)
|
|
upper = search_file_assets(self.vault, "BULK-LIMIT", db_path=self.db_path, limit=1000)
|
|
|
|
self.assertEqual(len(lower), 1)
|
|
self.assertEqual(len(upper), 100)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|