Added encryption to client + fix to service

Author: marshallmcdonnell

src/intersect_sdk/client.py

@@ -12,14 +12,21 @@
 
 from __future__ import annotations
 
+from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.hazmat.primitives import serialization
 import time
 from typing import TYPE_CHECKING
 from uuid import uuid4
 
 from pydantic import ValidationError
 from typing_extensions import Self, final
 
-from intersect_sdk._internal.generic_serializer import GENERIC_MESSAGE_SERIALIZER
+from ._internal.encryption.models.public_key import IntersectEncryptionPublicKey
+from ._internal.encryption import (
+    intersect_payload_decrypt,
+    intersect_payload_encrypt,
+)
+from ._internal.generic_serializer import GENERIC_MESSAGE_SERIALIZER
 
 from ._internal.control_plane.control_plane_manager import (
     ControlPlaneManager,
@@ -149,6 +156,24 @@ def __init__(
 
         self._campaign_id = uuid4()
 
+        # 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(),
+        )
+
     @final
     def startup(self) -> Self:
         """This function connects the client to all INTERSECT systems.
@@ -225,6 +250,69 @@ def considered_unrecoverable(self) -> bool:
         """
         return self._control_plane_manager.considered_unrecoverable()
 
+    def _fetch_service_public_key(self, destination: str) -> str | None:
+        """Fetch the public key from a destination service via intersect_sdk.get_public_key.
+        
+        Args:
+            destination: The hierarchy string of the destination service
+            
+        Returns:
+            The PEM-encoded public key from the service, or None if fetching failed
+        """
+        # Create a temporary request ID for tracking the public key request
+        public_key_request_id = uuid4()
+        
+        # Build and send the public key request
+        headers = create_userspace_message_headers(
+            source=self._hierarchy.hierarchy_string('.'),
+            destination=destination,
+            operation_id='intersect_sdk.get_public_key',
+            campaign_id=uuid4(),
+            request_id=public_key_request_id,
+            encryption_scheme='NONE',
+        )
+        
+        # Send the request to get the public key
+        request_channel = f'{destination.replace(".", "/")}/request'
+        self._control_plane_manager.publish_message(
+            request_channel, b'null', 'application/json', headers, persist=False
+        )
+        
+        # Poll for the response with timeout
+        start_time = time.time()
+        timeout = 30.0
+        poll_interval = 0.1
+        response_data = None
+        
+        while time.time() - start_time < timeout:
+            # Check if we have received a response for this request
+            # We'll store responses temporarily in a dict keyed by request_id
+            if hasattr(self, '_public_key_responses') and str(public_key_request_id) in self._public_key_responses:
+                response_data = self._public_key_responses.pop(str(public_key_request_id))
+                break
+            time.sleep(poll_interval)
+        
+        if response_data is None:
+            logger.error(f'Failed to retrieve public key from {destination} within timeout')
+            return None
+        
+        # Extract the public key from the response
+        try:
+            if isinstance(response_data, dict):
+                public_key = response_data.get('public_key')
+            else:
+                # If it's already parsed as IntersectEncryptionPublicKey
+                public_key = response_data.public_key if hasattr(response_data, 'public_key') else None
+            
+            if public_key is None:
+                logger.error(f'No public_key field in response from {destination}')
+                return None
+                
+            return public_key
+        except (KeyError, AttributeError) as e:
+            logger.error(f'Failed to extract public key from response: {e}')
+            return None
+
     def _handle_userspace_message(
         self, payload: bytes, content_type: str, raw_headers: dict[str, str]
     ) -> None:
@@ -265,11 +353,30 @@ def _handle_userspace_message(
             request_params = self._data_plane_manager.incoming_message_data_handler(
                 payload, headers.data_handler
             )
+            
+            # Check if this is a response to a public key request (before processing as user message)
+            if headers.operation_id == 'intersect_sdk.get_public_key':
+                if not hasattr(self, '_public_key_responses'):
+                    self._public_key_responses = {}
+                self._public_key_responses[str(headers.message_id)] = request_params
+                return
+            
             if not headers.has_error:
                 match headers.encryption_scheme:
                     case 'RSA':
-                        # TODO - decrypt and reassign request_params here
-                        pass
+                        from pydantic import BaseModel
+                        
+                        # Create a simple wrapper model for bytes deserialization
+                        class _BytesWrapper(BaseModel):
+                            data: str = ""
+                        
+                        decrypted = intersect_payload_decrypt(
+                            rsa_private_key=self._private_key,
+                            encrypted_payload=request_params,
+                            model=_BytesWrapper,
+                        )
+                        # Extract the decrypted data and return as bytes
+                        request_params = decrypted.model.data.encode()
                     case _:
                         pass
             if content_type == 'application/json':
@@ -428,8 +535,20 @@ def _send_userspace_message(self, params: IntersectDirectMessageParams) -> None:
         # TWO: encrypt message
         match params.encryption_scheme:
             case 'RSA':
-                # TODO reassign serialized_msg here to encrypted value
-                pass
+                # Fetch the destination service's public key
+                public_key_pem = self._fetch_service_public_key(params.destination)
+                if public_key_pem is None:
+                    logger.error(f'Failed to fetch public key from {params.destination}')
+                    return
+                
+                # Convert bytes to string if needed (since intersect_payload_encrypt expects a string)
+                unencrypted_string = serialized_msg if isinstance(serialized_msg, str) else serialized_msg.decode()
+                serialized_msg = intersect_payload_encrypt(
+                    key_payload=IntersectEncryptionPublicKey(
+                        public_key=public_key_pem
+                    ),
+                    unencrypted_model=unencrypted_string,
+                )
             case _:
                 pass
 

src/intersect_sdk/service.py

@@ -1039,7 +1039,9 @@ def _send_client_message(self, request_id: UUID, params: IntersectDirectMessageP
                 # 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,
+                    key_payload=IntersectEncryptionPublicKey(
+                        public_key=key_payload.public_key,
+                    ),
                     unencrypted_model=unencrypted_string,
                 )
             case _: