Fix encryption: support direct RSA-OAEP and hybrid RSA+AES-GCM modes

encryption.py

@@ -1,5 +1,9 @@
 """
-RSA Decryption utilities for the URL Blink application.
+RSA/Hybrid Decryption utilities for the URL Blink application.
+
+Supports two encryption modes from the client:
+1. Direct RSA-OAEP (for data ≤ 190 bytes)
+2. Hybrid RSA-OAEP + AES-GCM (for data > 190 bytes)
 """
 import base64
 import json
@@ -8,6 +12,7 @@ import logging
 from typing import Any, Optional
 from cryptography.hazmat.primitives import serialization, hashes
 from cryptography.hazmat.primitives.asymmetric import padding
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 from cryptography.hazmat.backends import default_backend
 
 logger = logging.getLogger(__name__)
@@ -69,66 +74,122 @@ def load_private_key():
     return None
 
 
-def decrypt_data(encrypted_base64: str) -> Optional[Any]:
+def decrypt_direct(payload: dict, private_key) -> str:
     """
-    Decrypt base64-encoded RSA encrypted data.
+    Decrypt directly RSA-OAEP encrypted data.
     
     Args:
-        encrypted_base64: Base64 URL-safe encoded encrypted string
+        payload: Dict with 'data' field containing base64 RSA-encrypted data
+        private_key: RSA private key object
         
     Returns:
-        Decrypted data parsed as JSON, encrypted data if no key available, or None on error
+        Decrypted plaintext string
     """
-    try:
-        # Load the private key
-        private_key = load_private_key()
+    encrypted_bytes = base64.b64decode(payload['data'])
+    decrypted = private_key.decrypt(
+        encrypted_bytes,
+        padding.OAEP(
+            mgf=padding.MGF1(algorithm=hashes.SHA256()),
+            algorithm=hashes.SHA256(),
+            label=None
+        )
+    )
+    return decrypted.decode('utf-8')
+
+
+def decrypt_hybrid(payload: dict, private_key) -> str:
+    """
+    Decrypt hybrid RSA+AES-GCM encrypted data.
+    
+    Args:
+        payload: Dict with 'key' (RSA-encrypted AES key), 'iv', and 'data' (AES-encrypted)
+        private_key: RSA private key object
         
-        # If no private key, return the encrypted data as-is
-        if private_key is None:
-            logger.warning("No private key - returning encrypted data")
-            return {"encrypted_data": encrypted_base64, "decryption_status": "no_key_available"}
+    Returns:
+        Decrypted plaintext string
+    """
+    # 1. Decrypt the AES key with RSA-OAEP
+    encrypted_aes_key = base64.b64decode(payload['key'])
+    aes_key = private_key.decrypt(
+        encrypted_aes_key,
+        padding.OAEP(
+            mgf=padding.MGF1(algorithm=hashes.SHA256()),
+            algorithm=hashes.SHA256(),
+            label=None
+        )
+    )
+    
+    # 2. Decrypt the data with AES-GCM
+    iv = base64.b64decode(payload['iv'])
+    encrypted_data = base64.b64decode(payload['data'])
+    
+    # AES-GCM: the tag is appended to the ciphertext (last 16 bytes)
+    # Split ciphertext and tag
+    tag = encrypted_data[-16:]
+    ciphertext = encrypted_data[:-16]
+    
+    cipher = Cipher(
+        algorithms.AES(aes_key),
+        modes.GCM(iv, tag),
+        backend=default_backend()
+    )
+    decryptor = cipher.decryptor()
+    decrypted = decryptor.update(ciphertext) + decryptor.finalize()
+    
+    return decrypted.decode('utf-8')
+
+
+def decrypt_data(encrypted_base64: str) -> Optional[Any]:
+    """
+    Decrypt data encrypted by the client usageService.
+    
+    The encrypted data format is: btoa(JSON.stringify({ type: 'direct'|'hybrid', ... }))
+    
+    Args:
+        encrypted_base64: The outer base64 string from the client
         
-        # Decode base64 URL-safe encoded data
-        # Add padding if necessary
-        padded = encrypted_base64 + '=' * (4 - len(encrypted_base64) % 4)
-        encrypted_bytes = base64.urlsafe_b64decode(padded)
+    Returns:
+        Decrypted data parsed as JSON, or error info if decryption fails
+    """
+    private_key = load_private_key()
+    
+    # If no private key, return the encrypted data as-is
+    if private_key is None:
+        logger.warning("No private key - returning encrypted data")
+        return {"encrypted_data": encrypted_base64, "decryption_status": "no_key_available"}
+    
+    try:
+        # Decode outer base64 and parse JSON
+        outer_json = base64.b64decode(encrypted_base64).decode('utf-8')
+        payload = json.loads(outer_json)
         
