Adds encryption to service w/ testing + fixes

Author: marshallmcdonnell

src/intersect_sdk/_internal/encryption/__init__.py

@@ -1,10 +1,10 @@
 from .aes_cypher import AESCipher
-from .client_encryption import intersect_client_encryption
-from .service_decryption import intersect_service_decryption
+from .client_encryption import intersect_payload_encrypt
+from .service_decryption import intersect_payload_decrypt
 
 
 __all__ = (
     'AESCipher',
-    'intersect_client_encryption',
-    'intersect_service_decryption',
+    'intersect_payload_encrypt',
+    'intersect_payload_decrypt',
 )

src/intersect_sdk/_internal/encryption/client_encryption.py

@@ -14,7 +14,7 @@
 from .models import IntersectEncryptedPayload, IntersectEncryptionPublicKey
 from ..logger import logger
 
-def intersect_client_encryption(
+def intersect_payload_encrypt(
     key_payload: IntersectEncryptionPublicKey,
     unencrypted_model: str,
 ) -> IntersectEncryptedPayload:

src/intersect_sdk/_internal/encryption/service_decryption.py

@@ -16,7 +16,7 @@
 from ..logger import logger
 
 
-def intersect_service_decryption(
+def intersect_payload_decrypt(
     rsa_private_key: rsa.RSAPrivateKey,
     encrypted_payload: IntersectEncryptedPayload,
     model: Type[BaseModel],

src/intersect_sdk/capability/universal_capability/universal_capability.py

@@ -6,11 +6,9 @@
 """
 
 import datetime
-from cryptography.hazmat.primitives.asymmetric import rsa
-from cryptography.hazmat.primitives import serialization
 import os
 import time
-from typing import final
+from typing import Dict, final
 
 import psutil
 
@@ -39,24 +37,6 @@ def __init__(self) -> None:  # noqa: D107
         self.process = psutil.Process(os.getpid())
         """psutil.Process caches most functions it calls after it calls the function once, so just save the object itself"""
 
-        # Generate a key pair for encryption
-        self._private_key: rsa.RSAPrivateKey = rsa.generate_private_key(
-            public_exponent=65537, key_size=2048
-        )
-        self._public_key = self._private_key.public_key()
-
-        # Get the PEM encoded public key
-        self._public_key_pem = self._public_key.public_bytes(
-            encoding=serialization.Encoding.PEM,
-            format=serialization.PublicFormat.SubjectPublicKeyInfo,
-        ).decode()
-
-        self._private_key.private_bytes(
-            encoding=serialization.Encoding.PEM,
-            format=serialization.PrivateFormat.TraditionalOpenSSL,
-            encryption_algorithm=serialization.NoEncryption(),
-        )
-
 
     @intersect_status
     def system_capability(self) -> IntersectCoreStatus:

src/intersect_sdk/service.py

@@ -18,6 +18,8 @@
 
 import time
 from collections import defaultdict
+from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.hazmat.primitives import serialization
 from threading import Lock
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Literal
@@ -31,6 +33,8 @@
     ControlPlaneManager,
 )
 from ._internal.data_plane.data_plane_manager import DataPlaneManager
+from ._internal.encryption import intersect_payload_decrypt, intersect_payload_encrypt
+from ._internal.encryption.models import IntersectEncryptionPublicKey
 from ._internal.exceptions import IntersectApplicationError, IntersectError
 from ._internal.generic_serializer import GENERIC_MESSAGE_SERIALIZER
 from ._internal.interfaces import IntersectEventObserver
@@ -308,6 +312,24 @@ def __init__(
             self._client_channel_name, {self._handle_client_message}, persist=True
         )
 
+        # Generate a key pair for encryption
+        self._private_key: rsa.RSAPrivateKey = rsa.generate_private_key(
+            public_exponent=65537, key_size=2048
+        )
+        self._public_key = self._private_key.public_key()
+
+        # Get the PEM encoded public key
+        self._public_key_pem = self._public_key.public_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo,
+        ).decode()
+
+        self._private_key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.TraditionalOpenSSL,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+
     def _get_capability(self, target: str) -> IntersectBaseCapabilityImplementation | None:
         for cap in self.capabilities:
             if cap.intersect_sdk_capability_name == target:
@@ -825,7 +847,11 @@ def _handle_service_message_inner(
                 )
             match headers.encryption_scheme:
                 case 'RSA':
-                    # TODO - decrypt request_params here
+                    request_params = intersect_payload_decrypt(
+                        rsa_private_key=self._private_key,
+                        encrypted_payload=request_params,
+                        model=bytes,
+                    )
                     pass
                 case _:
                     pass
@@ -922,7 +948,11 @@ def _handle_client_message(
                         # error messages should never be encrypted
                         match headers.encryption_scheme:
                             case 'RSA':
-                                # TODO - decrypt message here
+                                msg_payload = intersect_payload_decrypt(
+                                    rsa_private_key=self._private_key,
+                                    encrypted_payload=msg_payload,
+                                    model=bytes,
+                                )
                                 pass
                             case _:
                                 pass
@@ -954,7 +984,7 @@ def _send_client_message(self, request_id: UUID, params: IntersectDirectMessageP
         if params.content_type == 'application/json':
             request = GENERIC_MESSAGE_SERIALIZER.dump_json(params.payload, warnings=False)
         else:
-            if not isinstance(params.content_type, bytes):
+            if not isinstance(params.payload, bytes):
                 logger.error(
                     'service-to-service message must be bytes if content-type is not application/json'
                 )
@@ -963,8 +993,55 @@ def _send_client_message(self, request_id: UUID, params: IntersectDirectMessageP
 
         match params.encryption_scheme:
             case 'RSA':
-                # TODO encrypt message
-                pass
+                # Get the public key from the destination service
+                # The destination service has the public key through its universal capability
+                public_key_request = IntersectDirectMessageParams(
+                    destination=params.destination,
+                    operation='intersect_sdk.get_public_key',
+                    payload=None,
+                )
+                
+                # Create the external request to fetch the public key
+                public_key_request_id = self.create_external_request(
+                    public_key_request,
+                    response_handler=None,
+                    timeout=30.0,  # shorter timeout for key fetching
+                )
+                
+                # Poll for the response with timeout
+                start_time = time.time()
+                timeout = 30.0
+                poll_interval = 0.1
+                
+                public_key_payload = None
+                while time.time() - start_time < timeout:
+                    extreq = self._get_external_request(public_key_request_id)
+                    if extreq and extreq.request_state == 'received':
+                        public_key_payload = extreq.response_payload
+                        extreq.request_state = 'finalized'
+                        break
+                    time.sleep(poll_interval)
+                
+                if public_key_payload is None:
+                    logger.error(
+                        f'Failed to retrieve public key from {params.destination} within timeout'
+                    )
+                    return False
+                
+                # public_key_payload should be an IntersectEncryptionPublicKey instance (as a dict)
+                # Parse it if it's a dict, otherwise use it as-is
+                if isinstance(public_key_payload, dict):
+                    key_payload = IntersectEncryptionPublicKey(**public_key_payload)
+                else:
+                    key_payload = public_key_payload
+                
+                # Now encrypt the request using the retrieved public key
+                # Convert bytes back to string if needed (since intersect_payload_encrypt expects a string)
+                unencrypted_string = request if isinstance(request, str) else request.decode()
+                request = intersect_payload_encrypt(
+                    key_payload=key_payload,
+                    unencrypted_model=unencrypted_string,
+                )
             case _:
                 pass
 

tests/fixtures/example_schema.json

@@ -9,7 +9,45 @@
   "defaultContentType": "application/json",
   "capabilities": {
     "intersect_sdk": {
-      "endpoints": {},
+      "endpoints": {
+        "get_public_key": {
+          "publish": {
+            "message": {
+              "schemaFormat": "application/vnd.aai.asyncapi+json;version=2.6.0",
+              "contentType": "application/json",
+              "encryption_schemes": [
+                "NONE",
+                "RSA"
+              ],
+              "traits": {
+                "$ref": "#/components/messageTraits/commonHeaders"
+              },
+              "payload": {
+                "additionalProperties": {
+                  "type": "string"
+                },
+                "type": "object",
+                "title": "get_public_key"
+              }
+            },
+            "description": "Returns the public key for clients / services to use for encryption"
+          },
+          "subscribe": {
+            "message": {
+              "schemaFormat": "application/vnd.aai.asyncapi+json;version=2.6.0",
+              "contentType": "application/json",
+              "encryption_schemes": [
+                "NONE",
+                "RSA"
+              ],
+              "traits": {
+                "$ref": "#/components/messageTraits/commonHeaders"
+              }
+            },
+            "description": "Returns the public key for clients / services to use for encryption"
+          }
+        }
+      },
       "events": {},
       "status": {
         "$ref": "#/components/schemas/IntersectCoreStatus"
@@ -1473,4 +1511,4 @@
       }
     }
   }
-}
+}