Adds encryption module to _internal

Author: marshallmcdonnell

src/intersect_sdk/_internal/encryption/__init__.py

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

src/intersect_sdk/_internal/encryption/aes_cypher.py

@@ -0,0 +1,40 @@
+"""
+AESCipher class for intersect_sdk._internal.encryption.aes_cypher
+"""
+import base64
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+from cryptography.hazmat.primitives.padding import PKCS7
+from cryptography.hazmat.backends import default_backend
+
+
+class AESCipher:
+    def __init__(self, key: bytes, initialization_vector: bytes):
+        if len(key) * 8 != 256:
+            raise ValueError("Invalid key size (must be 256 bit)")
+        self._key: bytes = key
+        # No need to encrypt initialization vectors
+        # https://cryptography.io/en/latest/hazmat/primitives/symmetric-encryption/#cryptography.hazmat.primitives.ciphers.modes.CBC
+        self._initial_vector: bytes = initialization_vector
+        self._cipher: Cipher = Cipher(
+            algorithms.AES(self._key),
+            mode=modes.CBC(self._initial_vector),
+            backend=default_backend(),
+        )
+
+    def encrypt(self, plaintext: bytes) -> bytes:
+        encryptor = self._cipher.encryptor()
+        bytes_to_bits = 8
+        padder = PKCS7(len(self._initial_vector) * bytes_to_bits).padder()
+        padded_data = padder.update(plaintext) + padder.finalize()
+        ciphertext = encryptor.update(padded_data) + encryptor.finalize()
+        encodedciphertext = base64.b64encode(ciphertext)
+        return encodedciphertext
+
+    def decrypt(self, ciphertext: bytes) -> bytes:
+        decryptor = self._cipher.decryptor()
+        decodedciphertext = base64.b64decode(ciphertext)
+        padded_data = decryptor.update(decodedciphertext) + decryptor.finalize()
+        bytes_to_bits = 8
+        unpadder = PKCS7(len(self._initial_vector) * bytes_to_bits).unpadder()
+        plaintext = unpadder.update(padded_data) + unpadder.finalize()
+        return plaintext

src/intersect_sdk/_internal/encryption/client_encryption.py

@@ -0,0 +1,57 @@
+"""
+Client encryption function for intersect_sdk._internal.encryption.client_encryption
+RSA asymmetric encryption and AES symmetric encryption
+"""
+# RSA asymmetric encryption and AES symmetric encryption
+import base64
+from cryptography.hazmat.primitives.asymmetric import padding
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.backends import default_backend
+import os
+
+
+from .aes_cypher import AESCipher
+from .models import IntersectEncryptedPayload, IntersectEncryptionPublicKey
+from ..logger import logger
+
+def intersect_client_encryption(
+    key_payload: IntersectEncryptionPublicKey,
+    unencrypted_model: str,
+) -> IntersectEncryptedPayload:
+    # Decode and deserialize the public key for asymmetric encrypt
+    public_rsa_key = serialization.load_pem_public_key(
+        key_payload.public_key.encode(),
+        backend=default_backend(),
+    )
+
+    # Setup AES
+    aes_key: bytes = os.urandom(32)
+    aes_initialization_vector: bytes = os.urandom(16)
+    cipher = AESCipher(
+        key=aes_key,
+        initialization_vector=aes_initialization_vector,
+    )
+
+    # Encrypt using AES
+    logger.info("AES encrypting payload...")
+    encrypted_data = cipher.encrypt(unencrypted_model.encode())
+    logger.info("AES encrypted the payload!")
+
+    # Asymmetric RSA encrypts AES symmetric key using RSA public key
+    logger.info("RSA encrypting AES key...")
+    encrypted_aes_key: bytes = public_rsa_key.encrypt(
+        aes_key,
+        padding.OAEP(
+            mgf=padding.MGF1(algorithm=hashes.SHA256()),
+            algorithm=hashes.SHA256(),
+            label=None,
+        ),
+    )
+    logger.info("RSA encrypted AES key!")
+    logger.info("Done encrypting!")
+
+    return IntersectEncryptedPayload(
+        key=base64.b64encode(encrypted_aes_key).decode(),
+        initial_vector=base64.b64encode(aes_initialization_vector).decode(),
+        data=base64.b64encode(encrypted_data).decode(),
+    )

