Merge pull request #441 from kentakayama/add-aes-ctr-and-cbc-support

dajiaji · web-flow · commit 98d1ef9e52b0 · 2023-10-11T04:57:52.000+09:00
Add AES-CTR and AES-CBC support
diff --git a/README.md b/README.md
@@ -1744,6 +1744,7 @@ Python CWT is (partially) compliant with following specifications:
 - [RFC8747: Proof-of-Possession Key Semantics for CBOR Web Tokens (CWTs)](https://tools.ietf.org/html/rfc8747)
 - [RFC8392: CWT (CBOR Web Token)](https://tools.ietf.org/html/rfc8392)
 - [RFC8230: Using RSA Algorithms with COSE Messages](https://tools.ietf.org/html/rfc8230)
+- [RFC9459: CBOR Object Signing and Encryption (COSE): AES-CTR and AES-CBC](https://www.rfc-editor.org/rfc/rfc9459.html) - experimental
 - [RFC8152: CBOR Object Signing and Encryption (COSE)](https://tools.ietf.org/html/rfc8152)
 - [draft-05: Use of HPKE with COSE](https://www.ietf.org/archive/id/draft-ietf-cose-hpke-06.html) - experimental
 - [draft-06: CWT Claims in COSE Headers](https://www.ietf.org/archive/id/draft-ietf-cose-cwt-claims-in-headers-06.html) - experimental
diff --git a/cwt/algs/non_aead.py b/cwt/algs/non_aead.py
@@ -0,0 +1,91 @@
+import os
+
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+
+
+class AESCTR:
+    _MAX_SIZE = 2**31 - 1
+
+    def __init__(self, key: bytes):
+        if len(key) not in (16, 24, 32):
+            raise ValueError("AESCTR key must be 128, 192, or 256 bits.")
+
+        self._key = key
+
+    @classmethod
+    def generate_key(cls, bit_length: int) -> bytes:
+        if not isinstance(bit_length, int):
+            raise TypeError("bit_length must be an integer")
+
+        if bit_length not in (128, 192, 256):
+            raise ValueError("bit_length must be 128, 192, or 256")
+
+        return os.urandom(bit_length // 8)
+
+    def encrypt(
+        self,
+        nonce: bytes,
+        data: bytes,
+    ) -> bytes:
+        encryptor = Cipher(
+            algorithms.AES(self._key),
+            modes.CTR(nonce),
+        ).encryptor()
+
+        return encryptor.update(data) + encryptor.finalize()
+
+    def decrypt(
+        self,
+        nonce: bytes,
+        data: bytes,
+    ) -> bytes:
+        decryptor = Cipher(
+            algorithms.AES(self._key),
+            modes.CTR(nonce),
+        ).decryptor()
+
+        return decryptor.update(data) + decryptor.finalize()
+
+
+class AESCBC:
+    _MAX_SIZE = 2**31 - 1
+
+    def __init__(self, key: bytes):
+        if len(key) not in (16, 24, 32):
+            raise ValueError("AESCBC key must be 128, 192, or 256 bits.")
+
+        self._key = key
+
+    @classmethod
+    def generate_key(cls, bit_length: int) -> bytes:
+        if not isinstance(bit_length, int):
+            raise TypeError("bit_length must be an integer")
+
+        if bit_length not in (128, 192, 256):
+            raise ValueError("bit_length must be 128, 192, or 256")
+
+        return os.urandom(bit_length // 8)
+
+    def encrypt(
+        self,
+        nonce: bytes,
+        data: bytes,
+    ) -> bytes:
+        encryptor = Cipher(
+            algorithms.AES(self._key),
+            modes.CBC(nonce),
+        ).encryptor()
+
+        return encryptor.update(data) + encryptor.finalize()
+
+    def decrypt(
+        self,
+        nonce: bytes,
+        data: bytes,
+    ) -> bytes:
+        decryptor = Cipher(
+            algorithms.AES(self._key),
+            modes.CBC(nonce),
+        ).decryptor()
+
+        return decryptor.update(data) + decryptor.finalize()
diff --git a/cwt/algs/symmetric.py b/cwt/algs/symmetric.py
@@ -9,12 +9,14 @@
 from ..const import COSE_KEY_OPERATION_VALUES
 from ..cose_key_interface import COSEKeyInterface
 from ..exceptions import DecodeError, EncodeError, VerifyError
+from .non_aead import AESCBC, AESCTR
 
 _CWT_DEFAULT_KEY_SIZE_HMAC256 = 32  # bytes
 _CWT_DEFAULT_KEY_SIZE_HMAC384 = 48
 _CWT_DEFAULT_KEY_SIZE_HMAC512 = 64
 _CWT_NONCE_SIZE_AESGCM = 12
 _CWT_NONCE_SIZE_CHACHA20_POLY1305 = 12
+_CWT_NONCE_SIZE_AES = 16
 
 
 class SymmetricKey(COSEKeyInterface):
@@ -353,3 +355,107 @@ def unwrap_key(self, wrapped_key: bytes) -> bytes:
             return aes_key_unwrap(self._key, wrapped_key)
         except Exception as err:
             raise DecodeError("Failed to unwrap key.") from err
+
+
+class AESCTRKey(ContentEncryptionKey):
+    """ """
+
+    def __init__(self, params: Dict[int, Any]):
+        """ """
+        super().__init__(params)
+
+        self._cipher: AESCTR
+
+        # Validate alg.
+        if self._alg == -65534:  # A128CTR
+            if not self._key:
+                self._key = AESCTR.generate_key(bit_length=128)
+            if len(self._key) != 16:
+                raise ValueError("The length of A128CTR key should be 16 bytes.")
+        elif self._alg == -65533:  # A192CTR
+            if not self._key:
+                self._key = AESCTR.generate_key(bit_length=192)
+            if len(self._key) != 24:
+                raise ValueError("The length of A192CTR key should be 24 bytes.")
+        elif self._alg == -65532:  # A256CTR
+            if not self._key:
+                self._key = AESCTR.generate_key(bit_length=256)
+            if len(self._key) != 32:
+                raise ValueError("The length of A256CTR key should be 32 bytes.")
+        else:
+            raise ValueError(f"Unsupported or unknown alg(3) for AES CTR: {self._alg}.")
+
+        self._cipher = AESCTR(self._key)
+        return
+
+    def generate_nonce(self):
+        return token_bytes(_CWT_NONCE_SIZE_AES)
+
+    def encrypt(self, msg: bytes, nonce: bytes, aad: Optional[bytes] = None) -> bytes:
+        """ """
+        try:
+            return self._cipher.encrypt(nonce, msg)
+        except Exception as err:
+            raise EncodeError("Failed to encrypt.") from err
+
+    def decrypt(self, msg: bytes, nonce: bytes, aad: Optional[bytes] = None) -> bytes:
+        """ """
+        try:
+            return self._cipher.decrypt(nonce, msg)
+        except Exception as err:
+            raise DecodeError("Failed to decrypt.") from err
+
+
+class AESCBCKey(ContentEncryptionKey):
+    """ """
+
+    def __init__(self, params: Dict[int, Any]):
+        """ """
+        super().__init__(params)
+
+        self._cipher: AESCBC
+
+        # Validate alg.
+        if self._alg == -65531:  # A128CBC
+            if not self._key:
+                self._key = AESCBC.generate_key(bit_length=128)
+            if len(self._key) != 16:
+                raise ValueError("The length of A128CBC key should be 16 bytes.")
+        elif self._alg == -65530:  # A192CBC
+            if not self._key:
+                self._key = AESCBC.generate_key(bit_length=192)
+            if len(self._key) != 24:
+                raise ValueError("The length of A192CBC key should be 24 bytes.")
+        elif self._alg == -65529:  # A256CBC
+            if not self._key:
+                self._key = AESCBC.generate_key(bit_length=256)
+            if len(self._key) != 32:
+                raise ValueError("The length of A256CBC key should be 32 bytes.")
+        else:
+            raise ValueError(f"Unsupported or unknown alg(3) for AES CBC: {self._alg}.")
+
+        self._cipher = AESCBC(self._key)
+        return
+
+    def generate_nonce(self):
+        return token_bytes(_CWT_NONCE_SIZE_AES)
+
+    def encrypt(self, msg: bytes, nonce: bytes, aad: Optional[bytes] = None) -> bytes:
+        """ """
+        try:
+            # Add padding (see RFC 9459 and 5652)
+            padding_value = 16 - len(msg) % 16
+            padding_length = 16 if padding_value == 0 else padding_value
+            padding = (padding_value).to_bytes(1, "big") * padding_length
+            return self._cipher.encrypt(nonce, msg + padding)
+        except Exception as err:
+            raise EncodeError("Failed to encrypt.") from err
+
+    def decrypt(self, msg: bytes, nonce: bytes, aad: Optional[bytes] = None) -> bytes:
+        """ """
+        try:
+            decrypted = self._cipher.decrypt(nonce, msg)
+            # Remove padding (see RFC 9459 and 5652)
+            return decrypted[0 : -(decrypted[-1])]
+        except Exception as err:
+            raise DecodeError("Failed to decrypt.") from err
diff --git a/cwt/const.py b/cwt/const.py
@@ -120,6 +120,12 @@
 
 # COSE Algorithms for Content Encryption Key (CEK).
 COSE_ALGORITHMS_CEK = {
+    "A128CTR": -65534,  # AES-CTR mode w/ 128-bit key (Deprecated)
+    "A192CTR": -65533,  # AES-CTR mode w/ 192-bit key (Deprecated)
+    "A256CTR": -65532,  # AES-CTR mode w/ 256-bit key (Deprecated)
+    "A128CBC": -65531,  # AES-CBC mode w/ 128-bit key (Deprecated)
+    "A192CBC": -65530,  # AES-CBC mode w/ 192-bit key (Deprecated)
+    "A256CBC": -65529,  # AES-CBC mode w/ 256-bit key (Deprecated)
     "A128GCM": 1,  # AES-GCM mode w/ 128-bit key, 128-bit tag
     "A192GCM": 2,  # AES-GCM mode w/ 192-bit key, 128-bit tag
     "A256GCM": 3,  # AES-GCM mode w/ 256-bit key, 128-bit tag
@@ -136,6 +142,12 @@
 }
 
 COSE_KEY_LEN = {
+    -65534: 128,  # AES-CTR w/ 128-bit key (Deprecated)
+    -65533: 192,  # AES-CTR w/ 192-bit key (Deprecated)
+    -65532: 256,  # AES-CTR w/ 256-bit key (Deprecated)
+    -65531: 128,  # AES-CBC w/ 128-bit key (Deprecated)
+    -65530: 192,  # AES-CBC w/ 192-bit key (Deprecated)
+    -65529: 256,  # AES-CBC w/ 256-bit key (Deprecated)
     -5: 256,  # AES Key Wrap w/ 256-bit key
     -4: 192,  # AES Key Wrap w/ 192-bit key
     -3: 128,  # AES Key Wrap w/ 128-bit key
diff --git a/cwt/cose_key.py b/cwt/cose_key.py
@@ -16,7 +16,15 @@
 from .algs.okp import OKPKey
 from .algs.raw import RawKey
 from .algs.rsa import RSAKey
-from .algs.symmetric import AESCCMKey, AESGCMKey, AESKeyWrap, ChaCha20Key, HMACKey
+from .algs.symmetric import (
+    AESCBCKey,
+    AESCCMKey,
+    AESCTRKey,
+    AESGCMKey,
+    AESKeyWrap,
+    ChaCha20Key,
+    HMACKey,
+)
 from .const import (
     COSE_ALGORITHMS_CKDM_KEY_AGREEMENT,
     COSE_ALGORITHMS_RSA,
@@ -76,6 +84,10 @@ def new(params: Dict[int, Any]) -> COSEKeyInterface:
             return ChaCha20Key(params)
         if params[COSEKeyParams.ALG] in [-3, -4, -5]:
             return AESKeyWrap(params)
+        if params[COSEKeyParams.ALG] in [-65534, -65533, -65532]:
+            return AESCTRKey(params)
+        if params[COSEKeyParams.ALG] in [-65531, -65530, -65529]:
+            return AESCBCKey(params)
         raise ValueError(f"Unsupported or unknown alg(3): {params[3]}.")
 
     @classmethod
diff --git a/cwt/enums.py b/cwt/enums.py
@@ -60,6 +60,12 @@ class COSEKeyParams(enum.IntEnum):
 
 
 class COSEAlgs(enum.IntEnum):
+    A128CTR = -65534
+    A192CTR = -65533
+    A256CTR = -65532
+    A128CBC = -65531
+    A192CBC = -65530
+    A256CBC = -65529
     RS512 = -259
     RS384 = -258
     RS256 = -257
diff --git a/docs/algorithms.rst b/docs/algorithms.rst
@@ -30,11 +30,30 @@ COSE Key Types
 COSE Algorithms
 ---------------
 
+    -65534: 128,  # AES-CTR w/ 128-bit key (Deprecated)
+    -65533: 192,  # AES-CTR w/ 192-bit key (Deprecated)
+    -65532: 256,  # AES-CTR w/ 256-bit key (Deprecated)
+    -65531: 128,  # AES-CBC w/ 128-bit key (Deprecated)
+    -65530: 192,  # AES-CBC w/ 192-bit key (Deprecated)
+    -65529: 256,  # AES-CBC w/ 256-bit key (Deprecated)
+
 +------------------------+--------+-------+-----------------------------------------------------+
 | Name                   | Status | Value | Description                                         |
 +========================+========+=======+=====================================================+
 | RS1                    | ➖     | -65535| RSASSA-PKCS1-v1_5 using SHA-1                       |
 +------------------------+--------+-------+-----------------------------------------------------+
+| A128CTR                | ✅     | -65534| AES-CTR w/ 128-bit key                              |
++------------------------+--------+-------+-----------------------------------------------------+
+| A192CTR                | ✅     | -65533| AES-CTR w/ 192-bit key                              |
++------------------------+--------+-------+-----------------------------------------------------+
+| A256CTR                | ✅     | -65532| AES-CTR w/ 256-bit key                              |
++------------------------+--------+-------+-----------------------------------------------------+
+| A128CBC                | ✅     | -65531| AES-CBC w/ 128-bit key                              |
++------------------------+--------+-------+-----------------------------------------------------+
+| A192CBC                | ✅     | -65530| AES-CBC w/ 192-bit key                              |
++------------------------+--------+-------+-----------------------------------------------------+
+| A256CBC                | ✅     | -65529| AES-CBC w/ 256-bit key                              |
++------------------------+--------+-------+-----------------------------------------------------+
 | WalnutDSA              |        | -260  | WalnutDSA signature                                 |
 +------------------------+--------+-------+-----------------------------------------------------+
 | RS512                  | ✅     | -259  | RSASSA-PKCS1-v1_5 using SHA-512                     |
diff --git a/tests/test_algs_symmetric.py b/tests/test_algs_symmetric.py
diff --git a/tests/test_recipient.py b/tests/test_recipient.py