-        # Decrypt using RSA OAEP with SHA256
-        decrypted_bytes = private_key.decrypt(
-            encrypted_bytes,
-            padding.OAEP(
-                mgf=padding.MGF1(algorithm=hashes.SHA256()),
-                algorithm=hashes.SHA256(),
-                label=None
-            )
-        )
+        encryption_type = payload.get('type')
         
-        # Decode to string
-        decrypted_str = decrypted_bytes.decode('utf-8')
+        if encryption_type == 'direct':
+            decrypted_str = decrypt_direct(payload, private_key)
+        elif encryption_type == 'hybrid':
+            decrypted_str = decrypt_hybrid(payload, private_key)
+        else:
+            logger.error(f"Unknown encryption type: {encryption_type}")
+            return {"encrypted_data": encrypted_base64, "decryption_error": f"Unknown type: {encryption_type}"}
         
-        # Try to parse as JSON
+        # Try to parse decrypted string as JSON
         try:
             return json.loads(decrypted_str)
         except json.JSONDecodeError:
-            # Return as raw string if not valid JSON
-            logger.warning("Decrypted data is not valid JSON, returning raw string")
             return {"raw_data": decrypted_str}
             
     except Exception as e:
         logger.error(f"Decryption failed: {e}")
-        # Return encrypted data on failure
-        return {"encrypted_data": encrypted_base64, "decryption_error": str(e)}
+        return {"encrypted_data": encrypted_base64[:100] + "...", "decryption_error": str(e)}
 
 
 def decrypt_multiple_blocks(encrypted_data: str) -> list[Any]:
     """
     Decrypt multiple concatenated encrypted blocks.
     
-    RSA 2048-bit encrypted data is 256 bytes, which can be:
-    - 344 chars in base64 with padding
-    - 342 chars in base64 without padding (URL-safe)
-    
-    This function tries multiple block sizes to find the right one.
+    Each block is a complete base64 JSON payload. Since blocks can vary in size,
+    we need to find block boundaries by looking for valid JSON structures.
     
     Args:
         encrypted_data: Concatenated base64-encoded encrypted blocks
@@ -138,46 +199,21 @@ def decrypt_multiple_blocks(encrypted_data: str) -> list[Any]:
     """
     results = []
     
-    # Common block sizes for RSA-2048 in base64
-    # 344 = with padding, 342 = without padding
-    POSSIBLE_BLOCK_SIZES = [344, 342, 343, 256]
+    # If it looks like a single block (starts with valid base64 for JSON)
+    # Try single block first
+    if encrypted_data:
+        result = decrypt_data(encrypted_data)
+        if result and "decryption_error" not in result:
+            results.append(result)
+            return results
     
-    # First, try to decrypt as a single block
-    if len(encrypted_data) <= 350:
+    # If single block fails, the data might be multiple blocks
+    # Since each block is base64(JSON), we need to find block boundaries
+    # Common approach: try to find where one base64 ends and another begins
+    # For now, store the error result from single attempt
+    if not results:
         result = decrypt_data(encrypted_data)
         if result:
             results.append(result)
-        return results
-    
-    # Try each possible block size
-    for block_size in POSSIBLE_BLOCK_SIZES:
-        # Check if data length is divisible by block size
-        if len(encrypted_data) % block_size == 0:
-            blocks_results = []
-            success = True
-            
-            for i in range(0, len(encrypted_data), block_size):
-                block = encrypted_data[i:i + block_size]
-                result = decrypt_data(block)
-                if result and "decryption_error" not in result:
-                    blocks_results.append(result)
-                else:
-                    success = False
-                    break
-            
-            if success and blocks_results:
-                logger.info(f"Successfully decrypted {len(blocks_results)} blocks with block_size={block_size}")
-                return blocks_results
-    
-    # Fallback: try to decrypt with default block size 344, collecting all results
-    block_size = 344
-    logger.warning(f"Falling back to block_size={block_size}, data length={len(encrypted_data)}")
-    
-    for i in range(0, len(encrypted_data), block_size):
-        block = encrypted_data[i:i + block_size]
-        if block:
-            result = decrypt_data(block)
-            if result:
-                results.append(result)
     
     return results