[CSL3-2519] Fix for block size and encryption with JWK in JOSE (#567)

* Fix for block size in in Jose * Fix for block size in in Jose * Fix for block size in in Jose * Fix build issue * Fix build issue * Fix build issue * Fix build issue * Fix JWK * Remove duplicating test case --------- Co-authored-by: Viacheslav Rud <[email protected]>
verygoodsecurity · Jan 9, 2025 · c04e6a8 · c04e6a8
1 parent 210d7b8
commit c04e6a8
Show file tree

Hide file tree

Showing 5 changed files with 159 additions and 53 deletions.
diff --git a/larky/src/main/resources/vendor/jose/backends/pycrypto_backend.star b/larky/src/main/resources/vendor/jose/backends/pycrypto_backend.star
@@ -157,38 +157,6 @@ def RSAKey(key, algorithm):
                                SHA512=SHA512,
                                SHA1=SHA1)
 
-    def __init__(key, algorithm):
-        if not operator.contains(ALGORITHMS.RSA, algorithm):
-            return Error("JWKError: hash_alg: %s is not a valid hash algorithm" % algorithm).unwrap()
-
-        self.hash_alg = {
-            ALGORITHMS.RS256: self.SHA256,
-            ALGORITHMS.RS384: self.SHA384,
-            ALGORITHMS.RS512: self.SHA512,
-            ALGORITHMS.RSA1_5: self.SHA1,
-            ALGORITHMS.RSA_OAEP: self.SHA1,
-            ALGORITHMS.RSA_OAEP_256: self.SHA256,
-        }.get(algorithm)
-        self._algorithm = algorithm
-
-        if types.is_instance(key, _RSAKey):
-            self.prepared_key = key
-            return self
-
-        if types.is_dict(key):
-            self._process_jwk(key)
-            return self
-
-        if types.is_string(key):
-            key = codecs.encode(key, encoding='utf-8')
-
-        if types.is_bytelike(key):
-            self.prepared_key = RSA.importKey(key)
-            return self
-
-        return Error("JWKError: Unable to parse an RSA_JWK from key: %s" % key).unwrap()
-    self = __init__(key, algorithm)
-
     def _process_jwk(jwk_dict):
         if not jwk_dict.get('kty') == 'RSA':
             return Error("JWKError: Incorrect key type. Expected: 'RSA', Received: %s" % jwk_dict.get('kty'))
@@ -227,6 +195,38 @@ def RSAKey(key, algorithm):
         return self.prepared_key
     self._process_jwk = _process_jwk
 
+    def __init__(key, algorithm):
+        if not operator.contains(ALGORITHMS.RSA, algorithm):
+            return Error("JWKError: hash_alg: %s is not a valid hash algorithm" % algorithm).unwrap()
+
+        self.hash_alg = {
+            ALGORITHMS.RS256: self.SHA256,
+            ALGORITHMS.RS384: self.SHA384,
+            ALGORITHMS.RS512: self.SHA512,
+            ALGORITHMS.RSA1_5: self.SHA1,
+            ALGORITHMS.RSA_OAEP: self.SHA1,
+            ALGORITHMS.RSA_OAEP_256: self.SHA256,
+        }.get(algorithm)
+        self._algorithm = algorithm
+
+        if types.is_instance(key, _RSAKey):
+            self.prepared_key = key
+            return self
+
+        if types.is_dict(key):
+            self._process_jwk(key)
+            return self
+
+        if types.is_string(key):
+            key = codecs.encode(key, encoding='utf-8')
+
+        if types.is_bytelike(key):
+            self.prepared_key = RSA.importKey(key)
+            return self
+
+        return Error("JWKError: Unable to parse an RSA_JWK from key: %s" % key).unwrap()
+    self = __init__(key, algorithm)
+
     def _process_cert(key):
         pemLines = key.replace(b' ', b'').split()
         certDer = base64url_decode(b''.join(pemLines[1:-1]))
@@ -377,6 +377,7 @@ def AESKey(key, algorithm):
         ALGORITHMS.A256GCM: AES.MODE_GCM,
     }
 
+    self.IV_BYTE_LENGTH_MODE_MAP = {AES.MODE_CBC: AES.block_size // 8, AES.MODE_GCM: 96 // 8}
 
     def __init__(key, algorithm):
         if not operator.contains(ALGORITHMS.AES, algorithm):
@@ -416,7 +417,8 @@ def AESKey(key, algorithm):
     def encrypt(plain_text, aad=None):
         plain_text = six.ensure_binary(plain_text)
         def _try_encrypt(self, plain_text, aad):
-            iv = get_random_bytes(AES.block_size)
+            iv_byte_length = self.IV_BYTE_LENGTH_MODE_MAP.get(self._mode, AES.block_size)
+            iv = get_random_bytes(iv_byte_length)
             cipher = AES.new(self._key, self._mode, iv)
             if self._mode == AES.MODE_CBC:
                 padded_plain_text = self._pad(AES.block_size, plain_text)

diff --git a/larky/src/main/resources/vendor/jose/utils.star b/larky/src/main/resources/vendor/jose/utils.star
@@ -1,8 +1,10 @@
 load("@stdlib//base64", base64="base64")
+load("@stdlib//binascii", hexlify="hexlify")
 load("@stdlib//builtins", builtins="builtins")
 load("@stdlib//codecs", codecs="codecs")
 load("@stdlib//struct", struct="struct")
 load("@stdlib//types", types="types")
+load("@vendor//Crypto/Util/py3compat", tostr="tostr")
 
 # Piggyback of the backends implementation of the function that converts a long
 # to a bytes stream. Some plumbing is necessary to have the signatures match.
@@ -25,7 +27,7 @@ def long_to_base64(data, size=0):
 
 
 def int_arr_to_long(arr):
-    return int(''.join(["%02x" % byte for byte in arr]), 16)
+    return int(tostr(hexlify(arr)), 16)
 
 
 def base64_to_long(data):
@@ -34,7 +36,7 @@ def base64_to_long(data):
 
     # urlsafe_b64decode will happily convert b64encoded data
     _d = base64.urlsafe_b64decode(bytes(data) + b"==")
-    return int_arr_to_long(struct.unpack("%sB" % len(_d), _d))
+    return int_arr_to_long(_d)
 
 
 def calculate_at_hash(access_token, hash_alg):

diff --git a/larky/src/test/resources/vendor_tests/jose/test_jose.star b/larky/src/test/resources/vendor_tests/jose/test_jose.star
@@ -23,6 +23,45 @@ def test_encrypt_and_decrypt_jwe_with_defaults():
     asserts.eq(decrypted_data, b'533')
 
 
+def test_decrypt_jwk():
+    jwk_data = json.loads("""
+    {
+      "kty": "RSA",
+      "kid": "qfNlZYnIBmF_1JyVr14wlu34pDNQ9ggEjC6HiDQkDks",
+      "n": "uo80-yn_zR3vrRWFunyLTm7NFKO0aWFH1LsdkRPm62w13AA1btO89IxncW56L8qIAX7xBlCt92TB4mVnGEzWUbPLyn-OPyg5sgltsoAjFcYeFc0uvp0-iRMse4udjwk0YqrCgAOwlcL5IQu9KmqH3KdH8o1Um9QtoBxHe4ACWMdu8I2Xktm2vi076opkT2AquqLXRzD4gujAALgTMnGmBmNSqhKa2GK4iSjqUSTkYL6wF_OtPfdmlrlri_qaChRk-kK1VzIDk8jG_-I8d47_Pdln7FmADhO1O4_-4IoOiM1ESarNkO46aFV4I4xcFWvh0JnJQIoSwgPoqCPH3v7NiQ",
+      "e": "AQAB",
+      "alg": "RSA-OAEP-256",
+      "use": "enc"
+    }
+    """)
+
+    jwe_data = """
+    {
+      "user": {
+        "name": {
+          "first_name": "Test",
+          "last_name": "Test"
+        },
+        "address": {
+          "street": "Street1",
+          "street2": "Apt.#1",
+          "city": "NewYork",
+          "region": "NY",
+          "postal_code": "12345",
+          "country": "US"
+        },
+        "phone_number": "+15174242424"
+      },
+      "card": {
+        "number": "4242424242424242",
+        "expiration": "07/26",
+        "cvv": "012"
+      }
+    }
+    """
+
+    jwe.encrypt(jwe_data, jwk_data, 'A256GCM', jwk_data.get("alg"), None, None, jwk_data.get("kid"))
+
 def test_decrypt_GCM256_AES_wrapped_key_jwe():
     jweString = b"eyJlbmMiOiJBMjU2R0NNIiwidGFnIjoiazhaNnpTNjRJclllaUNpNV9JaWY5QSIsImFsZyI6IkEyNTZHQ01LVyIsIml2IjoieW9sTk8xLVFXSVg3R1poSCJ9.knPF3qV22v0pE-N6oUzlSIoBUEjr_k3sfFyYX-XuSH8.XkADjU0P2phiynPA.cvCd.Wp5oFTRAzEIFN8pImDnhmw"
     encryptionKey = b'96a18c1acc0b48beb9b24479355b70b5'
@@ -308,6 +347,7 @@ def test_encrypt_with_extra_headers():
 def _testsuite():
     _suite = unittest.TestSuite()
 
+    _suite.addTest(unittest.FunctionTestCase(test_decrypt_jwk))
     _suite.addTest(unittest.FunctionTestCase(test_encrypt_and_decrypt_jwe_with_defaults))
     _suite.addTest(unittest.FunctionTestCase(test_decrypt_GCM256_AES_wrapped_key_jwe))
     larky.parametrize(

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -13,8 +13,9 @@ readme = "pylarky/README.md"
 [tool.poetry.dependencies]
 python = "^3.10"
 
-[tool.poetry.dev-dependencies]
+[tool.poetry.group.dev.dependencies]
 pytest = "^7.4.4"
+setuptools = "^74.1.2"
 
 [build-system]
 requires = ["setuptools", "poetry-core>=1.0.0"]