From c533ed6b6c56051807d9ed624f027e7b64989f0e Mon Sep 17 00:00:00 2001
From: Todd Wildey <twwildey@amazon.com>
Date: Wed, 29 May 2024 15:49:10 -0700
Subject: [PATCH] Updating `CryptographyAESKey::encrypt` to generate 96 bit IVs
 for GCM block cipher mode to adhere to the RFC for JWA in
 `jose/backends/cryptography_backend.py`

See https://www.rfc-editor.org/rfc/rfc7518.html#section-5.3 for the official RFC requirements for JWA

See https://github.com/panva/jose/issues/678 for related discussion on this issue
---
 jose/backends/_asn1.py                | 1 +
 jose/backends/cryptography_backend.py | 5 ++++-
 tests/test_asn1.py                    | 1 +
 tests/test_backends.py                | 1 +
 4 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/jose/backends/_asn1.py b/jose/backends/_asn1.py
index af5fa8bc..87e3df1b 100644
--- a/jose/backends/_asn1.py
+++ b/jose/backends/_asn1.py
@@ -2,6 +2,7 @@
 
 Required by rsa_backend but not cryptography_backend.
 """
+
 from pyasn1.codec.der import decoder, encoder
 from pyasn1.type import namedtype, univ
 
diff --git a/jose/backends/cryptography_backend.py b/jose/backends/cryptography_backend.py
index abd24260..945349b8 100644
--- a/jose/backends/cryptography_backend.py
+++ b/jose/backends/cryptography_backend.py
@@ -439,6 +439,8 @@ class CryptographyAESKey(Key):
         ALGORITHMS.A256KW: None,
     }
 
+    IV_BYTE_LENGTH_MODE_MAP = {"CBC": algorithms.AES.block_size // 8, "GCM": 96 // 8}
+
     def __init__(self, key, algorithm):
         if algorithm not in ALGORITHMS.AES:
             raise JWKError("%s is not a valid AES algorithm" % algorithm)
@@ -468,7 +470,8 @@ def to_dict(self):
     def encrypt(self, plain_text, aad=None):
         plain_text = ensure_binary(plain_text)
         try:
-            iv = get_random_bytes(algorithms.AES.block_size // 8)
+            iv_byte_length = self.IV_BYTE_LENGTH_MODE_MAP.get(self._mode.name, algorithms.AES.block_size)
+            iv = get_random_bytes(iv_byte_length)
             mode = self._mode(iv)
             if mode.name == "GCM":
                 cipher = aead.AESGCM(self._key)
diff --git a/tests/test_asn1.py b/tests/test_asn1.py
index 64f2d4b1..6e1b1039 100644
--- a/tests/test_asn1.py
+++ b/tests/test_asn1.py
@@ -1,4 +1,5 @@
 """Tests for ``jose.backends._asn1``."""
+
 import base64
 
 import pytest
diff --git a/tests/test_backends.py b/tests/test_backends.py
index 10ef390b..4ce71a7d 100644
--- a/tests/test_backends.py
+++ b/tests/test_backends.py
@@ -1,4 +1,5 @@
 """Test the default import handling."""
+
 try:
     from jose.backends.rsa_backend import RSAKey as PurePythonRSAKey
 except ImportError: