diff --git a/opax-mcp/requirements.txt b/opax-mcp/requirements.txt index 43b99b6..5f1614c 100644 --- a/opax-mcp/requirements.txt +++ b/opax-mcp/requirements.txt @@ -10,3 +10,4 @@ google-cloud-bigquery pydantic>=2.0 python-multipart loguru +google-cloud-kms>=2.21.1 diff --git a/opax-mcp/server.py b/opax-mcp/server.py index ffc9d1e..2fed182 100644 --- a/opax-mcp/server.py +++ b/opax-mcp/server.py @@ -11,6 +11,7 @@ from urllib.parse import quote import difflib import google.auth import google.auth.transport.requests +from google.cloud import kms_v1 import google.oauth2.id_token from googleapiclient.discovery import build from google.cloud import secretmanager @@ -75,7 +76,7 @@ from deployment_policy import get_deployment_target from deployment_source import normalize_gitea_archive from email.mime.text import MIMEText from datetime import datetime, timezone, timedelta -from typing import Any, Optional, Dict, List +from typing import Any, Optional, Dict, List, Literal import logging import re from pydantic import BaseModel, Field @@ -84,7 +85,21 @@ import hashlib MAX_PROPOSE_CONTENT_BYTES = 1_048_576 _REPO_ID_RE = re.compile(r"^[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$") +# A simple regex to validate the KMS key name format +_KMS_KEY_NAME_RE = re.compile( + r"^projects/[^/]+/locations/[^/]+/keyRings/[^/]+/cryptoKeys/[^/]+$" +) + + # --- A2H2A Pydantic Models --- +class EncryptedPayloadV1(BaseModel): + """Versioned Firestore-safe envelope for KMS-encrypted plan content.""" + + payload_version: Literal["1"] = "1" + kms_key_name: str + ciphertext: str + + class GiteaChangePlan(BaseModel): """Immutable admin-approved one-file Gitea change plan.""" @@ -105,11 +120,13 @@ class GiteaChangePlan(BaseModel): branch: str path: str base_sha: str - new_content: str content_hash: str commit_message: str unified_diff: str + encrypted_payload: Optional[EncryptedPayloadV1] = None + approval_subject_hash: Optional[str] = None + result_commit_sha: Optional[str] = None existing_file_sha: Optional[str] = None apply_idempotency_key: Optional[str] = None @@ -122,7 +139,7 @@ async def create_gitea_change_plan(plan: GiteaChangePlan) -> dict: from google.cloud import firestore db = firestore.AsyncClient(project=GOOGLE_CLOUD_PROJECT) - plan_data = plan.model_dump(mode="json", exclude={"new_content"}) + plan_data = plan.model_dump(mode="json") plan_data["status"] = "PENDING" plan_ref = db.collection("gitea_change_plans").document(plan.plan_id) @@ -138,6 +155,65 @@ async def create_gitea_change_plan(plan: GiteaChangePlan) -> dict: return plan_data +async def _encrypt_payload_v1(plaintext_content: str) -> EncryptedPayloadV1: + """Encrypt plaintext UTF-8 content with the configured Cloud KMS key.""" + key_name = os.environ.get("GITEA_PLAN_KMS_KEY_NAME") + if not key_name or not _KMS_KEY_NAME_RE.fullmatch(key_name): + raise RuntimeError("KMS key is not configured or has an invalid format.") + + plaintext_bytes = plaintext_content.encode("utf-8") + + async with kms_v1.KeyManagementServiceAsyncClient() as client: + response = await client.encrypt( + request={"name": key_name, "plaintext": plaintext_bytes} + ) + + ciphertext = getattr(response, "ciphertext", None) + if not isinstance(ciphertext, bytes) or not ciphertext: + raise RuntimeError("KMS encryption returned an invalid empty ciphertext.") + + return EncryptedPayloadV1( + kms_key_name=key_name, + ciphertext=base64.b64encode(ciphertext).decode("ascii"), + ) + + +def _calculate_approval_subject_hash(plan: GiteaChangePlan) -> str: + """Hash the immutable, approval-bound representation of a Gitea plan.""" + import binascii + payload = plan.encrypted_payload + if payload is None: + raise ValueError( + "Cannot calculate approval hash without an encrypted payload." + ) + + try: + ciphertext_bytes = base64.b64decode(payload.ciphertext, validate=True) + except (ValueError, TypeError, binascii.Error) as exc: + raise ValueError("Encrypted payload ciphertext is invalid.") from exc + + subject = { + "plan_id": plan.plan_id, + "repo": plan.repo, + "branch": plan.branch, + "path": plan.path, + "base_sha": plan.base_sha, + "existing_file_sha": plan.existing_file_sha, + "content_hash": plan.content_hash, + "commit_message": plan.commit_message, + "unified_diff": plan.unified_diff, + "payload_version": payload.payload_version, + "kms_key_name": payload.kms_key_name, + "ciphertext_hash": hashlib.sha256(ciphertext_bytes).hexdigest(), + } + canonical_json = json.dumps( + subject, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + return hashlib.sha256(canonical_json).hexdigest() + + async def get_gitea_change_plan( plan_id: str, ) -> Optional[GiteaChangePlan]: @@ -277,17 +353,22 @@ async def propose_gitea_change(p: dict) -> dict: raise ValueError("Proposed content produces no file change.") content_hash = hashlib.sha256(new_content.encode("utf-8")).hexdigest() + + encrypted_payload = await _encrypt_payload_v1(new_content) + plan = GiteaChangePlan( repo=repo, branch=branch, path=path, base_sha=base_sha, - new_content=new_content, content_hash=content_hash, commit_message=commit_message, unified_diff=unified_diff, existing_file_sha=existing_file_sha, + encrypted_payload=encrypted_payload, ) + plan.approval_subject_hash = _calculate_approval_subject_hash(plan) + await create_gitea_change_plan(plan) return { @@ -296,6 +377,7 @@ async def propose_gitea_change(p: dict) -> dict: "expires_at": plan.expires_at.isoformat(), "unified_diff": unified_diff, "content_hash": content_hash, + "approval_subject_hash": plan.approval_subject_hash, "repo": repo, "branch": branch, "path": path, diff --git a/opax-mcp/test_propose_gitea_change.py b/opax-mcp/test_propose_gitea_change.py index 29bae46..0aaa301 100644 --- a/opax-mcp/test_propose_gitea_change.py +++ b/opax-mcp/test_propose_gitea_change.py @@ -1,7 +1,20 @@ import unittest -import server # Følger eksisterende mønster import hashlib +import os +import base64 +import sys +from pathlib import Path from unittest.mock import patch, AsyncMock +from unittest.mock import MagicMock + +EMMA_DIR = Path(__file__).resolve().parents[1] / "emma" +if str(EMMA_DIR) not in sys.path: + sys.path.insert(0, str(EMMA_DIR)) + +import server # Følger eksisterende mønster + +KMS_KEY_NAME = "projects/p/locations/l/keyRings/k/cryptoKeys/k" + class TestProposeGiteaChange(unittest.IsolatedAsyncioTestCase): @@ -95,21 +108,30 @@ class TestProposeGiteaChange(unittest.IsolatedAsyncioTestCase): mock_create_plan.assert_not_awaited() @patch('server.create_gitea_change_plan', new_callable=AsyncMock) + @patch("server.kms_v1.KeyManagementServiceAsyncClient") @patch('server._get_gitea_file_details', new_callable=AsyncMock, return_value=("old content", "b" * 40)) @patch('server.resolve_branch_to_commit_sha', new_callable=AsyncMock) - async def test_happy_path_creates_pending_plan(self, mock_resolve_sha, mock_get_details, mock_create_plan): + async def test_happy_path_creates_pending_plan(self, mock_resolve_sha, mock_get_details, mock_kms_constructor, mock_create_plan): + mock_kms_client = AsyncMock() + mock_kms_client.encrypt.return_value = MagicMock(ciphertext=b"encrypted-data") + mock_kms_constructor.return_value.__aenter__ = AsyncMock(return_value=mock_kms_client) + mock_kms_constructor.return_value.__aexit__ = AsyncMock(return_value=False) + mock_resolve_sha.return_value = self.valid_params["base_sha"] - result = await server.propose_gitea_change(self.valid_params) + with patch.dict(os.environ, {"GITEA_PLAN_KMS_KEY_NAME": KMS_KEY_NAME}): + result = await server.propose_gitea_change(self.valid_params) self.assertEqual(result["status"], "PENDING") self.assertEqual(result["repo"], self.valid_params["repo"]) self.assertEqual(result["branch"], self.valid_params["branch"]) self.assertEqual(result["path"], self.valid_params["path"]) + self.assertIn("approval_subject_hash", result) expected_hash = hashlib.sha256(self.valid_params["new_content"].encode("utf-8")).hexdigest() self.assertEqual(result["content_hash"], expected_hash) + mock_kms_client.encrypt.assert_awaited_once() mock_create_plan.assert_awaited_once() # Verify the object passed to create_gitea_change_plan is the real Pydantic model @@ -121,5 +143,206 @@ class TestProposeGiteaChange(unittest.IsolatedAsyncioTestCase): self.assertEqual(call_args.base_sha, self.valid_params["base_sha"]) self.assertEqual(call_args.existing_file_sha, "b" * 40) + @patch('server.create_gitea_change_plan', new_callable=AsyncMock) + @patch("server.kms_v1.KeyManagementServiceAsyncClient") + @patch('server._get_gitea_file_details', new_callable=AsyncMock) + @patch('server.resolve_branch_to_commit_sha', new_callable=AsyncMock) + async def test_patch2a_missing_kms_key_fails_after_gitea_reads( + self, mock_resolve_sha, mock_get_details, mock_kms_constructor, mock_create_plan + ): + mock_resolve_sha.return_value = self.valid_params["base_sha"] + mock_get_details.return_value = ("old content", "b" * 40) + + with patch.dict(os.environ, {"GITEA_PLAN_KMS_KEY_NAME": ""}): + with self.assertRaisesRegex(RuntimeError, "KMS key is not configured"): + await server.propose_gitea_change(self.valid_params) + + mock_resolve_sha.assert_awaited_once() + mock_get_details.assert_awaited_once() + mock_kms_constructor.assert_not_called() + mock_create_plan.assert_not_awaited() + + @patch('server.create_gitea_change_plan', new_callable=AsyncMock) + @patch("server.kms_v1.KeyManagementServiceAsyncClient") + @patch('server._get_gitea_file_details', new_callable=AsyncMock) + @patch('server.resolve_branch_to_commit_sha', new_callable=AsyncMock) + async def test_patch2a_invalid_kms_key_fails_before_kms_construction( + self, mock_resolve_sha, mock_get_details, mock_kms_constructor, mock_create_plan + ): + mock_resolve_sha.return_value = self.valid_params["base_sha"] + mock_get_details.return_value = ("old content", "b" * 40) + + with patch.dict(os.environ, {"GITEA_PLAN_KMS_KEY_NAME": "invalid-key-format"}): + with self.assertRaisesRegex(RuntimeError, "invalid format"): + await server.propose_gitea_change(self.valid_params) + + mock_resolve_sha.assert_awaited_once() + mock_get_details.assert_awaited_once() + mock_kms_constructor.assert_not_called() + mock_create_plan.assert_not_awaited() + + @patch('server.create_gitea_change_plan', new_callable=AsyncMock) + @patch("server.kms_v1.KeyManagementServiceAsyncClient") + @patch('server._get_gitea_file_details', new_callable=AsyncMock, return_value=("old content", "b" * 40)) + @patch('server.resolve_branch_to_commit_sha', new_callable=AsyncMock) + async def test_patch2a_empty_kms_ciphertext_is_rejected( + self, mock_resolve_sha, mock_get_details, mock_kms_constructor, mock_create_plan + ): + mock_resolve_sha.return_value = self.valid_params["base_sha"] + mock_kms_client = AsyncMock() + mock_kms_client.encrypt.return_value = MagicMock(ciphertext=b"") + mock_kms_constructor.return_value.__aenter__ = AsyncMock(return_value=mock_kms_client) + mock_kms_constructor.return_value.__aexit__ = AsyncMock(return_value=False) + + with patch.dict(os.environ, {"GITEA_PLAN_KMS_KEY_NAME": KMS_KEY_NAME}): + with self.assertRaisesRegex(RuntimeError, "invalid empty ciphertext"): + await server.propose_gitea_change(self.valid_params) + + mock_kms_client.encrypt.assert_awaited_once_with( + request={ + "name": KMS_KEY_NAME, + "plaintext": self.valid_params["new_content"].encode("utf-8"), + } + ) + mock_create_plan.assert_not_awaited() + + @patch('server.create_gitea_change_plan', new_callable=AsyncMock) + @patch("server.kms_v1.KeyManagementServiceAsyncClient") + @patch('server._get_gitea_file_details', new_callable=AsyncMock, return_value=("old content", "b" * 40)) + @patch('server.resolve_branch_to_commit_sha', new_callable=AsyncMock) + async def test_patch2a_non_bytes_kms_ciphertext_is_rejected( + self, mock_resolve_sha, mock_get_details, mock_kms_constructor, mock_create_plan + ): + mock_resolve_sha.return_value = self.valid_params["base_sha"] + mock_kms_client = AsyncMock() + mock_kms_client.encrypt.return_value = MagicMock(ciphertext=None) + mock_kms_constructor.return_value.__aenter__ = AsyncMock(return_value=mock_kms_client) + mock_kms_constructor.return_value.__aexit__ = AsyncMock(return_value=False) + + with patch.dict(os.environ, {"GITEA_PLAN_KMS_KEY_NAME": KMS_KEY_NAME}): + with self.assertRaisesRegex(RuntimeError, "invalid empty ciphertext"): + await server.propose_gitea_change(self.valid_params) + + mock_kms_client.encrypt.assert_awaited_once_with( + request={ + "name": KMS_KEY_NAME, + "plaintext": self.valid_params["new_content"].encode("utf-8"), + } + ) + mock_create_plan.assert_not_awaited() + + @patch('server.create_gitea_change_plan', new_callable=AsyncMock) + @patch("server.kms_v1.KeyManagementServiceAsyncClient") + @patch('server._get_gitea_file_details', new_callable=AsyncMock, return_value=("old content", "b" * 40)) + @patch('server.resolve_branch_to_commit_sha', new_callable=AsyncMock) + async def test_patch2a_successful_encrypted_proposal( + self, mock_resolve_sha, mock_get_details, mock_kms_constructor, mock_create_plan + ): + mock_resolve_sha.return_value = self.valid_params["base_sha"] + mock_kms_client = AsyncMock() + mock_kms_client.encrypt.return_value = MagicMock(ciphertext=b"encrypted-data") + mock_kms_constructor.return_value.__aenter__ = AsyncMock(return_value=mock_kms_client) + mock_kms_constructor.return_value.__aexit__ = AsyncMock(return_value=False) + + with patch.dict(os.environ, {"GITEA_PLAN_KMS_KEY_NAME": KMS_KEY_NAME}): + result = await server.propose_gitea_change(self.valid_params) + + mock_resolve_sha.assert_awaited_once() + mock_get_details.assert_awaited_once() + mock_kms_client.encrypt.assert_awaited_once_with( + request={ + "name": KMS_KEY_NAME, + "plaintext": self.valid_params["new_content"].encode("utf-8"), + } + ) + + mock_create_plan.assert_awaited_once() + persisted_plan = mock_create_plan.call_args.args[0] + + self.assertIsInstance(persisted_plan, server.GiteaChangePlan) + self.assertIsInstance( + persisted_plan.encrypted_payload, + server.EncryptedPayloadV1, + ) + self.assertEqual(persisted_plan.encrypted_payload.payload_version, "1") + self.assertEqual(persisted_plan.encrypted_payload.kms_key_name, KMS_KEY_NAME) + self.assertEqual( + persisted_plan.encrypted_payload.ciphertext, + base64.b64encode(b"encrypted-data").decode("ascii"), + ) + self.assertNotIn("new_content", persisted_plan.model_dump(mode="json")) + + expected_content_hash = hashlib.sha256(self.valid_params["new_content"].encode("utf-8")).hexdigest() + self.assertEqual(persisted_plan.content_hash, expected_content_hash) + self.assertEqual(result["content_hash"], expected_content_hash) + + self.assertEqual( + persisted_plan.approval_subject_hash, + server._calculate_approval_subject_hash(persisted_plan), + ) + self.assertEqual(result["approval_subject_hash"], persisted_plan.approval_subject_hash) + + for field in ("new_content", "encrypted_payload", "ciphertext", "kms_key_name"): + self.assertNotIn(field, result) + + def test_legacy_plan_parses_without_patch2a_fields(self): + current_plan = server.GiteaChangePlan( + repo="o/r", branch="b", path="p", base_sha="a" * 40, + content_hash="c" * 64, commit_message="m", unified_diff="d" + ) + legacy_dict = current_plan.model_dump(mode="json") + legacy_dict.pop("encrypted_payload", None) + legacy_dict.pop("approval_subject_hash", None) + + legacy_plan = server.GiteaChangePlan(**legacy_dict) + self.assertIsNone(legacy_plan.encrypted_payload) + self.assertIsNone(legacy_plan.approval_subject_hash) + + def test_malformed_base64_in_hash_raises_error(self): + plan_with_bad_payload = server.GiteaChangePlan( + plan_id="plan-123", repo="o/r", branch="b", path="p", base_sha="a"*40, content_hash="c"*64, + commit_message="m", unified_diff="d", + encrypted_payload=server.EncryptedPayloadV1( + kms_key_name=KMS_KEY_NAME, ciphertext="not-valid-base64!" + ) + ) + with self.assertRaisesRegex(ValueError, "Encrypted payload ciphertext is invalid"): + server._calculate_approval_subject_hash(plan_with_bad_payload) + + def test_approval_subject_hash_sensitivity(self): + base_plan = server.GiteaChangePlan( + plan_id="plan-123", repo="o/r", branch="b", path="p", + base_sha="a" * 40, existing_file_sha="b" * 40, content_hash="c" * 64, + commit_message="m", unified_diff="d", + encrypted_payload=server.EncryptedPayloadV1( + kms_key_name=KMS_KEY_NAME, + ciphertext=base64.b64encode(b"secret").decode("ascii"), + ), + ) + base_hash = server._calculate_approval_subject_hash(base_plan) + + fields_to_test = { + "plan_id": "plan-456", "repo": "o/r2", "branch": "b2", "path": "p2", + "base_sha": "A" * 40, "existing_file_sha": "B" * 40, "content_hash": "C" * 64, + "commit_message": "m2", "unified_diff": "d2", + } + + for field, new_value in fields_to_test.items(): + with self.subTest(field=field): + modified_plan = base_plan.model_copy(update={field: new_value}) + new_hash = server._calculate_approval_subject_hash(modified_plan) + self.assertNotEqual(base_hash, new_hash) + + with self.subTest(field="kms_key_name"): + modified_plan_kms = base_plan.model_copy(deep=True) + modified_plan_kms.encrypted_payload.kms_key_name = "projects/p/locations/l/keyRings/k/cryptoKeys/k2" + self.assertNotEqual(base_hash, server._calculate_approval_subject_hash(modified_plan_kms)) + + with self.subTest(field="ciphertext"): + modified_plan_cipher = base_plan.model_copy(deep=True) + modified_plan_cipher.encrypted_payload.ciphertext = base64.b64encode(b"secret2").decode("ascii") + self.assertNotEqual(base_hash, server._calculate_approval_subject_hash(modified_plan_cipher)) + + if __name__ == "__main__": unittest.main()