fix!: move JWSAlgModel, JWEEncModel, JWEZipModel to jwa module

Author: lepture

src/joserfc/jwa.py

@@ -1,3 +1,15 @@
+from ._rfc7515.registry import JWSRegistry
+from ._rfc7515.model import JWSAlgModel
+from ._rfc7516.registry import JWERegistry
+from ._rfc7516.models import (
+    JWEDirectEncryption,
+    JWEKeyEncryption,
+    JWEKeyWrapping,
+    JWEKeyAgreement,
+    JWEAlgModel,
+    JWEEncModel,
+    JWEZipModel,
+)
 from ._rfc7518.jws_algs import (
     NoneAlgorithm,
     HMACAlgorithm,
@@ -25,10 +37,12 @@
 )
 from ._rfc8037.jws_eddsa import EdDSA, EdDSAAlgorithm
 from ._rfc8812 import ES256K
+from ._keys import KeySet
 
 __all__ = [
     # JWS algorithms
     "JWS_ALGORITHMS",
+    "JWSAlgModel",
     "NoneAlgorithm",
     "HMACAlgorithm",
     "RSAAlgorithm",
@@ -39,18 +53,46 @@
     "JWE_ALG_MODELS",
     "JWE_ENC_MODELS",
     "JWE_ZIP_MODELS",
+    "JWEAlgModel",
+    "JWEDirectEncryption",
+    "JWEKeyEncryption",
+    "JWEKeyWrapping",
+    "JWEKeyAgreement",
     "DirectAlgEncryption",
     "AESAlgKeyWrapping",
     "ECDHESAlgKeyAgreement",
     "AESGCMAlgKeyWrapping",
     "PBES2HSAlgKeyEncryption",
+    "JWEEncModel",
     "CBCHS2EncModel",
     "GCMEncModel",
+    "JWEZipModel",
     "DeflateZipModel",
+    # setup methods
+    "setup_jws_algorithms",
+    "setup_jwe_algorithms",
 ]
 
 JWS_ALGORITHMS = [
     *_JWS_ALGORITHMS,
     EdDSA,
     ES256K,
 ]
+
+
+def setup_jws_algorithms() -> None:
+    for _alg in JWS_ALGORITHMS:
+        JWSRegistry.register(_alg)
+        KeySet.algorithm_keys[_alg.name] = [_alg.key_type]
+
+
+def setup_jwe_algorithms() -> None:
+    for _alg in JWE_ALG_MODELS:
+        KeySet.algorithm_keys[_alg.name] = _alg.key_types
+        JWERegistry.register(_alg)
+
+    for _enc in JWE_ENC_MODELS:
+        JWERegistry.register(_enc)
+
+    for _zip in JWE_ZIP_MODELS:
+        JWERegistry.register(_zip)

src/joserfc/jwe.py

@@ -9,8 +9,6 @@
     CompactEncryption as CompactEncryption,
     GeneralJSONEncryption as GeneralJSONEncryption,
     FlattenedJSONEncryption as FlattenedJSONEncryption,
-    JWEEncModel,
-    JWEZipModel,
 )
 from ._rfc7516.registry import (
     JWERegistry as JWERegistry,
@@ -24,7 +22,7 @@
     extract_general_json,
     extract_flattened_json,
 )
-from .jwa import JWE_ALG_MODELS, JWE_ENC_MODELS, JWE_ZIP_MODELS
+from .jwa import setup_jwe_algorithms
 from .jwk import Key, KeySet, ECKey, OKPKey, KeyFlexible, guess_key
 from .util import to_bytes
 from .registry import Header
@@ -35,8 +33,6 @@
     "FlattenedJSONSerialization",
     # modules
     "JWERegistry",
-    "JWEEncModel",
-    "JWEZipModel",
     "Recipient",
     "CompactEncryption",
     "GeneralJSONEncryption",
@@ -48,25 +44,8 @@
     "decrypt_json",
     # consts
     "default_registry",
-    "JWE_ALG_MODELS",
-    "JWE_ENC_MODELS",
-    "JWE_ZIP_MODELS",
 ]
-
-
-def __setup_jwe() -> None:
-    for _alg in JWE_ALG_MODELS:
-        KeySet.algorithm_keys[_alg.name] = _alg.key_types
-        JWERegistry.register(_alg)
-
-    for _enc in JWE_ENC_MODELS:
-        JWERegistry.register(_enc)
-
-    for _zip in JWE_ZIP_MODELS:
-        JWERegistry.register(_zip)
-
-
-__setup_jwe()
+setup_jwe_algorithms()
 
 
 def encrypt_compact(

src/joserfc/jws.py

@@ -10,6 +10,7 @@
 from ._rfc7515.registry import (
     JWSRegistry as JWSRegistry,
     construct_registry,
+    default_registry,
 )
 from ._rfc7515.compact import (
     sign_compact,
@@ -39,8 +40,8 @@
     extract_rfc7797_json as extract_flattened_json,
 )
 from .errors import BadSignatureError, MissingKeyError
-from .jwk import Key, KeyFlexible, KeySet, guess_key
-from .jwa import JWS_ALGORITHMS
+from .jwk import Key, KeyFlexible, guess_key
+from .jwa import setup_jws_algorithms
 from .util import to_bytes
 from .registry import Header
 
@@ -50,7 +51,6 @@
     "GeneralJSONSerialization",
     "FlattenedJSONSerialization",
     # modules
-    "JWSAlgModel",
     "JWSRegistry",
     "HeaderMember",
     "CompactSignature",
@@ -64,16 +64,11 @@
     "serialize_json",
     "deserialize_json",
     "detach_content",
+    # consts
+    "default_registry",
 ]
 
-
-def __setup_jws() -> None:
-    for _alg in JWS_ALGORITHMS:
-        JWSRegistry.register(_alg)
-        KeySet.algorithm_keys[_alg.name] = [_alg.key_type]
-
-
-__setup_jws()
+setup_jws_algorithms()
 
 
 def serialize_compact(

tests/jwe/test_compact.py

@@ -4,8 +4,8 @@
     encrypt_compact,
     decrypt_compact,
     CompactEncryption,
-    JWE_ENC_MODELS,
 )
+from joserfc.jwa import JWE_ENC_MODELS
 from joserfc.jwk import RSAKey, ECKey, OctKey, OKPKey, KeySet
 from joserfc.errors import (
     InvalidKeyLengthError,

tests/jwe/test_ecdh_1pu.py

@@ -6,8 +6,8 @@
     decrypt_compact,
     encrypt_json,
     decrypt_json,
-    JWE_ENC_MODELS,
 )
+from joserfc.jwa import JWE_ENC_MODELS
 from joserfc.jwk import KeySet
 from joserfc.errors import InvalidEncryptionAlgorithmError
 from joserfc.drafts.jwe_ecdh_1pu import JWE_ALG_MODELS, register_ecdh_1pu