diff --git a/python/packages/core/agent_framework/nti_security.py b/python/packages/core/agent_framework/nti_security.py new file mode 100644 index 00000000000..958954221e2 --- /dev/null +++ b/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" diff --git a/python/packages/core/tests/test_nti_security.py b/python/packages/core/tests/test_nti_security.py new file mode 100644 index 00000000000..47b5ca6ea6d --- /dev/null +++ b/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