Merge pull request #11 from bitovi/more-pr-feedback

More pr feedback

Author: phillipskevin

encryption_jwt/codec_server.py

@@ -6,6 +6,7 @@
 import requests
 from aiohttp import hdrs, web
 from google.protobuf import json_format
+from jwt import PyJWK
 from jwt.algorithms import RSAAlgorithm
 from temporalio.api.cloud.cloudservice.v1 import GetUsersRequest
 from temporalio.api.common.v1 import Payloads
@@ -18,9 +19,9 @@
 
 TEMPORAL_CLIENT_CLOUD_API_VERSION = "2024-05-13-00"
 
-temporal_ops_address = "saas-api.tmprl.cloud:443"
-if os.environ.get("TEMPORAL_OPS_ADDRESS"):
-    temporal_ops_address = os.environ.get("TEMPORAL_OPS_ADDRESS")
+temporal_ops_address = (
+    os.environ.get("TEMPORAL_OPS_ADDRESS") or "saas-api.tmprl.cloud:443"
+)
 
 
 def build_codec_server() -> web.Application:
@@ -76,8 +77,8 @@ async def decryption_authorized(email: str, namespace: str) -> bool:
 
     def make_handler(fn: str):
         async def handler(req: web.Request):
-            namespace = req.headers.get("x-namespace")
-            auth_header = req.headers.get("Authorization")
+            namespace = req.headers.get("x-namespace") or "default"
+            auth_header = req.headers.get("Authorization") or ""
             _bearer, encoded = auth_header.split(" ")
 
             # Extract the kid from the Auth header
@@ -90,20 +91,20 @@ async def handler(req: web.Request):
             jwks = requests.get(jwks_url).json()
 
             # Extract Temporal Cloud's public key
-            public_key = None
+            pyjwk = None
             for key in jwks["keys"]:
                 if key["kid"] == kid:
                     # Convert JWKS key to PEM format
-                    public_key = RSAAlgorithm.from_jwk(key)
+                    pyjwk = PyJWK.from_dict(key)
                     break
 
-            if public_key is None:
+            if pyjwk is None:
                 raise ValueError("Public key not found in JWKS")
 
             # Decode the jwt, verifying against Temporal Cloud's public key
             decoded = jwt.decode(
                 encoded,
-                public_key,
+                pyjwk.key,
                 algorithms=[algorithm],
                 audience=[
                     "https://saas-api.tmprl.cloud",
@@ -156,7 +157,7 @@ async def handler(req: web.Request):
         ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
         ssl_context.check_hostname = False
         ssl_context.load_cert_chain(
-            os.environ.get("SSL_PEM"), os.environ.get("SSL_KEY")
+            os.environ.get("SSL_PEM") or "", os.environ.get("SSL_KEY") or ""
         )
 
     web.run_app(

encryption_jwt/encryptor.py

@@ -1,3 +1,4 @@
+import asyncio
 import base64
 import logging
 import os
@@ -38,16 +39,20 @@ async def encrypt(self, data: bytes) -> tuple[bytes, bytes]:
 
         nonce = os.urandom(12)
         encryptor = AESGCM(data_key_plaintext)
-        return nonce + encryptor.encrypt(nonce, data, None), base64.b64encode(
-            data_key_encrypted
+        encrypted = asyncio.get_running_loop().run_in_executor(
+            None, encryptor.encrypt, nonce, data, None
         )
+        return nonce + await encrypted, base64.b64encode(data_key_encrypted)
 
     async def decrypt(self, data_key_encrypted_base64, data: bytes) -> bytes:
         """Encrypt data using a key from KMS."""
         data_key_encrypted = base64.b64decode(data_key_encrypted_base64)
         data_key_plaintext = await self.__decrypt_data_key(data_key_encrypted)
         encryptor = AESGCM(data_key_plaintext)
-        return encryptor.decrypt(data[:12], data[12:], None)
+        decrypted = await asyncio.get_running_loop().run_in_executor(
+            None, encryptor.decrypt, data[:12], data[12:], None
+        )
+        return decrypted
 
     async def __create_data_key(self, namespace: str):
         """Get a set of keys from AWS KMS that can be used to encrypt data."""