src/intersect_sdk/_internal/encryption/models/__init__.py

@@ -0,0 +1,9 @@
+from .decrypted_payload import IntersectDecryptedPayload
+from .encrypted_payload import IntersectEncryptedPayload
+from .public_key import IntersectEncryptionPublicKey
+
+__all__ = (
+    "IntersectDecryptedPayload",
+    "IntersectEncryptedPayload",
+    "IntersectEncryptionPublicKey",
+)

src/intersect_sdk/_internal/encryption/models/decrypted_payload.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+class IntersectDecryptedPayload(BaseModel):
+    model: BaseModel
+    aes_key: bytes
+    aes_initialization_vector: bytes

src/intersect_sdk/_internal/encryption/models/encrypted_payload.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel
+
+
+class IntersectEncryptedPayload(BaseModel):
+    key: str
+    initial_vector: str
+    data: str

src/intersect_sdk/_internal/encryption/models/public_key.py

@@ -0,0 +1,13 @@
+"""Definitions supporting the 'core status' functionality of the core capability."""
+
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+from typing import Annotated
+
+
+class IntersectEncryptionPublicKey(BaseModel):
+    """Public key information for encryption within the INTERSECT-SDK Service."""
+
+    public_key: Annotated[str, Field(title='Public Key PEM')]
+    """The PEM encoded public key for asymmetric encryption."""

src/intersect_sdk/_internal/encryption/service_decryption.py

@@ -0,0 +1,57 @@
+"""
+Service decryption function for intersect_sdk._internal.encryption.service_decryption
+RSA asymmetric encryption and AES symmetric encryption
+"""
+
+import json
+from typing import Dict, Type
+from cryptography.hazmat.primitives.asymmetric import rsa, padding
+from cryptography.hazmat.primitives import hashes
+import base64
+
+from pydantic import BaseModel
+
+from .aes_cypher import AESCipher
+from .models import IntersectEncryptedPayload, IntersectDecryptedPayload
+from ..logger import logger
+
+
+def intersect_service_decryption(
+    rsa_private_key: rsa.RSAPrivateKey,
+    encrypted_payload: IntersectEncryptedPayload,
+    model: Type[BaseModel],
+) -> IntersectDecryptedPayload:
+    # Decode base64 encoded values from payload
+    encrypted_aes_key = base64.b64decode(encrypted_payload.key.encode())
+    initial_vector = base64.b64decode(encrypted_payload.initial_vector.encode())
+    encrypted_data = base64.b64decode(encrypted_payload.data.encode())
+
+    # Decrypt AES key using RSA
+    logger.info("RSA decrypting AES key...")
+    decrypted_aes_key: bytes = rsa_private_key.decrypt(
+        encrypted_aes_key,
+        padding.OAEP(
+            mgf=padding.MGF1(algorithm=hashes.SHA256()),
+            algorithm=hashes.SHA256(),
+            label=None,
+        ),
+    )
+    logger.info("RSA decrypted AES key!")
+
+    # Encrypt using AES
+    logger.info("AES decrypting payload...")
+    cipher = AESCipher(
+        key=decrypted_aes_key,
+        initialization_vector=initial_vector,
+    )
+
+    decrypted_payload: bytes = cipher.decrypt(encrypted_data)
+    unencrypted_payload = decrypted_payload.decode()
+    logger.info("AES decrypted payload!")
+
+    # Cast unencrypted payload to model and return AES key and IV for possible re-use
+    return IntersectDecryptedPayload(
+        model=model(**json.loads(unencrypted_payload)),
+        aes_key=decrypted_aes_key,
+        aes_initialization_vector=initial_vector,
+    )