feat: Add NTI Security Wrapper for post-quantum agent governance

python/packages/core/agent_framework/nti_security.py

@@ -0,0 +1,49 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+import uuid
+from typing import Any
+
+from agent_framework._experimental import experimental
+from ube_foundation import TrustEngine, PqcKeyPair
+
+@experimental
+class NTISecurityWrapper:
+    """
+    NTI Security Wrapper for Microsoft Agent Framework.
+    Verifies message exchanges and tool calls with post-quantum signatures.
+    """
+    def __init__(self, agent_id: str) -> None:
+        self.agent_id = agent_id
+        self.engine = TrustEngine()
+        self.pqc_key = PqcKeyPair.generate()
+
+    def grant_capability(self, capability: str) -> None:
+        self.engine.grant(self.agent_id, capability)
+
+    def verify_message(
+        self, 
+        capability: str, 
+        message: dict[str, Any], 
+        sender_public_key: str, 
+        sender_signature: str
+    ) -> bool:
+        """
+        Verifies an incoming message from a remote agent using their public key and signature.
+        """
+        req = {
+            "id": f"req-{uuid.uuid4()}",
+            "actor": self.agent_id,
+            "capability": capability,
+            "action": "message_exchange",
+            "input": message,
+            "signature": None,
+            "pqc_signature": sender_signature,
+            "public_key": None,
+            "pqc_public_key": sender_public_key,
+            "token": None,
+            "identity_claim": None
+        }
+
+        decision = json.loads(self.engine.evaluate(json.dumps(req)))
+        return decision.get("decision") == "Allow"

python/packages/core/tests/test_nti_security.py

@@ -0,0 +1,55 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from unittest.mock import MagicMock, patch
+import pytest
+from agent_framework.nti_security import NTISecurityWrapper
+
+@patch("agent_framework.nti_security.TrustEngine")
+@patch("agent_framework.nti_security.PqcKeyPair")
+def test_nti_security_wrapper_init(mock_pqc, mock_trust):
+    wrapper = NTISecurityWrapper(agent_id="test_agent")
+    assert wrapper.agent_id == "test_agent"
+    mock_trust.assert_called_once()
+    mock_pqc.generate.assert_called_once()
+
+@patch("agent_framework.nti_security.TrustEngine")
+@patch("agent_framework.nti_security.PqcKeyPair")
+def test_verify_message_allowed(mock_pqc, mock_trust):
+    # Setup mocks
+    mock_engine_instance = mock_trust.return_value
+    mock_engine_instance.evaluate.return_value = '{"decision": "Allow"}'
+
+    wrapper = NTISecurityWrapper(agent_id="test_agent")
+
+    result = wrapper.verify_message(
+        capability="test_cap",
+        message={"data": "test"},
+        sender_public_key="pub_key",
+        sender_signature="sig"
+    )
+
+    assert result is True
+    mock_engine_instance.evaluate.assert_called_once()
+
+    # Verify the payload sent to the Rust engine contains the sender's info
+    call_args = mock_engine_instance.evaluate.call_args[0][0]
+    assert "test_cap" in call_args
+    assert "pub_key" in call_args
+    assert "sig" in call_args
+
+@patch("agent_framework.nti_security.TrustEngine")
+@patch("agent_framework.nti_security.PqcKeyPair")
+def test_verify_message_denied(mock_pqc, mock_trust):
+    mock_engine_instance = mock_trust.return_value
+    mock_engine_instance.evaluate.return_value = '{"decision": "Deny"}'
+
+    wrapper = NTISecurityWrapper(agent_id="test_agent")
+
+    result = wrapper.verify_message(
+        capability="test_cap",
+        message={"data": "test"},
+        sender_public_key="pub_key",
+        sender_signature="sig"
+    )
+
+    assert result is False