From a26db01ac3329871e8484c8a85c9d6b1089db864 Mon Sep 17 00:00:00 2001 From: NexusLightvale Date: Sat, 24 Oct 2026 17:38:49 +0700 Subject: [PATCH] Ruff messed up my formatting, EdDSA simplified, ML-DSA initial support, rotor machine ciphers initial implementation (with Claude Sonnet 5.5 support for implementation and testing as a non-agent), and AI use disclaimer in README --- README.md | 16 +- __init__.py | 10 +- src/ciphers.py | 6 +- .../asymmetric_encryption/eddsa.py | 130 ++-- .../asymmetric_encryption/elliptic_curve.py | 1 + .../asymmetric_encryption/mldsa.py | 241 ++++++ .../asymmetric_encryption/rsa.py | 222 ++++-- .../asymmetric_encryption/serialization.py | 27 +- .../asymmetric_encryption/xdh.py | 1 + src/cryptography_primitives/auth_encrypt.py | 6 +- src/cryptography_primitives/const_time.py | 2 + src/cryptography_primitives/hashing.py | 2 + src/cryptography_primitives/kdf.py | 46 +- src/cryptography_primitives/key_wrapper.py | 3 + src/cryptography_primitives/symm_padding.py | 3 + .../symmetrical_encrypt.py | 6 +- src/debugging_nodes.py | 6 +- src/error_correction.py | 3 +- src/fernet.py | 3 +- src/morse_code.py | 2 +- src/rotor_machines.py | 719 ++++++++++++++++++ src/spectrogram.py | 77 +- src/steganography.py | 7 +- src/utils.py | 38 +- tests/conftest.py | 2 +- tests/nodes/test_asymmetric_encryption.py | 52 +- tests/nodes/test_auth_encrypt.py | 1 + tests/nodes/test_ciphers.py | 28 +- tests/nodes/test_error_correction.py | 2 +- tests/nodes/test_hashing.py | 1 + tests/nodes/test_morse_code.py | 1 + tests/nodes/test_rotor_machines.py | 302 ++++++++ tests/nodes/test_spectrogram.py | 408 ++-------- tests/nodes/test_utils.py | 1 + 34 files changed, 1720 insertions(+), 655 deletions(-) create mode 100644 src/cryptography_primitives/asymmetric_encryption/mldsa.py create mode 100644 src/rotor_machines.py create mode 100644 tests/nodes/test_rotor_machines.py diff --git a/README.md b/README.md index 59239bf..fa9c188 100644 --- a/README.md +++ b/README.md @@ -30,7 +30,11 @@ Might also double as a poor man's cryptanalysis and steganography lab if absolut ## Testing methodology -Since this project started as an amateur project, the test files only cover the functions and most common use cases, not the edge cases or any automated testing yet. Many of the tests done are manual functionality tests done on ComfyUI and not unit tests. This can and will change on future releases. +Since this project started as an amateur project, the test files only cover the functions and most common use cases, not the edge cases or any automated testing yet. Many of the tests done are manual functionality tests done on ComfyUI and not unit tests. This can and will change on future releases to add more tests, including workflow tests and examples. + +## AI Use Disclosure and Disclaimer + +Many of the nodes in this project were made with help of generative AI tools like LLM agents (through Cline) and direct questioning and prompting without agentic access manually. This includes unit tests, documentations, and certain nodes and features. I am not inherently against AI use in coding, and every major changes and code pieces in this project all have human reviews on what they are doing. This project is not something that I can do by myself without help, and especially in any reasonable amount of time without AI assistance cutting down the major code works. If you want to take on similar projects by yourself without AI use, feel free to do so. ## Other information @@ -47,12 +51,12 @@ If you like the idea behind this node pack, want to explore cryptography, stegan ### Companion node packs -- [ComfyUI-Custom-Scripts](https://github.com/pythongosssss/ComfyUI-Custom-Scripts) - Provides a lot of useful features and nodes in general for ComfyUI. -- [ComfyUI-Crystools](https://github.com/crystian/ComfyUI-Crystools) - System monitor extension for ComfyUI. -- [ComfyUI_Invisible_Watermark](https://github.com/web3nomad/ComfyUI_Invisible_Watermark) - The main reason I even started this project in the first place, for making an "invisible watermark" node that left such an impression on me I decided to make my own. -- [ComfyLiterals](https://github.com/M1kep/ComfyLiterals) - Optional custom nodes for use with these nodes to provide or bypass certain limitations of ComfyUI. +- [ComfyUI-Custom-Scripts](https://github.com/pythongosssss/ComfyUI-Custom-Scripts) - Provides a lot of useful features and nodes in general for ComfyUI, though the main standout feature (Show Text) is now less useful thanks to newer releases of ComfyUI utilizing Preview as String. +- [ComfyUI-Crystools](https://github.com/crystian/ComfyUI-Crystools) - System monitor extension for ComfyUI. Useful to check system status when running certain nodes that has the capability to freeze the workflow for long periods of time. +- [ComfyUI_Invisible_Watermark](https://github.com/web3nomad/ComfyUI_Invisible_Watermark) - The main reason I even started this project in the first place, for making an "invisible watermark" node that left such a frankly terrible impression on me that I decided to make my own. +- [ComfyLiterals](https://github.com/M1kep/ComfyLiterals) - Optional custom nodes for use with these nodes to provide or bypass certain limitations of ComfyUI widgets. ## License -This project is licensed under the GNU General Public License 3.0 (GPLv3). +This project is licensed under the GNU General Public License 3.0 (GPLv3). Additional libraries upstream of this project follow their own open-source licenses. See the LICENSE file for full license details. diff --git a/__init__.py b/__init__.py index 4cc7f34..64ae3b8 100644 --- a/__init__.py +++ b/__init__.py @@ -28,20 +28,14 @@ for root, _, files in os.walk(SRC_FOLDER): if getattr(module, "NODE_CLASS_MAPPINGS", None): for key, value in module.NODE_CLASS_MAPPINGS.items(): if key in NODE_CLASS_MAPPINGS: - print( - f"[ComfyUI-ARG-Toolkit] [Duplicate NODE_CLASS_MAPPINGS] '{key}' " - f"overwritten by {module_name} (was {NODE_CLASS_SOURCES[key]})" - ) + print(f"[ComfyUI-ARG-Toolkit] [Duplicate NODE_CLASS_MAPPINGS] '{key}' overwritten by {module_name} (was {NODE_CLASS_SOURCES[key]})") NODE_CLASS_MAPPINGS[key] = value NODE_CLASS_SOURCES[key] = module_name if getattr(module, "NODE_DISPLAY_NAME_MAPPINGS", None): for key, value in module.NODE_DISPLAY_NAME_MAPPINGS.items(): if key in NODE_DISPLAY_NAME_MAPPINGS: - print( - f"[ComfyUI-ARG-Toolkit] [Duplicate NODE_DISPLAY_NAME_MAPPINGS] '{key}' " - f"overwritten by {module_name} (was {NODE_DISPLAY_NAME_SOURCES[key]})" - ) + print(f"[ComfyUI-ARG-Toolkit] [Duplicate NODE_DISPLAY_NAME_MAPPINGS] '{key}' overwritten by {module_name} (was {NODE_DISPLAY_NAME_SOURCES[key]})") NODE_DISPLAY_NAME_MAPPINGS[key] = value NODE_DISPLAY_NAME_SOURCES[key] = module_name diff --git a/src/ciphers.py b/src/ciphers.py index 11f8a70..bbda968 100644 --- a/src/ciphers.py +++ b/src/ciphers.py @@ -1,3 +1,5 @@ +# Classical ciphers (pen and paper, ancient ciphers,...) + import secretpy from secretpy import alphabets as al import math @@ -507,9 +509,7 @@ class Polybius(BaseCipherNode): def execute(self, text, alphabet, key, mode, keep_formatting): if key != math.sqrt(len(alphabet)): - raise ValueError( - f"Your given key length ({len(key)} does not match the square root of your alphabet size ({math.sqrt(len(alphabet))})). Please check your alphabet and key again" - ) + raise ValueError(f"Your given key length ({len(key)} does not match the square root of your alphabet size ({math.sqrt(len(alphabet))})). Please check your alphabet and key again") processed_alphabet = self.alphabet_checker(alphabet, as_tuple=True) return self.execute_cipher(text, processed_alphabet, key, mode, keep_formatting) diff --git a/src/cryptography_primitives/asymmetric_encryption/eddsa.py b/src/cryptography_primitives/asymmetric_encryption/eddsa.py index d185dfd..34f68b5 100644 --- a/src/cryptography_primitives/asymmetric_encryption/eddsa.py +++ b/src/cryptography_primitives/asymmetric_encryption/eddsa.py @@ -1,3 +1,5 @@ +# EdDSA signature algorithms + from cryptography.hazmat.primitives.asymmetric import ed25519, ed448 from cryptography.hazmat.primitives import serialization @@ -5,27 +7,24 @@ from cryptography.exceptions import InvalidSignature class InitNode: - def __init__(self): - pass - CATEGORY = "ARG Toolkit/Cryptography/Modern/Asymmetric" FUNCTION = "execute" @staticmethod - def get_private_key(key_source, key_type, private_key_bytes=None): - if key_type == "Ed25519": - pk_cls = ed25519.Ed25519PrivateKey - elif key_type == "Ed448": - pk_cls = ed448.Ed448PrivateKey - else: - raise ValueError(f"Unsupported key type: {key_type}") + def get_private_format(sig_alg): + algorithm = { + "Ed25519": ed25519.Ed25519PrivateKey, + "Ed448": ed448.Ed448PrivateKey, + } + return algorithm[sig_alg] - if key_source == "Fresh Key": - return pk_cls.generate() - elif key_source == "From Private Bytes": - if private_key_bytes is None: - raise ValueError("Private key bytes are required when key_source is 'From Private Bytes'") - return pk_cls.from_private_bytes(private_key_bytes) + @staticmethod + def get_public_format(sig_alg): + algorithm = { + "Ed25519": ed25519.Ed25519PublicKey, + "Ed448": ed448.Ed448PublicKey, + } + return algorithm[sig_alg] class EdDSAPrivateKeyFormat(InitNode): @@ -33,9 +32,9 @@ class EdDSAPrivateKeyFormat(InitNode): def INPUT_TYPES(cls): return { "required": { - "key_type": ( + "sig_alg": ( ["Ed25519", "Ed448"], - {"default": "Ed25519", "description": "The key type to use for EdDSA-based asymmetric signing algorithms."}, + {"default": "Ed25519", "description": "The signature algorithm for EdDSA-based asymmetric signing."}, ), "encoding": ( ["PEM", "DER", "Raw"], @@ -68,8 +67,8 @@ class EdDSAPrivateKeyFormat(InitNode): RETURN_TYPES = ("BYTESLIKE", "KEYOBJ") RETURN_NAMES = ("private_bytes", "private_key") - def execute(self, key_type, encoding, formatting, encryption: str, encryption_password=""): - private_key = self.get_private_key(key_source="Fresh Key", key_type=key_type) + def execute(self, sig_alg, encoding, formatting, encryption: str, encryption_password=""): + private_key = InitNode.get_private_format(sig_alg=sig_alg).generate() encode = getattr(serialization.Encoding, encoding) formatting = getattr(serialization.PrivateFormat, formatting) @@ -79,7 +78,7 @@ class EdDSAPrivateKeyFormat(InitNode): enc_alg = serialization.NoEncryption() output = private_key.private_bytes(encoding=encode, format=formatting, encryption_algorithm=enc_alg) - return (output, private_key,) + return (output, private_key) class EdDSASignature(InitNode): @@ -88,22 +87,23 @@ class EdDSASignature(InitNode): return { "required": { "key_source": ( - ["Fresh Key", "From Private Bytes", "From Loaded Key"], + "BOOLEAN", { - "tooltip": "The source of the private key to be used to sign the message", - "default": "Fresh Key", + "label_on": "From Key Object", + "label_off": "From Private Bytes", + "tooltip": "The source of the private key to be used to verify the signature of the message", + "default": True, }, ), - "key_type": ( + "sig_alg": ( ["Ed25519", "Ed448"], - {"default": "Ed25519", "description": "The key type to use for EdDSA-based asymmetric signing algorithms."}, + {"default": "Ed25519", "description": "The signature algorithm for EdDSA-based asymmetric signing."}, ), "message": ( - "STRING", + "BYTESLIKE", { - "default": "Hello World!", - "multiline": True, - "placeholder": "Type your message here...", + "forceInput": True, + "tooltip": "The messange to sign.", }, ), }, @@ -115,19 +115,16 @@ class EdDSASignature(InitNode): "forceInput": True, }, ), - "serialized_key": ("KEYOBJ", {"forceInput": True, "tooltip": "Only applicable if key_source is `From Key Object`"}), + "private_key": ("KEYOBJ", {"forceInput": True, "tooltip": "Only applicable if key_source is `From Key Object`"}), }, } RETURN_TYPES = ("BYTESLIKE",) RETURN_NAMES = ("signature",) - def execute(self, key_source, key_type, message, serialized_key, private_bytes=None): - if key_source == "From Loaded Key": - private_key_obj = serialized_key - else: - private_key_obj = self.get_private_key(key_source, key_type, private_bytes) - signature = private_key_obj.sign(message.encode("utf-8")) + def execute(self, key_source, sig_alg, message, private_key, private_bytes): + private_key_obj = InitNode.get_private_format(sig_alg).from_private_bytes(private_bytes) if not key_source else private_key + signature = private_key_obj.sign(message) return (signature,) @@ -137,15 +134,17 @@ class EdDSAPublicKeyFormat(InitNode): return { "required": { "key_source": ( - ["Fresh Key", "From Private Bytes", "From Loaded Key"], + "BOOLEAN", { - "tooltip": "The source of the private key to be used to sign the message", - "default": "Fresh Key", + "label_on": "From Key Object", + "label_off": "From Public Bytes", + "default": False, + "tooltip": "The source of the private key to be used to verify the signature of the message", }, ), - "key_type": ( + "sig_alg": ( ["Ed25519", "Ed448"], - {"default": "Ed25519", "description": "The key type to use for EdDSA-based asymmetric signing algorithms."}, + {"default": "Ed25519", "description": "The signature algorithm for EdDSA-based asymmetric signing."}, ), "encoding": ( [ @@ -168,31 +167,22 @@ class EdDSAPublicKeyFormat(InitNode): ), }, "optional": { - "private_key": ( + "public_bytes": ( "BYTESLIKE", { - "tooltip": "Only applicable if key_source is 'From Private Bytes'.", + "tooltip": "Only applicable if key_source is 'From Public Bytes'.", "forceInput": True, }, ), - "serialized_key": ("KEYOBJ", - { - "forceInput": True, - "tooltip": "Only applicable is key_source is `From Key Object`." - } - ), + "private_key": ("KEYOBJ", {"forceInput": True, "tooltip": "Only applicable is key_source is `From Key Object`."}), }, } RETURN_TYPES = ("BYTESLIKE", "KEYOBJ") RETURN_NAMES = ("public_bytes", "public_key") - def execute(self, key_source, key_type, encoding, formatting, serialized_key, private_key=None): - if key_source == "From Loaded Key": - private_key_obj = serialized_key - else: - private_key_obj = self.get_private_key(key_source, key_type, private_key) - public_key = private_key_obj.public_key() + def execute(self, key_source, sig_alg, encoding, formatting, private_key, public_bytes=None): + public_key = InitNode.get_public_format(sig_alg).from_public_bytes(public_bytes) if not key_source else private_key.public_key() encode = getattr(serialization.Encoding, encoding) formatting = getattr(serialization.PublicFormat, formatting) @@ -205,42 +195,40 @@ class EdDSAVerify(InitNode): def INPUT_TYPES(cls): return { "required": { - "key_type": ( + "sig_alg": ( ["Ed25519", "Ed448"], - {"default": "Ed25519", "description": "The key type to use for EdDSA-based asymmetric signing algorithms."}, + {"default": "Ed25519", "description": "The signature algorithm for EdDSA-based asymmetric signing."}, ), "key_source": ( "BOOLEAN", { "label_on": "From Private Bytes", "label_off": "From Loaded Key", - "tooltip": "The source of the private key to be used to verify the signature of the message", + "tooltip": "The source of the public key to be used to verify the signature of the message", }, ), "message": ( "BYTESLIKE", {"forceInput": True}, ), - "signature": ("BYTESLIKE", {"forceInput": True,}), + "signature": ( + "BYTESLIKE", + { + "forceInput": True, + }, + ), }, "optional": { "public_bytes": ("BYTESLIKE", {"forceInput": True}), - "public_key": ("KEYOBJ", {"forceInput": True, "tooltip": "Only applicable if key_source is 'From Loaded Key'."},)}, + "public_key": ("KEYOBJ", {"forceInput": True, "tooltip": "Only applicable if key_source is 'From Loaded Key'."}), + }, } RETURN_TYPES = ("BOOLEAN",) RETURN_NAMES = ("verify_status",) - def execute(self, key_type, public_bytes, public_key, message, signature, key_source): - if key_source: - if key_type == "Ed25519": - pk_loader = ed25519.Ed25519PublicKey.from_public_bytes - public_key_obj = pk_loader(public_bytes) - elif key_type == "Ed448": - pk_loader = ed448.Ed448PublicKey.from_public_bytes - public_key_obj = pk_loader(public_bytes) - elif not key_source: - public_key_obj = public_key + def execute(self, sig_alg, public_bytes, public_key, message, signature, key_source): + public_key_obj = InitNode.get_public_format(sig_alg).from_public_bytes(public_bytes) if not key_source else public_key try: public_key_obj.verify(signature, message) output = True diff --git a/src/cryptography_primitives/asymmetric_encryption/elliptic_curve.py b/src/cryptography_primitives/asymmetric_encryption/elliptic_curve.py index c3e4226..0b4cc4c 100644 --- a/src/cryptography_primitives/asymmetric_encryption/elliptic_curve.py +++ b/src/cryptography_primitives/asymmetric_encryption/elliptic_curve.py @@ -1,4 +1,5 @@ # Elliptic curve encryption + from cryptography.exceptions import InvalidSignature from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.asymmetric import ec diff --git a/src/cryptography_primitives/asymmetric_encryption/mldsa.py b/src/cryptography_primitives/asymmetric_encryption/mldsa.py new file mode 100644 index 0000000..fd81f77 --- /dev/null +++ b/src/cryptography_primitives/asymmetric_encryption/mldsa.py @@ -0,0 +1,241 @@ +# ML-DSA signature algorithms + +# As of `cryptography` v51.0.0, DSA has been deprecated and replaced with ML-DSA (Module-Lattice-Based Digital Signature Algorithm) on the implementation side. Their recommendation is to use EdDSA or ML-DSA instead. + +from cryptography.hazmat.primitives.asymmetric import mldsa +from cryptography.hazmat.primitives import serialization + +from cryptography.exceptions import InvalidSignature + + +class InitNode: + CATEGORY = "ARG Toolkit/Cryptography/Modern/Asymmetric" + FUNCTION = "execute" + + @staticmethod + def get_private_format(sig_alg): + algorithm = { + "ML-DSA-44": mldsa.MLDSA44PrivateKey, + "ML-DSA-65": mldsa.MLDSA65PrivateKey, + "ML-DSA-87": mldsa.MLDSA87PrivateKey, + } + return algorithm[sig_alg] + + @staticmethod + def get_public_format(sig_alg): + algorithm = { + "ML-DSA-44": mldsa.MLDSA44PublicKey, + "ML-DSA-65": mldsa.MLDSA65PublicKey, + "ML-DSA-87": mldsa.MLDSA87PublicKey, + } + return algorithm[sig_alg] + + +# ML-DSA Private and Public Key loaders +class MLDSAPrivateKeyFormat(InitNode): + RETURN_TYPES = ( + "BYTESLIKE", + "KEYOBJ", + ) + RETURN_NAMES = ("private_bytes", "private_key") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "sig_alg": (["ML-DSA-44", "ML-DSA-65", "ML-DSA-87"], {"default": "ML-DSA-65", "tooltip": "The algorithm to use for the private key."}), + "key_source": ("BOOLEAN", {"default": True, "label_on": "Fresh Key", "label_off": "Seed Bytes", "tooltip": "The source to generate the private key."}), + "encoding": ( + ["PEM", "DER", "Raw"], + { + "default": "PEM", + "description": "Encoding type for the private key.", + }, + ), + "formatting": ( + ["PKCS8", "Raw"], + { + "default": "PKCS8", + "description": "What format to serialize the key in.", + }, + ), + "encryption": (["Best Available", "None"], {"default": "Best Available"}), + }, + "optional": { + "encryption_password": ( + "STRING", + { + "default": "", + "multiline": False, + "tooltip": "The password to use to encrypt the private key. Required if encryption is used.", + }, + ), + "seed_bytes": ( + "BYTESLIKE", + {"forceInput": True, "tooltip": "The seed bytes to use for generating the private key only when the source is 'Seed Bytes'. Has to be 32 bytes."}, + ), + }, + } + + def execute(self, sig_alg, key_source, encoding, formatting, encryption, encryption_password, seed_bytes): + private_key = self.get_private_format(sig_alg=sig_alg).generate() if key_source else self.get_private_format(sig_alg=sig_alg).from_seed_bytes(seed_bytes) + encode = getattr(serialization.Encoding, encoding) + formatting = getattr(serialization.PrivateFormat, formatting) + + if encryption == "Best Available" and encryption_password: + enc_alg = serialization.BestAvailableEncryption(encryption_password.encode("utf-8")) + else: + enc_alg = serialization.NoEncryption() + + private_bytes = private_key.private_bytes(encoding=encode, format=formatting, encryption_algorithm=enc_alg) + return (private_bytes, private_key) + + +class MLDSAPublicKeyFormat(InitNode): + RETURN_TYPES = ( + "BYTESLIKE", + "KEYOBJ", + ) + RETURN_NAMES = ("public_bytes", "public_key") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "sig_alg": (["ML-DSA-44", "ML-DSA-65", "ML-DSA-87"], {"default": "ML-DSA-65", "tooltip": "The algorithm to use for the public key."}), + "key_source": ("BOOLEAN", {"default": True, "label_on": "From Key Object", "label_off": "From Public Bytes", "tooltip": "The source to load the public key."}), + "encoding": ( + ["PEM", "DER", "Raw"], + { + "default": "PEM", + "description": "Encoding type for the public key.", + }, + ), + "formatting": ( + ["SubjectPublicKeyInfo", "Raw"], + { + "default": "SubjectPublicKeyInfo", + "description": "What format to serialize the key in.", + }, + ), + }, + "optional": { + "public_bytes": ( + "BYTESLIKE", + { + "tooltip": "Only applicable if key_source is `From Public Bytes`. Refer to documentation for exact byte size for each algorithm.", + "forceInput": True, + } + ), + "private_key": ("KEYOBJ", {"forceInput": True, "tooltip": "Only applicable is key_source is `From Key Object`."}), + }, + } + + def execute(self, sig_alg, key_source, encoding, formatting, public_bytes, private_key): + public_key = InitNode.get_public_format(sig_alg).from_public_bytes(public_bytes) if not key_source else private_key.public_key() + encode = getattr(serialization.Encoding, encoding) + formatting = getattr(serialization.PublicFormat, formatting) + + output = public_key.public_bytes(encoding=encode, format=formatting) + return (output, public_key) + + +## ML-DSA Signing and Verification (with context support) +class MLDSASignature(InitNode): + RETURN_TYPES = "BYTESLIKE" + RETURN_NAMES = "signature" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "key_source": ( + "BOOLEAN", + { + "label_on": "From Key Object", + "label_off": "From Private Bytes", + "tooltip": "The source of the private key to be used to verify the signature of the message", + "default": True, + }, + ), + "sig_alg": (["ML-DSA-44", "ML-DSA-65", "ML-DSA-87"], {"default": "ML-DSA-65", "tooltip": "The algorithm to use for the private key."}), + "message": ( + "BYTESLIKE", + { + "forceInput": True, + "tooltip": "The messange to sign.", + }, + ), + }, + "optional": { + "private_bytes": ( + "BYTESLIKE", + { + "tooltip": "Only applicable if key_source is 'From Private Bytes'.", + "forceInput": True, + }, + ), + "private_key": ("KEYOBJ", {"forceInput": True, "tooltip": "Only applicable if key_source is `From Key Object`"}), + "context": ("BYTESLIKE", {"forceInput": True, "tooltip": "Optional context string, up to 255 bytes."}), + }, + } + + def execute(self, key_source, sig_alg, message, private_bytes, private_key, context: bytes | None): + private_key = InitNode.get_private_format(sig_alg).from_private_bytes(private_bytes) if not key_source else private_key + signature = private_key.sign(message, context) + return (signature,) + + +class MLDSAVerify(InitNode): + RETURN_TYPES = ("BOOLEAN",) + RETURN_NAMES = ("verify_status",) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "sig_alg": (["ML-DSA-44", "ML-DSA-65", "ML-DSA-87"], {"default": "ML-DSA-65", "tooltip": "The algorithm to use for the public key."}), + "key_source": ( + "BOOLEAN", + { + "label_on": "From Private Bytes", + "label_off": "From Loaded Key", + "tooltip": "The source of the public key to be used to verify the signature of the message", + }, + ), + "message": ( + "BYTESLIKE", + {"forceInput": True}, + ), + "signature": ( + "BYTESLIKE", + { + "forceInput": True, + }, + ), + }, + "optional": {"public_bytes": ("BYTESLIKE", {"forceInput": True}), "public_key": ("KEYOBJ", {"forceInput": True, "tooltip": "Only applicable if key_source is 'From Loaded Key'."}), "context": ("BYTESLIKE", {"forceInput": True, "tooltip": "Optional context string, up to 255 bytes."})}, + } + + def execute(self, sig_alg, public_bytes, public_key, message, signature, key_source, context: bytes | None): + public_key_obj = InitNode.get_public_format(sig_alg).from_public_bytes(public_bytes) if not key_source else public_key + try: + public_key_obj.verify(signature, message, context) + output = True + except (InvalidSignature, ValueError): + output = False + return (output,) + + +NODE_CLASS_MAPPINGS = { + "MLDSAPrivateKeyFormat": MLDSAPrivateKeyFormat, + "MLDSASignature": MLDSASignature, + "MLDSAPublicKeyFormat": MLDSAPublicKeyFormat, + "MLDSAVerify": MLDSAVerify, +} +NODE_DISPLAY_NAME_MAPPINGS = { + "MLDSAPrivateKeyFormat": "ML-DSA Private Key Bytes", + "MLDSASignature": "ML-DSA Signature Generator", + "MLDSAPublicKeyFormat": "ML-DSA Public Key Bytes", + "MLDSAVerify": "ML-DSA Message Verification", +} diff --git a/src/cryptography_primitives/asymmetric_encryption/rsa.py b/src/cryptography_primitives/asymmetric_encryption/rsa.py index f7fb3bd..6627a1e 100644 --- a/src/cryptography_primitives/asymmetric_encryption/rsa.py +++ b/src/cryptography_primitives/asymmetric_encryption/rsa.py @@ -1,12 +1,15 @@ -# RSA asymmetric encryption +# RSA asymmetric encryption and signing + from cryptography.hazmat.primitives.asymmetric import rsa, padding, utils from cryptography.hazmat.primitives import serialization, hashes + # Simple init node to deduplicate some bits class InitNode: CATEGORY = "ARG Toolkit/Cryptography/Modern/Asymmetric" FUNCTION = "execute" + # Actual RSA functions class RSAPrivateKeyFormat(InitNode): RETURN_TYPES = ("KEYOBJ", "BYTESLIKE", "NUMOBJ") @@ -16,13 +19,23 @@ class RSAPrivateKeyFormat(InitNode): def INPUT_TYPES(cls): return { "required": { - "public_exponent": ("INT", { - "min": 1, - "max": 4294967296, - "default": 65537, - } + "public_exponent": ( + "INT", + { + "min": 1, + "max": 4294967296, + "default": 65537, + }, + ), + "key_size": ( + "INT", + { + "min": 512, + "max": 8192, + "step": 128, + "default": 2048, + }, ), - "key_size": ("INT", {"min": 512, "max": 8192, "step": 128, "default": 2048,}), "encoding": ( ["PEM", "DER"], { @@ -31,7 +44,11 @@ class RSAPrivateKeyFormat(InitNode): }, ), "formatting": ( - ["PKCS8", "OpenSSH", "TraditionalOpenSSL",], + [ + "PKCS8", + "OpenSSH", + "TraditionalOpenSSL", + ], { "default": "PKCS8", "description": "What format to serialize the key in. OpenSSH requires PEM encoding.", @@ -66,6 +83,7 @@ class RSAPrivateKeyFormat(InitNode): private_bytes = private_key.private_bytes(encoding=encoding, format=formatting, encryption_algorithm=enc_alg) return (private_key, private_bytes, private_key.private_numbers()) + class RSAPublicKeyFormat(InitNode): RETURN_TYPES = ("KEYOBJ", "BYTESLIKE", "NUMOBJ") RETURN_NAMES = ("key_object", "public_key", "public_numbers") @@ -100,33 +118,13 @@ class RSAPublicKeyFormat(InitNode): ), }, "optional": { - "serialized_key": ("KEYOBJ", - { - "forceInput": True, - "tooltip": "Only applicable if key_source is `From Loaded Key`." - } - ), - "fresh_key_size": ("INT", - { - "min": 512, - "max": 8192, - "step": 128, - "default": 2048, - "tooltip": "Only applicable if key_source is `From Fresh Key`." - } - ), - "public_exponent": ("INT", - { - "min": 1, - "max": 4294967296, - "default": 65537, - "tooltip": "Only applicable if key_sourcce if `From Fresh Key`." - } - ), + "serialized_key": ("KEYOBJ", {"forceInput": True, "tooltip": "Only applicable if key_source is `From Loaded Key`."}), + "fresh_key_size": ("INT", {"min": 512, "max": 8192, "step": 128, "default": 2048, "tooltip": "Only applicable if key_source is `From Fresh Key`."}), + "public_exponent": ("INT", {"min": 1, "max": 4294967296, "default": 65537, "tooltip": "Only applicable if key_sourcce if `From Fresh Key`."}), }, } - def execute(self, key_source, encoding, formatting, serialized_key:rsa.RSAPrivateKey | None, fresh_key_size, public_exponent): + def execute(self, key_source, encoding, formatting, serialized_key: rsa.RSAPrivateKey | None, fresh_key_size, public_exponent): match key_source: case "Fresh Key": p_key = rsa.generate_private_key(public_exponent, fresh_key_size) @@ -143,6 +141,7 @@ class RSAPublicKeyFormat(InitNode): public_bytes = public_key.public_bytes(encoding=encoding, format=formatting) return (public_key, public_bytes, public_key.public_numbers()) + # RSA signing and verifying of signatures class RSASign(InitNode): RETURN_TYPES = ("SIGNATURE",) @@ -152,11 +151,14 @@ class RSASign(InitNode): def INPUT_TYPES(cls): return { "required": { - "private_key": ("KEYOBJ", {"forceInput": True,}), - "data": ("BYTESLIKE", {"forceInput": True}), - "salt_length": ( - ["Max Length", "Digest Length", "Manual"], {"default": "Max Length", "tooltip": "The length of the salt."} + "private_key": ( + "KEYOBJ", + { + "forceInput": True, + }, ), + "data": ("BYTESLIKE", {"forceInput": True}), + "salt_length": (["Max Length", "Digest Length", "Manual"], {"default": "Max Length", "tooltip": "The length of the salt."}), "signature_algorithm": ( [ "SHA224", @@ -177,12 +179,9 @@ class RSASign(InitNode): ], {"tooltip": "The hashing algorithm used for the signature."}, ), - "padding_alg": (["PSS", "PKCS1v15"], {"default": "PSS"}) - }, - "optional": { - "prehashed": ("BOOLEAN", {"tooltip": "Whether the data is prehashed or not."}), - "salt_length_manual": ("INT", {"min": 1, "max": 65536, "step": 1, "tooltip": "For only when salt_length is `Manual`."}) + "padding_alg": (["PSS", "PKCS1v15"], {"default": "PSS"}), }, + "optional": {"prehashed": ("BOOLEAN", {"tooltip": "Whether the data is prehashed or not."}), "salt_length_manual": ("INT", {"min": 1, "max": 65536, "step": 1, "tooltip": "For only when salt_length is `Manual`."})}, } def execute(self, private_key, data, salt_length, salt_length_manual, signature_algorithm, padding_alg, prehashed): @@ -195,7 +194,7 @@ class RSASign(InitNode): salt_length = padding.PSS.DIGEST_LENGTH case "Manual": salt_length = salt_length_manual - padder = padding.PSS(mgf=padding.MGF1(algorithm),salt_length=salt_length) + padder = padding.PSS(mgf=padding.MGF1(algorithm), salt_length=salt_length) if padding_alg == "PSS" and not prehashed: signature = private_key.sign(data, padder, algorithm) elif padding_alg == "PSS" and prehashed: @@ -206,6 +205,7 @@ class RSASign(InitNode): signature = private_key.sign(data, padding.PKCS1v15(), utils.Prehashed(algorithm)) return (signature,) + class RSAVerify(InitNode): RETURN_TYPES = ("BOOLEAN",) RETURN_NAMES = ("verify_result",) @@ -214,12 +214,25 @@ class RSAVerify(InitNode): def INPUT_TYPES(cls): return { "required": { - "public_key": ("KEYOBJ", {"forceInput": True,}), - "data": ("BYTESLIKE", {"forceInput": True,}), - "signature": ("SIGNATURE", {"forceInput": True,}), - "salt_length": ( - ["Max Length", "Digest Length", "Manual"], {"default": "Max Length", "tooltip": "The length of the salt."} + "public_key": ( + "KEYOBJ", + { + "forceInput": True, + }, ), + "data": ( + "BYTESLIKE", + { + "forceInput": True, + }, + ), + "signature": ( + "SIGNATURE", + { + "forceInput": True, + }, + ), + "salt_length": (["Max Length", "Digest Length", "Manual"], {"default": "Max Length", "tooltip": "The length of the salt."}), "signature_algorithm": ( [ "SHA224", @@ -240,17 +253,15 @@ class RSAVerify(InitNode): ], {"tooltip": "The hashing algorithm used for the signature."}, ), - "padding_alg": (["PSS", "PKCS1v15"], {"default": "PSS"}) - }, - "optional": { - "prehashed": ("BOOLEAN", {"tooltip": "Whether the data is prehashed or not."}), - "salt_length_manual": ("INT", {"min": 1, "max": 65536, "step": 1, "tooltip": "For only when salt_length is `Manual`."}) + "padding_alg": (["PSS", "PKCS1v15"], {"default": "PSS"}), }, + "optional": {"prehashed": ("BOOLEAN", {"tooltip": "Whether the data is prehashed or not."}), "salt_length_manual": ("INT", {"min": 1, "max": 65536, "step": 1, "tooltip": "For only when salt_length is `Manual`."})}, } def execute(self, public_key, data, signature, salt_length, salt_length_manual, signature_algorithm, padding_alg, prehashed): algorithm = getattr(hashes, signature_algorithm)() from cryptography.exceptions import InvalidSignature + if padding_alg == "PSS": match salt_length: case "Max Length": @@ -259,7 +270,7 @@ class RSAVerify(InitNode): salt_length = padding.PSS.DIGEST_LENGTH case "Manual": salt_length = salt_length_manual - padder = padding.PSS(mgf=padding.MGF1(algorithm),salt_length=salt_length) + padder = padding.PSS(mgf=padding.MGF1(algorithm), salt_length=salt_length) if padding_alg == "PSS" and not prehashed: try: public_key.verify(signature, data, padder, algorithm) @@ -286,6 +297,7 @@ class RSAVerify(InitNode): verify_result = False return (verify_result,) + # RSA encryption and decryption class RSAEncrypt(InitNode): RETURN_TYPES = ("BYTESLIKE",) @@ -295,8 +307,18 @@ class RSAEncrypt(InitNode): def INPUT_TYPES(cls): return { "required": { - "public_key": ("KEYOBJ", {"forceInput": True,}), - "data": ("BYTESLIKE", {"forceInput": True,}), + "public_key": ( + "KEYOBJ", + { + "forceInput": True, + }, + ), + "data": ( + "BYTESLIKE", + { + "forceInput": True, + }, + ), "signature_algorithm": ( [ "SHA224", @@ -317,14 +339,14 @@ class RSAEncrypt(InitNode): ], {"tooltip": "The hashing algorithm used for the signature."}, ), - "padding_alg": (["OAEP", "PKCS1v15"], {"default": "OAEP"}) + "padding_alg": (["OAEP", "PKCS1v15"], {"default": "OAEP"}), }, "optional": { "label": ("BYTESLIKE", {"forceInput": True}), }, } - def execute(self, public_key: rsa.RSAPublicKey, data, signature_algorithm, padding_alg, label = None): + def execute(self, public_key: rsa.RSAPublicKey, data, signature_algorithm, padding_alg, label=None): algorithm = getattr(hashes, signature_algorithm)() if padding_alg == "OAEP": encrypted_data = public_key.encrypt(data, padding.OAEP(mgf=padding.MGF1(algorithm), algorithm=algorithm, label=label)) @@ -332,6 +354,7 @@ class RSAEncrypt(InitNode): encrypted_data = public_key.encrypt(data, padding.PKCS1v15()) return (encrypted_data,) + class RSADecrypt(InitNode): RETURN_TYPES = ("BYTESLIKE",) RETURN_NAMES = ("decrypted_data",) @@ -340,8 +363,18 @@ class RSADecrypt(InitNode): def INPUT_TYPES(cls): return { "required": { - "private_key": ("KEYOBJ", {"forceInput": True,}), - "data": ("BYTESLIKE", {"forceInput": True,}), + "private_key": ( + "KEYOBJ", + { + "forceInput": True, + }, + ), + "data": ( + "BYTESLIKE", + { + "forceInput": True, + }, + ), "signature_algorithm": ( [ "SHA224", @@ -362,14 +395,14 @@ class RSADecrypt(InitNode): ], {"tooltip": "The hashing algorithm used for the signature."}, ), - "padding_alg": (["OAEP", "PKCS1v15"], {"default": "OAEP"}) + "padding_alg": (["OAEP", "PKCS1v15"], {"default": "OAEP"}), }, "optional": { "label": ("BYTESLIKE", {"forceInput": True}), }, } - def execute(self, private_key: rsa.RSAPrivateKey, data, signature_algorithm, padding_alg, label = None): + def execute(self, private_key: rsa.RSAPrivateKey, data, signature_algorithm, padding_alg, label=None): algorithm = getattr(hashes, signature_algorithm)() if padding_alg == "OAEP": decrypted_data = private_key.decrypt(data, padding.OAEP(mgf=padding.MGF1(algorithm), algorithm=algorithm, label=label)) @@ -377,9 +410,13 @@ class RSADecrypt(InitNode): decrypted_data = private_key.decrypt(data, padding.PKCS1v15()) return (decrypted_data,) + # RSA numbers, for when standard serialized key loading fails class RSAPublicNumbersFormat(InitNode): - RETURN_TYPES = ("KEYOBJ", "NUMOBJ",) + RETURN_TYPES = ( + "KEYOBJ", + "NUMOBJ", + ) RETURN_NAMES = ("public_key", "public_numbers") @classmethod @@ -395,6 +432,7 @@ class RSAPublicNumbersFormat(InitNode): pub_num = rsa.RSAPublicNumbers(n=public_modulus, e=public_exponent) return (pub_num.public_key(), pub_num) + class RSAPubNumValue(InitNode): RETURN_TYPES = ("INT", "INT") RETURN_NAMES = ("public_modulus", "public_exponent") @@ -416,11 +454,21 @@ class RSAPubNumValue(InitNode): pnum = public_numbers else: raise ValueError("Invalid configuration given. Input either the public numbers object or public key object only.") - return (pnum.n, pnum.e,) + return ( + pnum.n, + pnum.e, + ) + class RSAPrivateNumbersFormat(InitNode): - RETURN_TYPES = ("KEYOBJ", "NUMOBJ",) - RETURN_NAMES = ("public_key", "private_numbers",) + RETURN_TYPES = ( + "KEYOBJ", + "NUMOBJ", + ) + RETURN_NAMES = ( + "public_key", + "private_numbers", + ) @classmethod def INPUT_TYPES(cls): @@ -430,7 +478,7 @@ class RSAPrivateNumbersFormat(InitNode): "var_p": ("INT", {"min": 1, "max": 4294967296, "step": 1, "default": 1, "tooltip": "One of the two primes composing the private modulus `n`"}), "var_q": ("INT", {"min": 1, "max": 4294967296, "step": 1, "default": 1, "tooltip": "One of the two primes composing the private modulus `n`"}), "private_exponent": ("INT", {"min": 1, "max": 4294967296, "step": 1, "default": 65537, "tooltip": "The private exponent (d)."}), - "skip_rsa_key_validation": ("BOOLEAN", {"default": False, "tooltip": "Unless you know what you're doing, do not change this. Disabling RSA key validation speeds up key loading, but is **unsafe** and may cause OpenSSL to misbehave or crash, potentially dragging ComfyUI down with it."}) + "skip_rsa_key_validation": ("BOOLEAN", {"default": False, "tooltip": "Unless you know what you're doing, do not change this. Disabling RSA key validation speeds up key loading, but is **unsafe** and may cause OpenSSL to misbehave or crash, potentially dragging ComfyUI down with it."}), }, } @@ -440,11 +488,31 @@ class RSAPrivateNumbersFormat(InitNode): pvn_dmq1 = rsa.rsa_crt_dmq1(private_exponent, var_q) private_num = rsa.RSAPrivateNumbers(p=var_p, q=var_q, d=private_exponent, dmp1=pvn_dmp1, dmq1=pvn_dmq1, iqmp=pvn_iqmp, public_numbers=public_numbers) private_key = private_num.private_key(unsafe_skip_rsa_key_validation=skip_rsa_key_validation) - return (private_num, private_key,) + return ( + private_num, + private_key, + ) + class RSAPrivNumValue(InitNode): - RETURN_TYPES = ("NUMOBJ", "INT", "INT", "INT", "INT", "INT", "INT",) - RETURN_NAMES = ("public_numbers", "var_p", "var_q", "private_exponent", "dmp1", "dmq1", "iqmp",) + RETURN_TYPES = ( + "NUMOBJ", + "INT", + "INT", + "INT", + "INT", + "INT", + "INT", + ) + RETURN_NAMES = ( + "public_numbers", + "var_p", + "var_q", + "private_exponent", + "dmp1", + "dmq1", + "iqmp", + ) @classmethod def INPUT_TYPES(cls): @@ -453,7 +521,7 @@ class RSAPrivNumValue(InitNode): "optional": { "private_numbers": ("NUMOBJ", {"forceInput": True}), "private_key": ("KEYOBJ", {"forceInput": True}), - } + }, } def execute(self, private_numbers: rsa.RSAPrivateNumbers | None, private_key: rsa.RSAPrivateKey | None): @@ -463,10 +531,20 @@ class RSAPrivNumValue(InitNode): pnum = private_numbers else: raise ValueError("Invalid configuration given. Input either the private numbers object or private key object only.") - return (pnum.public_numbers, pnum.p, pnum.q, pnum.d, pnum.dmp1, pnum.dmq1, pnum.iqmp,) + return ( + pnum.public_numbers, + pnum.p, + pnum.q, + pnum.d, + pnum.dmp1, + pnum.dmq1, + pnum.iqmp, + ) + # Recovery nodes for partial numbers given + class RSARestoreD(InitNode): RETURN_TYPES = ("INT",) RETURN_NAMES = ("d",) @@ -484,6 +562,7 @@ class RSARestoreD(InitNode): def execute(self, public_exponent, var_p, var_q): return (rsa.rsa_recover_private_exponent(public_exponent, var_p, var_q),) + class RSARestorePrime(InitNode): RETURN_TYPES = ("INT", "INT") RETURN_NAMES = ("p", "q") @@ -501,6 +580,7 @@ class RSARestorePrime(InitNode): def execute(self, private_modulus, public_exponent, private_exponent): return rsa.rsa_recover_prime_factors(private_modulus, public_exponent, private_exponent) + NODE_CLASS_MAPPINGS = { "RSAPrivateKeyFormat": RSAPrivateKeyFormat, "RSAPublicKeyFormat": RSAPublicKeyFormat, diff --git a/src/cryptography_primitives/asymmetric_encryption/serialization.py b/src/cryptography_primitives/asymmetric_encryption/serialization.py index b80d1a4..ce5c0d1 100644 --- a/src/cryptography_primitives/asymmetric_encryption/serialization.py +++ b/src/cryptography_primitives/asymmetric_encryption/serialization.py @@ -1,3 +1,5 @@ +# Serialized key loading and SSH parsing + from cryptography.hazmat.primitives import serialization, hashes @@ -12,11 +14,15 @@ class KeyInitNode: def INPUT_TYPES(cls): return { "required": { - "keyfile": ("BYTESLIKE", {"forceInput": True},), + "keyfile": ( + "BYTESLIKE", + {"forceInput": True}, + ), }, "optional": {}, } + # PEM keys class PEMPrivateKey(KeyInitNode): @classmethod @@ -35,6 +41,7 @@ class PEMPublicKey(KeyInitNode): loaded_key = serialization.load_pem_public_key(keyfile) return (loaded_key,) + # DER keys class DERPrivateKey(KeyInitNode): @classmethod @@ -53,12 +60,14 @@ class DERPublicKey(KeyInitNode): loaded_key = serialization.load_der_public_key(keyfile) return (loaded_key,) -# SSH keys and fingerprint + +# SSH keys and fingerprint class SSHPublicKey(KeyInitNode): def execute(self, keyfile): loaded_key = serialization.load_ssh_public_key(keyfile) return (loaded_key,) + class SSHPublicKeyFingerprint(KeyInitNode): RETURN_TYPES = ("BYTESLIKE",) RETURN_NAMES = ("fingerprint",) @@ -66,7 +75,12 @@ class SSHPublicKeyFingerprint(KeyInitNode): @classmethod def INPUT_TYPES(cls): class_input = super().INPUT_TYPES() - class_input["required"]["hash_algorithm"] = (["MD5", "SHA-1", "SHA-256"], {"default": "SHA-256",}) + class_input["required"]["hash_algorithm"] = ( + ["MD5", "SHA-1", "SHA-256"], + { + "default": "SHA-256", + }, + ) return class_input def execute(self, keyfile, hash_algorithm): @@ -81,6 +95,7 @@ class SSHPublicKeyFingerprint(KeyInitNode): signature = serialization.ssh_key_fingerprint(loaded_key, algorithm) return (bytes(signature),) + class SSHPrivateKey(KeyInitNode): @classmethod def INPUT_TYPES(cls): @@ -93,8 +108,12 @@ class SSHPrivateKey(KeyInitNode): loaded_key = serialization.load_ssh_private_key(keyfile, password, unsafe_skip_rsa_key_validation=skip_rsa_key_validation) return (loaded_key,) + class SSHPublicIdentity(KeyInitNode): - RETURN_TYPES = ("BYTESLIKE", "BYTESLIKE",) + RETURN_TYPES = ( + "BYTESLIKE", + "BYTESLIKE", + ) RETURN_NAMES = ("certificate", "key") def execute(self, keyfile): diff --git a/src/cryptography_primitives/asymmetric_encryption/xdh.py b/src/cryptography_primitives/asymmetric_encryption/xdh.py index 53899ec..8d9647b 100644 --- a/src/cryptography_primitives/asymmetric_encryption/xdh.py +++ b/src/cryptography_primitives/asymmetric_encryption/xdh.py @@ -1,4 +1,5 @@ # Elliptic curve (X/EC) Diffie-Hellman (DH) key exchange + from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import x25519, x448 diff --git a/src/cryptography_primitives/auth_encrypt.py b/src/cryptography_primitives/auth_encrypt.py index 2241293..6443c5c 100644 --- a/src/cryptography_primitives/auth_encrypt.py +++ b/src/cryptography_primitives/auth_encrypt.py @@ -1,3 +1,5 @@ +# Authenticated encryption with associated data (AEAD) + from cryptography.hazmat.primitives.ciphers import aead @@ -222,9 +224,7 @@ class AESAuthKeygen: elif aes_type in generators: key = generators[aes_type].generate_key(bit_length) else: - raise ValueError( - f"Invalid AES type chosen. Perhaps the node is broken? Try making a new one, as this is normally impossible. Currently, it's {aes_type}, which is wrong." - ) + raise ValueError(f"Invalid AES type chosen. Perhaps the node is broken? Try making a new one, as this is normally impossible. Currently, it's {aes_type}, which is wrong.") return (key,) diff --git a/src/cryptography_primitives/const_time.py b/src/cryptography_primitives/const_time.py index c82dd54..6127eca 100644 --- a/src/cryptography_primitives/const_time.py +++ b/src/cryptography_primitives/const_time.py @@ -1,3 +1,5 @@ +# Constant time compare, the cryptography way + from cryptography.hazmat.primitives import constant_time diff --git a/src/cryptography_primitives/hashing.py b/src/cryptography_primitives/hashing.py index cf9dd0e..2b4b5b3 100644 --- a/src/cryptography_primitives/hashing.py +++ b/src/cryptography_primitives/hashing.py @@ -1,3 +1,5 @@ +# Message digests/hashes + from cryptography.hazmat.primitives import hashes diff --git a/src/cryptography_primitives/kdf.py b/src/cryptography_primitives/kdf.py index 6572dfc..82d79d6 100644 --- a/src/cryptography_primitives/kdf.py +++ b/src/cryptography_primitives/kdf.py @@ -1,3 +1,5 @@ +# Key derivation functions + from cryptography.hazmat.primitives import hashes # For the algorithm functions # Import hell @@ -118,9 +120,7 @@ class Argon2id_Derive(KeyDerivationNodes): def _get_kdf(self, salt, length, iterations, parallel_lanes, memory_cost, ad, secret): min_memory_cost = 8 * parallel_lanes if memory_cost < min_memory_cost: - print( - f"[ComfyUI ARG Toolkit][WARNING]: Because defined memory cost ({memory_cost} KiB) is lower than the minimum required ({min_memory_cost} KiB), the value will be silently clamped to the minimum." - ) + print(f"[ComfyUI ARG Toolkit][WARNING]: Because defined memory cost ({memory_cost} KiB) is lower than the minimum required ({min_memory_cost} KiB), the value will be silently clamped to the minimum.") memory_cost = min_memory_cost return Argon2id( @@ -169,9 +169,7 @@ class Argon2id_Verify(Argon2id_Derive): except InvalidKey: output = False except AlreadyFinalized: - raise RuntimeError( - "The verification function is being called more than once on a finalized KDF target. This should not be happening." - ) + raise RuntimeError("The verification function is being called more than once on a finalized KDF target. This should not be happening.") else: # For some stupid reason, PHC-formatted strings have to be done in string form, not hex or base64... try: @@ -180,9 +178,7 @@ class Argon2id_Verify(Argon2id_Derive): except InvalidKey: output = False except AlreadyFinalized: - raise RuntimeError( - "The verification function is being called more than once on a finalized KDF target. This should not be happening." - ) + raise RuntimeError("The verification function is being called more than once on a finalized KDF target. This should not be happening.") return (output,) @@ -261,9 +257,7 @@ class PBKDF2HMAC_Verify(PBKDF2HMAC_Derive): except InvalidKey: output = False except AlreadyFinalized: - raise RuntimeError( - "The verification function is being called more than once on a finalized KDF target. This should not be happening." - ) + raise RuntimeError("The verification function is being called more than once on a finalized KDF target. This should not be happening.") return (output,) @@ -337,9 +331,7 @@ class Scrypt_Verify(Scrypt_Derive): except InvalidKey: output = False except AlreadyFinalized: - raise RuntimeError( - "The verification function is being called more than once on a finalized KDF target. This should not be happening." - ) + raise RuntimeError("The verification function is being called more than once on a finalized KDF target. This should not be happening.") return (output,) @@ -412,9 +404,7 @@ class ConcatKDFHash_Verify(ConcatKDFHash_Derive): except InvalidKey: output = False except AlreadyFinalized: - raise RuntimeError( - "The verification function is being called more than once on a finalized KDF target. This should not be happening." - ) + raise RuntimeError("The verification function is being called more than once on a finalized KDF target. This should not be happening.") return (output,) @@ -466,9 +456,7 @@ class ConcatKDFHMAC_Verify(ConcatKDFHMAC_Derive): except InvalidKey: output = False except AlreadyFinalized: - raise RuntimeError( - "The verification function is being called more than once on a finalized KDF target. This should not be happening." - ) + raise RuntimeError("The verification function is being called more than once on a finalized KDF target. This should not be happening.") return (output,) @@ -548,9 +536,7 @@ class HKDF_Verify(HKDF_Derive): except InvalidKey: output = False except AlreadyFinalized: - raise RuntimeError( - "The verification function is being called more than once on a finalized KDF target. This should not be happening." - ) + raise RuntimeError("The verification function is being called more than once on a finalized KDF target. This should not be happening.") return (output,) @@ -623,9 +609,7 @@ class HKDFExpand_Verify(HKDFExpand_Derive): except InvalidKey: output = False except AlreadyFinalized: - raise RuntimeError( - "The verification function is being called more than once on a finalized KDF target. This should not be happening." - ) + raise RuntimeError("The verification function is being called more than once on a finalized KDF target. This should not be happening.") return (output,) @@ -698,9 +682,7 @@ class X963KDF_Verify(X963KDF_Derive): except InvalidKey: output = False except AlreadyFinalized: - raise RuntimeError( - "The verification function is being called more than once on a finalized KDF target. This should not be happening." - ) + raise RuntimeError("The verification function is being called more than once on a finalized KDF target. This should not be happening.") return (output,) @@ -861,9 +843,7 @@ class KBKDF_Verify(KBKDF_Derive): except InvalidKey: output = False except AlreadyFinalized: - raise RuntimeError( - "The verification function is being called more than once on a finalized KDF target. This should not be happening." - ) + raise RuntimeError("The verification function is being called more than once on a finalized KDF target. This should not be happening.") return (output,) diff --git a/src/cryptography_primitives/key_wrapper.py b/src/cryptography_primitives/key_wrapper.py index 75e20ce..2eabaef 100644 --- a/src/cryptography_primitives/key_wrapper.py +++ b/src/cryptography_primitives/key_wrapper.py @@ -1,5 +1,8 @@ +# Key wrapping + from cryptography.hazmat.primitives import keywrap + class WrapKeyNodes: def __init__(self): pass diff --git a/src/cryptography_primitives/symm_padding.py b/src/cryptography_primitives/symm_padding.py index 635d7c5..5be5f72 100644 --- a/src/cryptography_primitives/symm_padding.py +++ b/src/cryptography_primitives/symm_padding.py @@ -1,5 +1,8 @@ +# Symmetric padding for block ciphers + from cryptography.hazmat.primitives import padding + class PaddingNode: CATEGORY = "ARG Toolkit/Cryptography/Modern/Padding" diff --git a/src/cryptography_primitives/symmetrical_encrypt.py b/src/cryptography_primitives/symmetrical_encrypt.py index e6eedff..ce8233b 100644 --- a/src/cryptography_primitives/symmetrical_encrypt.py +++ b/src/cryptography_primitives/symmetrical_encrypt.py @@ -1,3 +1,5 @@ +# Symmetric encryption, including all current decrepit modes and algorithms available + from cryptography.hazmat.decrepit.ciphers import algorithms as decrepit_algorithms from cryptography.hazmat.decrepit.ciphers import modes as decrepit_modes from cryptography.hazmat.primitives.ciphers import algorithms, Cipher @@ -121,7 +123,7 @@ class EncryptDecrypt(CipherNodes): if algorithm == "ChaCha20": algorithm = algorithms.ChaCha20(key, nonce) elif algorithm in ["ARC4", "Blowfish", "CAST5", "SEED", "IDEA", "Camellia"]: - warnings.warn(f'[ComfyUI-ARG-Toolkit][WARNING]: Current algorithm used in Symmetric Encrypt/Decrypt ({algorithm}) is a decrepit type and is only meant for full inclusion. Unless you have a specific use case for this, please switch to a more reasonable algorithm.') + warnings.warn(f"[ComfyUI-ARG-Toolkit][WARNING]: Current algorithm used in Symmetric Encrypt/Decrypt ({algorithm}) is a decrepit type and is only meant for full inclusion. Unless you have a specific use case for this, please switch to a more reasonable algorithm.") algorithm = getattr(decrepit_algorithms, algorithm)(key) else: algorithm = getattr(algorithms, algorithm)(key) @@ -133,7 +135,7 @@ class EncryptDecrypt(CipherNodes): elif modes == "XTS": modes = cipher_modes.XTS(tweak=iv) elif modes == "ECB": - warnings.warn(f'[ComfyUI-ARG-Toolkit][WARNING]: Mode used in Symmetric Encrypt/Decrypt ({modes}) is inherently insecure and should never be used in anything in production. Unless you have a specific use case for this, please switch to a more reasonable mode.') + warnings.warn(f"[ComfyUI-ARG-Toolkit][WARNING]: Mode used in Symmetric Encrypt/Decrypt ({modes}) is inherently insecure and should never be used in anything in production. Unless you have a specific use case for this, please switch to a more reasonable mode.") modes = cipher_modes.ECB() elif modes == "None": modes = None diff --git a/src/debugging_nodes.py b/src/debugging_nodes.py index b8de11f..208036a 100644 --- a/src/debugging_nodes.py +++ b/src/debugging_nodes.py @@ -1,3 +1,5 @@ +# Debugging booleans in case of weirdness + class BooleanOutputter: @classmethod def INPUT_TYPES(cls): @@ -23,9 +25,7 @@ class BooleanOutputter: elif not boolean: return (false_text,) else: - raise ValueError( - "How did this even happen? Unless Python or ComfyUI does not handle booleans properly, *this should never happen*. If this does happen, please report this as a bug immediately." - ) + raise ValueError("How did this even happen? Unless Python or ComfyUI does not handle booleans properly, *this should never happen*. If this does happen, please report this as a bug immediately.") NODE_CLASS_MAPPINGS = { diff --git a/src/error_correction.py b/src/error_correction.py index 120670b..89289a2 100644 --- a/src/error_correction.py +++ b/src/error_correction.py @@ -1,3 +1,5 @@ +# Reed-Solomon function on bytes + from reedsolo import RSCodec @@ -22,7 +24,6 @@ class InitNode: } - class ReedSolomonEncode(InitNode): RETURN_TYPES = ("BYTESLIKE",) RETURN_NAMES = ("encoded_data",) diff --git a/src/fernet.py b/src/fernet.py index b1aa5c4..b678927 100644 --- a/src/fernet.py +++ b/src/fernet.py @@ -1,4 +1,5 @@ -# Fernet encryption, the basics of basics +# Fernet encryption, the basics of basics in modern cryptography + from cryptography import fernet diff --git a/src/morse_code.py b/src/morse_code.py index db62d7c..e1ee01b 100644 --- a/src/morse_code.py +++ b/src/morse_code.py @@ -1,4 +1,4 @@ -# This implementation is lifted from https://github.com/Aayush9029/encodeDecode, which is left abandoned since 2020. This will attempt to roll that implementation in while implementing more languages that can be supported within reason (so no Chinese telegraph code for Chinese, unfortunately). +# This implementation of a Morse code encoder/decoder is lifted from https://github.com/Aayush9029/encodeDecode, which is left abandoned since 2020. This will attempt to roll that implementation in while implementing more languages that can be supported within reason (so no Chinese telegraph code for Chinese, unfortunately). import unicodedata diff --git a/src/rotor_machines.py b/src/rotor_machines.py new file mode 100644 index 0000000..a2f4d23 --- /dev/null +++ b/src/rotor_machines.py @@ -0,0 +1,719 @@ +# Rotor machine ciphers (Enigma, Typex, etc.) +from dataclasses import dataclass, replace +from abc import ABC, abstractmethod +from functools import lru_cache, cached_property + +# Base datatypes used for this node +Perm = tuple[int, ...] +Positions = tuple[int, ...] + +# Stage dataclass, for mapping symbols to other symbols +@dataclass(frozen=True) +class Stage(ABC): + @abstractmethod + def apply(self, pos: Positions, x: int) -> int: ... + + def validate(self, n: int) -> None: return None + + @abstractmethod + def inverted(self): ... + + @abstractmethod + def src_alphabet(self) -> str: ... + + @abstractmethod + def dst_alphabet(self) -> str: ... + + def bind(self, rotors): + return self + + +@dataclass(frozen=True) +class Fixed(Stage): + src: str; dst: str; perm: Perm + def apply(self, pos: Positions, x): return self.perm[x] + def src_alphabet(self): return self.src + def dst_alphabet(self): return self.dst + def inverted(self) -> Stage: + if len(self.src) != len(self.dst) or len(set(self.perm)) != len(self.dst): + raise ValueError("Stage given is not a bijection, so it has no inverse.") + return Fixed(self.dst, self.src, InitNode.perm_invert(self.perm)) + +@dataclass(frozen=True) +class Rotor: alphabet: str; fwd: Perm; inv: Perm; notches: frozenset[int]; ring: int = 0 + +@dataclass(frozen=True) +class Through(Stage): + rotor: Rotor + index: int + inverse: bool = False + def apply(self, pos, x): + n = len(self.rotor.alphabet) + s = pos[self.index] - self.rotor.ring + table = self.rotor.inv if self.inverse else self.rotor.fwd + return (table[(x + s) % n] - s) % n + def src_alphabet(self): return self.rotor.alphabet + def dst_alphabet(self): return self.rotor.alphabet + def inverted(self): return replace(self, inverse=not self.inverse) + +# Stepper dataclass +class Stepper(ABC): + @abstractmethod + def __call__(self, rotors: tuple[Rotor, ...], pos: Positions) -> Positions: ... + def validate(self, rotors: tuple[Rotor, ...]) -> None: return None + +@dataclass(frozen=True) +class Clause: + notched: tuple[int, ...] = () + +Rule = tuple[Clause, ...] + +@dataclass(frozen=True) +class SlotRef: + index: int + inverse: bool = False + def bind(self, rotors) -> Through: + if not 0 <= self.index < len(rotors): + raise ValueError(f"circuit references slot {self.index}, " + f"machine has {len(rotors)} rotors") + return Through(rotors[self.index], self.index, self.inverse) + +@dataclass(frozen=True) +class RuleStepper(Stepper): + rules: tuple[Rule, ...] + steps: tuple[int, ...] | None = None # +1 or -1 per rotor; None = all +1 + + def __call__(self, rotors, pos): + d = self.steps or (1,) * len(rotors) + at = tuple(p in r.notches for r, p in zip(rotors, pos)) + return tuple((p + d[i] * int(any(all(at[j] for j in c.notched) for c in rule))) % len(r.alphabet) for i, (p, r, rule) in enumerate(zip(pos, rotors, self.rules))) + + def validate(self, rotors: tuple[Rotor, ...],) -> None: + if len(self.rules) != len(rotors): + raise ValueError("Rules must have one entry per rotor.") + if self.steps is not None and (len(self.steps) != len(rotors) + or any(s not in (1, -1) for s in self.steps)): + raise ValueError("steps must be +1 or -1, one per rotor") + for rule in self.rules: + for c in rule: + if any(not 0 <= j < len(rotors) for j in c.notched): + raise ValueError("Clause references a nonexistent rotor.") + + @staticmethod + def parse(text: str) -> tuple[Rule, ...]: + entries = [e.strip() for e in text.replace(";", "\n").splitlines() if e.strip()] + rules: list[Rule] = [] + for slot, e in enumerate(entries): + low = e.lower() + if low == "never": + rules.append(()) + elif low == "always": + rules.append((Clause(),)) + else: + clauses = [] + for part in e.split("|"): + try: + idx = tuple(int(t) for t in part.split("&")) + except ValueError: + raise ValueError(f"Rotor {slot}: bad clause {part.strip()!r}") from None + clauses.append(Clause(idx)) + rules.append(tuple(clauses)) + return tuple(rules) + + @staticmethod + def parse_reverse(text: str, count: int) -> tuple[int, ...] | None: + toks = text.replace(",", " ").split() + if not toks: + return None + try: + idx = {int(t) for t in toks} + except ValueError: + raise ValueError(f"bad reverse list {text!r}; expected slot numbers") from None + bad = sorted(i for i in idx if not 0 <= i < count) + if bad: + raise ValueError(f"reverse slots out of range 0..{count - 1}: {bad}") + return tuple(-1 if i in idx else 1 for i in range(count)) + + @staticmethod + def format_rule(rule) -> str: + if not rule: + return "never" + if any(not c.notched for c in rule): + return "always" + return "|".join("&".join(map(str, c.notched)) for c in rule) + +@dataclass(frozen=True) +class CircuitStepper(Stepper): + base: RuleStepper # steps the side rotors by notch rules + circuit: tuple[Fixed | SlotRef, ...] + inputs: str # live wires entering the circuit + targets: tuple[tuple[str, int], ...] # (output symbol, slot it advances) + after: bool = False # evaluate on base-stepped positions? + + def __call__(self, rotors, pos): + stepped = self.base(rotors, pos) + ref = stepped if self.after else pos + hops = [h.bind(rotors) for h in self.circuit] + first, last = hops[0].src_alphabet(), hops[-1].dst_alphabet() + drive = dict(self.targets) + adv = set() + for c in self.inputs: + x = first.index(c) + for h in hops: + x = h.apply(ref, x) + if last[x] in drive: + adv.add(drive[last[x]]) + dirs = self.base.steps or (1,) * len(rotors) + return tuple((p + dirs[i] * int(i in adv)) % len(r.alphabet)for i, (p, r) in enumerate(zip(stepped, rotors))) + + def validate(self, rotors): + self.base.validate(rotors) + hops = [h.bind(rotors) for h in self.circuit] # also range-checks slots + for i, (a, b) in enumerate(zip(hops, hops[1:])): + if a.dst_alphabet() != b.src_alphabet(): + raise ValueError(f"circuit hop {i} outputs {a.dst_alphabet()!r} but " + f"hop {i + 1} expects {b.src_alphabet()!r}") + if not self.inputs or any(c not in hops[0].src_alphabet() for c in self.inputs): + raise ValueError("inputs must be a non-empty set of circuit-entry symbols") + for c, slot in self.targets: + if c not in hops[-1].dst_alphabet(): + raise ValueError(f"target symbol {c!r} not in the circuit's output alphabet") + if not 0 <= slot < len(rotors): + raise ValueError(f"target slot {slot} out of range") + if self.base.rules[slot]: + raise ValueError(f"slot {slot} is driven by the circuit, " + "so its base rule must be 'never'") + +# Machine dataclass +@dataclass(frozen=True) +class Machine: + rotors: tuple[Rotor, ...]; path: tuple[Stage, ...]; stepper: Stepper; roles: tuple[str, ...] = () + + @cached_property + def inv_path(self) -> tuple[Stage, ...]: + return tuple(s.inverted() for s in reversed(self.path)) + + @property + def in_alphabet(self): + return self.path[0].src_alphabet() + + @property + def out_alphabet(self): + return self.path[-1].dst_alphabet() + + def press(self, pos: Positions, x: int, decrypt=False) -> tuple[Positions, int]: + pos = self.stepper(self.rotors, pos) # step BEFORE encrypting + for st in (self.path if not decrypt else self.inv_path): + x = st.apply(pos, x) + return pos, x + + def encode_pos(self, s: str) -> Positions: + if len(s) != len(self.rotors): + raise ValueError(f"positions has {len(s)} entries, machine has {len(self.rotors)} rotors") + out = [] + for i, (r, c) in enumerate(zip(self.rotors, s)): + j = r.alphabet.find(c) + if j < 0: + raise ValueError(f"rotor {i}: {c!r} is not in its alphabet") + out.append(j) + return tuple(out) + + def decode_pos(self, pos: Positions) -> str: + return "".join(r.alphabet[p] for r, p in zip(self.rotors, pos)) + + @staticmethod + def describe_stage(s: Stage) -> str: + if isinstance(s, Through): + return f"rotor slot {s.index} {'inverse' if s.inverse else 'forward'}" + if isinstance(s, Fixed): + if len(s.src) != len(s.dst): + return f"fixed {len(s.src)}->{len(s.dst)}, changes alphabet" + n = len(s.perm) + tag = ("involution" if all(s.perm[s.perm[j]] == j for j in range(n)) + else "not an involution") + fp = sum(s.perm[j] == j for j in range(n)) + return f"fixed {n}->{n}, {tag}" + (f", {fp} fixed points" if fp else "") + return f"{type(s).__name__} {s.src_alphabet()!r} -> {s.dst_alphabet()!r}" + + + def describe(self) -> str: + ia, oa = self.in_alphabet, self.out_alphabet + out = [f"Alphabet: {ia}" if ia == oa else f"Input alphabet: {ia}\nOutput alphabet: {oa}"] + + # resolve the stepper to something printable + st = self.stepper + circ = st if isinstance(st, CircuitStepper) else None + base = circ.base if circ else st + if isinstance(base, PresetStepper): + base = _preset(base.kind, base.k, len(self.rotors)) + rules = base.rules if isinstance(base, RuleStepper) else None + dirs = (base.steps if isinstance(base, RuleStepper) else None) or (1,) * len(self.rotors) + driven = {} + if circ: + for sym, slot in circ.targets: + driven.setdefault(slot, []).append(sym) + + out.append(f"\nSlots ({len(self.rotors)}; the positions string needs " + f"{len(self.rotors)} characters, slot 0 first):") + out.append(" slot role size ring notches step rule") + for i, r in enumerate(self.rotors): + role = self.roles[i] if i < len(self.roles) else "?" + notches = "".join(r.alphabet[j] for j in sorted(r.notches)) or "-" + rule = RuleStepper.format_rule(rules[i]) if rules else "?" + if i in driven: + rule += f" + circuit({''.join(driven[i])})" + out.append(f" {i:<5} {role:<10} {len(r.alphabet):<5} {r.ring:<5} " + f"{notches:<10} {dirs[i]:+d} {rule}") + + if circ: + hops = " ".join(f"{h.index}{chr(39) if h.inverse else ''}" if isinstance(h, SlotRef) else f"map({len(h.src)}->{len(h.dst)})" for h in circ.circuit) + out.append(f"\nCircuit: {hops} Live inputs: {circ.inputs} after_base: {circ.after}") + + out.append(f"\nPath ({len(self.path)} stages):") + for k, s in enumerate(self.path, 1): + out.append(f" {k}. {Machine.describe_stage(s)}") + + notes = [] + if rules: + notes += [f"slot {i} never steps on its own" for i, ru in enumerate(rules) if not ru and i not in driven] + if not any(isinstance(s, Fixed) and s is not self.path[0] for s in self.path[1:]) and \ + len(self.path) == len(set(map(id, self.path))): + notes.append("path has no repeated stages: likely non-reflecting, so decrypt mode is required") + if notes: + out.append("\nNotes:\n" + "\n".join(f" - {n}" for n in notes)) + return "\n".join(out) + +@dataclass(frozen=True) +class StartPos: + text: str + + + +#================================================================== +# +# From this point onward, the ComfyUI node stuff begins +# +#=================================================================== + +class InitNode: + CATEGORY = "ARG Toolkit/Cryptography/Classical/Rotor Machines" + FUNCTION = "execute" + COMMON_ALIASES = ["rotor machine", "rotor cipher", "cipher machine", "enigma", "encryption"] + + # Check a set of wiring against the alphabet + @staticmethod + def perm_wiring(alphabet: str, string:str, type: str) -> Perm: + index = {char: i for i, char in enumerate(alphabet)} + bad = sorted({char for char in string if char not in index}) + if bad: + raise ValueError(f"{type}: symbols not in alphabet: {bad}") + return tuple(index[c] for c in string) + + # Inversion of a permutation for decrypting a message + @staticmethod + def perm_invert(permutation: Perm) -> Perm: + output = [0] * len(permutation) + for index, variant in enumerate(permutation): + output[variant] = index + return tuple(output) + + @staticmethod + def str_cleanup(string: str): + return "".join(string.split()) + + def odometer(k: int, total: int | None = None) -> RuleStepper: + """Rightmost k rotors carry on notch; the rest are fixed.""" + total = k if total is None else total + first = total - k + rules: list[Rule] = [()] * first + for i in range(first, total): + rules.append((Clause(tuple(range(i + 1, total))),)) + return RuleStepper(tuple(rules)) + + @staticmethod + def enigma(k: int = 3, total: int | None = None) -> RuleStepper: + total = total or k + first = total - k + rules: list[Rule] = [()] * first + for i in range(first, total): + if i == total - 1: + rules.append((Clause(),)) + elif i == first: + rules.append((Clause((i + 1,)),)) + else: + rules.append((Clause((i + 1,)), Clause((i,)))) + return RuleStepper(tuple(rules)) + + @staticmethod + def check_alphabet(alphabet: str) -> int: + if not alphabet or len(set(alphabet)) != len(alphabet): + raise ValueError("alphabet must be non-empty with unique symbols") + return len(alphabet) + + @staticmethod + def _fold(c: str) -> str: + u = c.upper() + return u if len(u) == 1 else c + + @staticmethod + def parse_circuit(text, maps): + hops = [] + for tok in text.split(): + if tok in maps: + if maps[tok] is None: + raise ValueError(f"Circuit uses {tok!r} but no map is connected to it") + hops.append(maps[tok]) + else: + try: + hops.append(SlotRef(int(tok.rstrip("'")), tok.endswith("'"))) + except ValueError: + raise ValueError(f"Bad circuit token {tok!r}") from None + if not hops: + raise ValueError("Circuit is empty.") + return tuple(hops) + + @staticmethod + def parse_targets(text): + out = [] + for group in text.split(): + syms, sep, slot = group.partition("=") + try: + if not sep: + raise ValueError + n = int(slot) + except ValueError: + raise ValueError(f"Bad target group {group!r}; expected symbols=slot") from None + out += [(c, n) for c in syms] + chars = [c for c, _ in out] + if len(set(chars)) != len(chars): + raise ValueError("An output symbol may drive only one rotor") + return tuple(out) + + +## Rotors +class RotorCipherRotor(InitNode): + RETURN_TYPES = ("ROTORS",) + RETURN_NAMES = ("rotor",) + SEARCH_ALIASES = InitNode.COMMON_ALIASES + ["rotor", "wheel", "scrambler"] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "alphabet": ("STRING", {"default": "ABCDEFGHIJKLMNOPQRSTUVWXYZ"}), + "wiring": ("STRING", {"default": "EKMFLGDQVZNTOWYHXUSPAIBRCJ"}), + "notches": ("STRING", {"default": "Q"}), + "ring": ("INT", {"default": 0, "min": 0, "max": 4096}), + "invert": ("BOOLEAN", {"default": False}), + }, + "optional": { + "previous_rotor": ("ROTORS", {"forceInput": True}) + } + } + + def execute(self, alphabet, wiring, notches, ring, invert, previous_rotor=None): + alphabet = InitNode.str_cleanup(alphabet) + n = InitNode.check_alphabet(alphabet=alphabet) + + w = InitNode.perm_wiring(alphabet, wiring, type="Wiring") + if len(w) != n or len(set(w)) != n: + raise ValueError(f"Wiring must use each of the {n} alphabet symbols exactly once") + + if not 0 <= ring < n: + raise ValueError(f"Ring must be in [0, {n - 1}]") + + fwd, inv = w, InitNode.perm_invert(w) + if invert: + fwd, inv = inv, fwd + + rotor= Rotor(alphabet, fwd, inv, frozenset(InitNode.perm_wiring(alphabet, InitNode.str_cleanup(notches), "Notches")), ring) + return (((*previous_rotor, rotor) if previous_rotor else (rotor,)),) + +## Fixed permutations +class RotorCipherPerm(InitNode): + RETURN_TYPES = ("PERMUTATION",) + RETURN_NAMES = ("perm_obj",) + SEARCH_ALIASES = InitNode.COMMON_ALIASES + ["reflector", "entry wheel", "etw", "ukw", "permutation", "wiring"] + + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "alphabet": ("STRING", {"default": "ABCDEFGHIJKLMNOPQRSTUVWXYZ"}), + "wiring": ("STRING", {"default": "EKMFLGDQVZNTOWYHXUSPAIBRCJ"}) + } + } + + def execute(self, alphabet, wiring): + alphabet = InitNode.str_cleanup(alphabet) + n = len(alphabet) + w = InitNode.perm_wiring(alphabet, InitNode.str_cleanup(wiring), "Wiring") + if len(w) != n or len(set(w)) != n: + raise ValueError (f"Wiring must use every symbols exactly once. Required {n}, got {len(w)}.") + return (Fixed(alphabet, alphabet, w),) + +## Mapper node for cases that needs to convert between different alphabet sizes +class RotorCipherMap(InitNode): + RETURN_TYPES = ("PERMUTATION",) + RETURN_NAMES = ("perm_obj",) + SEARCH_ALIASES = InitNode.COMMON_ALIASES + ["alphabet map", "grouping", "mapping stage", "relabel", "permutation"] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "src_alphabet": ("STRING", {"default": "ABCEDEFGHIJKLMNOPQRSTUVWXYZ"}), + "dst_alphabet": ("STRING", {"default": "ABCEDEFGHIJKLMNOPQRSTUVWXYZ"}), + "mapping": ("STRING", {"default": ""}) + } + } + + def execute(self, src_alphabet, dst_alphabet, mapping): + src, dst = map(InitNode.str_cleanup, (src_alphabet, dst_alphabet)) + InitNode.check_alphabet(src); InitNode.check_alphabet(dst) + m = InitNode.perm_wiring(dst, InitNode.str_cleanup(mapping), "Mapping") + if len(m) != len(src): + raise ValueError(f"Mapping needs one symbol per source symbol ({len(src)})") + return (Fixed(src, dst, m),) + +# Machine node to construct the necessary machine to encrypt/decrypt + +class RotorCipherMachine(InitNode): + RETURN_TYPES = ("MACHINE",) + RETURN_NAMES = ("cipher_machine",) + SEARCH_ALIASES = InitNode.COMMON_ALIASES + ["assemble machine", "rotor stack", "build machine"] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "rotors": ("ROTORS", {"forceInput": True}), + "stepper": ("STEPPER", {"forceInput": True}) + }, + "optional": { + "entry": ("PERMUTATION", {"forceInput": True, "tooltip": "Entry wheel (ETW). Applied first on the way in, and inverted on the way out. Identity for Wehrmacht Enigmas."}), + "reflector": ("PERMUTATION", {"forceInput": True, "tooltip": "Static reflector (UKW). Must be an involution. Use this OR reflector_rotor, not both."}), + "reflector_rotor": ("PERMUTATION", {"forceInput": True, "tooltip": "Reflector that takes a position (slot 0). A chain of exactly one rotor with an involutive wiring."}), + "plugboard": ("PERMUTATION", {"forceInput": True, "tooltip": "Steckerbrett. Applied outermost, before the entry wheel. Must be an involution."}), + "auxilliary_rotors": ("ROTORS", {"forceInput": True}), + }, + } + + def execute(self, rotors, stepper, entry=None, reflector=None, reflector_rotor=None, plugboard=None, auxilliary_rotors=None): + if reflector is not None and reflector_rotor is not None: + raise ValueError("Use either a static reflector or a reflector rotor, not both.") + + ref_rot = None + if reflector_rotor is not None: + if len(reflector_rotor) != 1 or not isinstance(reflector_rotor[0], Rotor): + raise ValueError("reflector_rotor must be a chain of exactly one rotor.") + ref_rot = reflector_rotor[0] + if any(ref_rot.fwd[ref_rot.fwd[i]] != i for i in range(len(ref_rot.fwd))): + raise ValueError("Reflector rotor wiring must be an involution.") + + slot = 1 if ref_rot else 0 # slot 0 reserved for the reflector + stages = [] + for item in rotors: + if isinstance(item, Rotor): + stages.append(Through(item, slot)); slot += 1 + else: + stages.append(item) + path_rotors = tuple(i for i in rotors if isinstance(i, Rotor)) + aux_rotors = tuple(auxilliary_rotors) if auxilliary_rotors else () + if not all(isinstance(r, Rotor) for r in aux_rotors): + raise ValueError("aux must be a chain of rotors only") + if aux_rotors and isinstance(stepper, PresetStepper): + raise ValueError("Preset steppers count every rotor, aux included. Please use a rule or circuit stepper with aux rotors.") + rotors = ((ref_rot,) if ref_rot else ()) + path_rotors + aux_rotors + if not rotors: + raise ValueError("Machine needs at least one rotor") + + for name, f in (("reflector", reflector), ("plugboard", plugboard)): + if f is not None: + if f.src != f.dst: + raise ValueError(f"{name} must map an alphabet to itself") + if any(f.perm[f.perm[i]] != i for i in range(len(f.perm))): + raise ValueError(f"{name} must be an involution") + + pre = [s for s in (plugboard, entry) if s is not None] + fwd = (*pre, *reversed(stages)) + mid = Through(ref_rot, 0) if ref_rot else reflector + path = fwd if mid is None else (*fwd, mid, *(s.inverted() for s in reversed(fwd))) + + for i, (a, b) in enumerate(zip(path, path[1:])): + if a.dst_alphabet() != b.src_alphabet(): + raise ValueError(f"stage {i} outputs {a.dst_alphabet()!r} but stage {i + 1} " + f"expects {b.src_alphabet()!r}; add a mapping stage between them") + stepper.validate(rotors) + roles = (("reflector",) if ref_rot else ()) + ("path",) * len(path_rotors) + ("aux",) * len(aux_rotors) + return (Machine(rotors, path, stepper, roles),) + +# Rotor stepping rules + +@lru_cache(maxsize=None) +def _preset(kind: str, k: int, total: int) -> RuleStepper: + return (InitNode.enigma if kind == "Enigma" else InitNode.odometer)(k, total) + +@dataclass(frozen=True) +class PresetStepper(Stepper): + kind: str # "Enigma" | "Odometer" + k: int + + def __call__(self, rotors, pos): + return _preset(self.kind, self.k, len(rotors))(rotors, pos) + + def validate(self, rotors): + if not 1 <= self.k <= len(rotors): + raise ValueError(f"k must be in [1, {len(rotors)}] for this machine") + _preset(self.kind, self.k, len(rotors)).validate(rotors) + +class RotorCipherPresetStepper(InitNode): + RETURN_TYPES = ("STEPPER",) + RETURN_NAMED = ("stepper",) + + @classmethod + def INPUT_TYPES(cls): + return {"required": { + "type": (["Enigma", "Odometer"], {"default": "Enigma",}), + "k": ("INT", {"default": 3, "min": 1, "max": 64, "tooltip": "The rightmost k rotor steps."}), + }} + + def execute(self, type, k): + return (PresetStepper(type, k),) + + +class RotorCipherRuleStepper(InitNode): + RETURN_TYPES = ("STEPPER",) + RETURN_NAMES = ("stepper",) + SEARCH_ALIASES = InitNode.COMMON_ALIASES + ["stepping", "turnover", "notch rules", "odometer", "double step"] + + @classmethod + def INPUT_TYPES(cls): + return {"required": { + "rules": ("STRING", {"multiline": True, "default": "1\n2|1\nalways"}), + "reverse": ("STRING", {"default": "", "multiline": False}) + }} + + def execute(self, rules: str, reverse=""): + parsed_rules = RuleStepper.parse(rules) + return (RuleStepper(parsed_rules, RuleStepper.parse_reverse(reverse, len(parsed_rules))),) + +class RotorCipherPosition(InitNode): + RETURN_TYPES = ("POSITIONS",) + RETURN_NAMES = ("rotor_pos",) + SEARCH_ALIASES = InitNode.COMMON_ALIASES + ["start positions", "key setting", "grundstellung", "window letters"] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "positions": ("STRING", {"default": "AAA", "tooltip": "From left to right, where the rotors will start at, whitespace permitting."}), + } + } + + def execute(self, positions): + position = self.str_cleanup(positions) + if not position: + raise ValueError("Positions must list at least one rotor.") + return (StartPos(position),) + +class RotorCipherCircuitStepper(InitNode): + RETURN_TYPES = ("STEPPER",) + RETURN_NAMES = ("circuit_stepper",) + SEARCH_ALIASES = InitNode.COMMON_ALIASES + ["irregular stepping", "control circuit", "driven stepping"] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "base": ("STEPPER", {"forceInput": True}), + "circuit": ("STRING", {"default": "3"}), + "inputs": ("STRING", {"default": "A"}), + "targets": ("STRING", {"default": "E=2"}), + "evaluate_after_base": ("BOOLEAN", {"default": False}), + }, + "optional": { + "map_a": ("PERMOBJ", {"forceInput": True}), + "map_b": ("PERMOBJ", {"forceInput": True}) + } + } + + def execute(self, base, circuit, inputs, targets, evaluate_after_base, map_a=None, map_b=None): + if not isinstance(base, RuleStepper): + raise ValueError("base must be a rule stepper (it steps the side rotors)") + return (CircuitStepper(base, InitNode.parse_circuit(circuit, {"a": map_a, "b": map_b}), InitNode.str_cleanup(inputs), InitNode.parse_targets(targets), evaluate_after_base),) + + +# Encrypt/Decrypt function node +class RotorCipherEncrypt(InitNode): + RETURN_TYPES = ("STRING", "POSITIONS",) + RETURN_NAMES = ("encrypted_msg", "rotor_pos",) + SEARCH_ALIASES = InitNode.COMMON_ALIASES + ["encrypt", "decrypt", "cipher text", "plaintext", "run machine"] + + @classmethod + def INPUT_TYPES(cls): + return {"required": { + "machine": ("MACHINE", {"forceInput": True,}), + "positions": ("POSITIONS", {"forceInput": True,}), + "text": ("STRING", {"multiline": True, "default": "Hello World!", "placeholder": "Type your message here..."}), + "preserve": ("BOOLEAN", {"default": False, "tooltip": "When set to true, attempt to preserve the form of the message, including case and punctuation marks."}), + "mode": ("BOOLEAN", {"default": False, "label_on": "Decrypt", "label_off": "Encrypt", "tooltip": "Whether to encrypt or decrypt the message."}) + }} + + def execute(self, machine: Machine, positions, text, preserve, mode=False): + src, dst = ((machine.out_alphabet, machine.in_alphabet) if mode else (machine.in_alphabet, machine.out_alphabet)) + for a in (src, dst): + if any(InitNode._fold(c) != c for c in a): + raise ValueError("alphabets must be uppercase/uncased: input is case-folded") + + idx = {c: i for i, c in enumerate(src)} + pos = machine.encode_pos(positions.text) + + out = [] + for ch in text: + i = idx.get(InitNode._fold(ch)) + if i is None: + if preserve: + out.append(ch) + continue + pos, y = machine.press(pos, i, mode) + c = dst[y] + out.append(c.lower() if preserve and ch != InitNode._fold(ch) else c) + return ("".join(out), StartPos(machine.decode_pos(pos))) + +class RotorCipherDescribe(InitNode): + RETURN_TYPES = ("STRING",) + RETURN_TYPES = ("description",) + SEARCH_ALIASES = InitNode.COMMON_ALIASES + ["describe", "inspect", "debug", "machine layout", "slots"] + OUTPUT_NODE = True + + @classmethod + def INPUT_TYPES(cls): return {"required": {"machine": ("MACHINE",)}} + + def execute(self, machine: Machine): + return (machine.describe(),) + +NODE_CLASS_MAPPINGS = { + "RotorCipherRotor": RotorCipherRotor, + "RotorCipherPerm": RotorCipherPerm, + "RotorCipherMachine": RotorCipherMachine, + "RotorCipherRuleStepper": RotorCipherRuleStepper, + "RotorCipherPresetStepper": RotorCipherPresetStepper, + "RotorCipherPosition": RotorCipherPosition, + "RotorCipherEncrypt": RotorCipherEncrypt, + "RotorCipherMapper": RotorCipherMap, + "RotorCipherDescribe": RotorCipherDescribe, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "RotorCipherRotor": "Rotor Machine Rotor", + "RotorCipherPerm": "Rotor Machine Permutation", + "RotorCipherMachine": "Rotor Machine State", + "RotorCipherRuleStepper": "Rotor Machine Rule-Based Stepper", + "RotorCipherPresetStepper": "Rotor Machine Preset Stepper", + "RotorCipherPosition": "Rotor Machine Positions", + "RotorCipherEncrypt": "Rotor Machine Encrypt/Decrypt", + "RotorCipherMapper": "Rotor Machine Rotor Mapper", + "RotorCipherDescribe": "Rotor Machine State Descriptor", +} diff --git a/src/spectrogram.py b/src/spectrogram.py index ef9aea0..c48ffde 100644 --- a/src/spectrogram.py +++ b/src/spectrogram.py @@ -1,3 +1,5 @@ +# Spectrogram encoder/decoder that cal also make images into spectrograms + import torch import torchaudio.transforms as torch_trans import torchaudio.functional as torch_fn @@ -185,12 +187,13 @@ class SpectrogramEncoder: def mel_to_hz(m): m = np.asarray(m, dtype=np.float64) return np.where(m < min_log_mel, m * f_sp, min_log_hz * np.exp(logstep * (m - min_log_mel))) + hz_per_bin = sampling_rate / p2_fft mel_max = hz_to_mel(nyquist) mel_num = int(np.round(mel_max / (2595.0 * np.log10(1.0 + hz_per_bin / 700.0)))) mel_num = int(np.clip(mel_num, 32, p2_fft // 2 + 1)) - mel_points = np.linspace(0, mel_max, mel_num +2) + mel_points = np.linspace(0, mel_max, mel_num + 2) hz_points = mel_to_hz(mel_points) stft_freqs = np.linspace(0.0, nyquist, freq_num, dtype=np.float32) @@ -230,13 +233,13 @@ class SpectrogramDecoder: def INPUT_TYPES(cls): return { "required": { - "image": ("IMAGE", {"forceInput": True,"tooltip": "The spectrogram image to decode into audio."}), + "image": ("IMAGE", {"forceInput": True, "tooltip": "The spectrogram image to decode into audio."}), "channel": ( ["colormap", "channel_red", "channel_green", "channel_blue", "luminance"], { "default": "luminance", "tooltip": "The color channel used to collapse and deconstruct the image. For spectrograms with a known colormap, use `colormap`, otherwise use `luminance` or any of the RGB channels.", - } + }, ), "window_size": ("FLOAT", {"default": 25.000, "min": 0.001, "max": 100000.000, "step": 0.0011, "tooltip": "The size of each 'frame' in miliseconds."}), "hop_size": ("FLOAT", {"default": 15.000, "min": 0.001, "max": 100000.000, "step": 0.001, "tooltip": "The size of the overlap between each 'frames' in miliseconds"}), @@ -252,13 +255,13 @@ class SpectrogramDecoder: { "default": "Logarithmic", "tooltip": "The scaling method used for the spectrogram.", - } + }, ), "autogain": ( "BOOLEAN", { "tooltip": "For Linear and Quadratic scaling method, dynamically scale the amplitude based on the peak. If disabled, use `max_amp` to define it instead. If used with Logarithmic Mel scaling, will enable area normalization if used with `mel_scale=slaney`.", - } + }, ), "top_db": ( "FLOAT", @@ -269,15 +272,42 @@ class SpectrogramDecoder: "tooltip": "The dynamic range (in negative dB) to scale against for both Logarithmic scaling methods.", }, ), - "sampling_rate": ("INT", {"max": 384000, "min":100, "default": 44100, "step": 1, "tooltip": "The target sample rate for the audio file."}), + "sampling_rate": ("INT", {"max": 384000, "min": 100, "default": 44100, "step": 1, "tooltip": "The target sample rate for the audio file."}), "max_iter": ("INT", {"min": 1, "max": 1024, "default": 32, "tooltip": "Maximum iteration to run the Griffin-Lim algorithm for."}), - "gl_momentum": ("FLOAT", {"min": 0.00, "max": 2.00, "default": 0.99, "step": 0.01, "tooltip": "For use with Fast Griffin-Lim. 0 disables Fast Griffin-Lim."}) + "gl_momentum": ("FLOAT", {"min": 0.00, "max": 2.00, "default": 0.99, "step": 0.01, "tooltip": "For use with Fast Griffin-Lim. 0 disables Fast Griffin-Lim."}), }, "optional": { - "colormap_low": ("STRING", {"default": "#000000", "tooltip": "The color of the lowest point of the spectrogram.",}), - "colormap_mid": ("STRING", {"default": "#B73779", "tooltip": "The color of the middle point of the spectrogram.",}), - "colormap_high": ("STRING", {"default": "#FCFDBF", "tooltip": "The color of the highest point of the spectrogram.",}), - "max_amp": ("FLOAT", {"step": 0.1, "min": 0.0, "max": 1.0, "default": 1.0, "tooltip": "If `autogain is disabled, this defines the maximum amplitude to scale with in Linear and Quadratic scaling.",}), + "colormap_low": ( + "STRING", + { + "default": "#000000", + "tooltip": "The color of the lowest point of the spectrogram.", + }, + ), + "colormap_mid": ( + "STRING", + { + "default": "#B73779", + "tooltip": "The color of the middle point of the spectrogram.", + }, + ), + "colormap_high": ( + "STRING", + { + "default": "#FCFDBF", + "tooltip": "The color of the highest point of the spectrogram.", + }, + ), + "max_amp": ( + "FLOAT", + { + "step": 0.1, + "min": 0.0, + "max": 1.0, + "default": 1.0, + "tooltip": "If `autogain is disabled, this defines the maximum amplitude to scale with in Linear and Quadratic scaling.", + }, + ), "mel_scale": ("BOOLEAN", {"label_off": "htk", "label_on": "slaney", "default": False, "tooltip": "Mel scaling method to use."}), }, } @@ -312,14 +342,13 @@ class SpectrogramDecoder: raw_frame_length = int(np.round((window_size / 1000.0) * sampling_rate)) hop_length = int(np.round((hop_size / 1000.0) * sampling_rate)) p2_fft = 2 ** int(np.ceil(np.log2(raw_frame_length))) - n_stft = (p2_fft//2)+1 - + n_stft = (p2_fft // 2) + 1 # Channel collapse if channel == "colormap": norm = self.colormap_collapse(image, colormap_low, colormap_mid, colormap_high) elif channel == "luminance": - weights = torch.tensor([0.299, 0.587, 0.114]) # Luminance value based on ITU-R BT.601 perceptual weights. + weights = torch.tensor([0.299, 0.587, 0.114]) # Luminance value based on ITU-R BT.601 perceptual weights. norm = torch.sum(image * weights, dim=-1) elif channel == "channel_red": norm = image[..., 0] @@ -335,25 +364,17 @@ class SpectrogramDecoder: elif scaling_method == "Quadratic": image_peak = torch.max(norm**2) - if autogain and image_peak>0: - magnitude = (norm/image_peak) * 1.0 + if autogain and image_peak > 0: + magnitude = (norm / image_peak) * 1.0 else: magnitude = norm * max_amp elif scaling_method == "Logarithmic": - log_mag = torch.pow(10.0, ((norm * top_db) - top_db)/20) + log_mag = torch.pow(10.0, ((norm * top_db) - top_db) / 20) magnitude = torch.clamp(log_mag, min=0.0) elif scaling_method == "Logarithmic Mel": - log_mag = torch.pow(10.0, ((norm * top_db) - top_db)/20) + log_mag = torch.pow(10.0, ((norm * top_db) - top_db) / 20) lin_mag = torch.clamp(log_mag, min=0.0) - mel_fb = torch_fn.melscale_fbanks( - n_freqs=n_stft, - f_min=0.0, - f_max=float(sampling_rate//2), - n_mels=norm.shape[1], - sample_rate=sampling_rate, - norm="slaney" if autogain==True else None, - mel_scale="htk" if mel_scale==False else "slaney" - ) + mel_fb = torch_fn.melscale_fbanks(n_freqs=n_stft, f_min=0.0, f_max=float(sampling_rate // 2), n_mels=norm.shape[1], sample_rate=sampling_rate, norm="slaney" if autogain == True else None, mel_scale="htk" if mel_scale == False else "slaney") mel_pinv = torch.linalg.pinv(mel_fb, rtol=1e-5, atol=1e-8) magnitude = torch.matmul(mel_pinv.T, lin_mag) @@ -369,7 +390,7 @@ class SpectrogramDecoder: griffin_lim = torch_trans.GriffinLim(n_fft=p2_fft, win_length=raw_frame_length, hop_length=hop_length, window_fn=windowing_func, n_iter=max_iter, momentum=gl_momentum, power=1.0) # ty: ignore[invalid-argument-type] waveform = griffin_lim(magnitude) max_val = torch.max(torch.abs(waveform)) - if max_val > 1e-8 : + if max_val > 1e-8: waveform = waveform / max_val waveform = waveform.unsqueeze(0).cpu() final_audio = {"waveform": waveform, "sample_rate": sampling_rate} diff --git a/src/steganography.py b/src/steganography.py index 044fa2b..7073eda 100644 --- a/src/steganography.py +++ b/src/steganography.py @@ -1,3 +1,5 @@ +# Steganography nodes, for hding text in images + import os import inspect import textwrap @@ -13,6 +15,7 @@ import folder_paths # This import is handled by ComfyUI. This is a very jank so # Every node is inspired heavily by their reference implementations in their GitHub repository, with changes made to best use PyTorch as possible as it's the main way ComfyUI stores data # For more information, check https://docs.comfy.org/custom-nodes/backend/images_and_masks#images + class Stegano_LSB_Encode: def __init__(self): pass @@ -413,9 +416,7 @@ class IMWatermarkDecode: encoding_format = self.encoding_selector(encoding_format, other_encoding_format) if types in ["bytes", "bits", "b16"]: # Bit-related types if length <= 0: - raise ValueError( - "Length is required for bytes-based watermark decoding. Please put in a non-zero length fitting to the format chosen." - ) + raise ValueError("Length is required for bytes-based watermark decoding. Please put in a non-zero length fitting to the format chosen.") if types == "bytes": length = int(length * 8) elif types == "b16": diff --git a/src/utils.py b/src/utils.py index 97c9c9b..85e6929 100644 --- a/src/utils.py +++ b/src/utils.py @@ -1,4 +1,5 @@ """Utilities for use with the ARG Toolkit""" + import os import base64 import ast @@ -6,10 +7,10 @@ import logging logger = logging.getLogger(__name__) -class SystemRandom: - def __init__(self): - pass +# The Random(pRNG) set of Utilities + +class SystemRandom: CATEGORY = "ARG Toolkit/Utilities/Random" DESCRIPTION = "Generates a random number of bytes using os.urandom. Required for use with the Cryptography nodes." @@ -32,6 +33,10 @@ class SystemRandom: RETURN_NAMES = ("rand_value",) FUNCTION = "SystemRandom" + @classmethod + def IS_CHANGED(self): + return float("NaN") # Due to the nature of os.urandom(), it will always change, to the point of unreproducible. Therefore, this only makes sense for it. Unfortunately, also due to the nature of it, this node will cause massively unoptimized workflows in large workflows as it forces reprocessing on everything. + def SystemRandom(self, byte_num): return (os.urandom(byte_num),) @@ -118,9 +123,7 @@ class String2Hex(ConverterNodes): class Hex2String(ConverterNodes): def execute(self, text, encoding_format, other_encoding_format): - return ( - bytes.fromhex(text.replace("0x", "").replace(" ", "")).decode(self.encoding_selector(encoding_format, other_encoding_format)), - ) + return (bytes.fromhex(text.replace("0x", "").replace(" ", "")).decode(self.encoding_selector(encoding_format, other_encoding_format)),) class String2Base64(ConverterNodes): @@ -232,14 +235,10 @@ class BitwiseNodes: raise ValueError(f"Unsupported datatype: {datatype}") def operate(self, text_1: str, text_2: str, datatype, encoding_format, func, other_encoding_format=None): - b1, return_fmt= self.detect_and_parse( - text=text_1, datatype=datatype, encoding_format=encoding_format, other_encoding_format=other_encoding_format - ) + b1, return_fmt = self.detect_and_parse(text=text_1, datatype=datatype, encoding_format=encoding_format, other_encoding_format=other_encoding_format) if text_2 != "": - b2, _ = self.detect_and_parse( - text=text_2, datatype=datatype, encoding_format=encoding_format, other_encoding_format=other_encoding_format - ) + b2, _ = self.detect_and_parse(text=text_2, datatype=datatype, encoding_format=encoding_format, other_encoding_format=other_encoding_format) # Ensure same length for operations that need it if len(b1) != len(b2): raise ValueError("Inputs must be same length for bitwise operations") @@ -378,7 +377,7 @@ class StringLooper: FUNCTION = "Looper" def Looper(self, text: str, loops: int): - return (text*loops,) + return (text * loops,) class ByteslikeEncode: @@ -387,12 +386,12 @@ class ByteslikeEncode: return { "required": { "text": ("STRING", {"multiline": True, "default": "Hello World!", "placeholder": "Type your message here..."}), - "encoding": (["Hexadecimal", "Base64", "UTF-8", "Binary", "Raw Bytes"], {"default": "UTF-8","tooltip": "The type of data to be encoded into bytes."}), + "encoding": (["Hexadecimal", "Base64", "UTF-8", "Binary", "Raw Bytes"], {"default": "UTF-8", "tooltip": "The type of data to be encoded into bytes."}), }, "optional": {}, "hidden": { "unique_id": "UNIQUE_ID", - } + }, } RETURN_TYPES = ("BYTESLIKE",) @@ -423,7 +422,7 @@ class ByteslikeEncode: case _: raise ValueError(f"Invalid encoding method used: {encoding}.") except (ValueError, TypeError, UnicodeDecodeError) as err: - logger.warning("Bytes-like Object Encode at node ID %s has failed to encode the string given under the mode of \"%s\" with %s. Falling back to raw string bytes.", unique_id, encoding, err) + logger.warning('Bytes-like Object Encode at node ID %s has failed to encode the string given under the mode of "%s" with %s. Falling back to raw string bytes.', unique_id, encoding, err) data = text.encode("utf-8") return (data,) @@ -432,14 +431,11 @@ class ByteslikeDecode: @classmethod def INPUT_TYPES(cls): return { - "required": { - "data": ("BYTESLIKE",{}), - "encoding": (["Hexadecimal", "Base64", "UTF-8", "Binary", "Raw Bytes"], {"default": "UTF-8", "tooltip": "The type of data the bytes will be decoded to."}) - }, + "required": {"data": ("BYTESLIKE", {}), "encoding": (["Hexadecimal", "Base64", "UTF-8", "Binary", "Raw Bytes"], {"default": "UTF-8", "tooltip": "The type of data the bytes will be decoded to."})}, "optional": {}, "hidden": { "unique_id": "UNIQUE_ID", - } + }, } RETURN_TYPES = ("STRING",) diff --git a/tests/conftest.py b/tests/conftest.py index 310609c..7ffe772 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,4 +3,4 @@ import sys # Add the project root directory to Python path # This allows the tests to import the project -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) diff --git a/tests/nodes/test_asymmetric_encryption.py b/tests/nodes/test_asymmetric_encryption.py index 44638e8..2b4d2e4 100644 --- a/tests/nodes/test_asymmetric_encryption.py +++ b/tests/nodes/test_asymmetric_encryption.py @@ -1,5 +1,4 @@ -"""Focused tests for the ``asymmetric_encryption`` node suite. -""" +"""Focused tests for the ``asymmetric_encryption`` node suite.""" import hashlib @@ -46,18 +45,10 @@ class TestEdDSAKeyFormat: loaded = serialization.load_pem_private_key(priv_bytes, b"hunter2") assert loaded is not None - def test_get_private_key_unsupported_type_raises(self): - with pytest.raises(ValueError, match="Unsupported key type"): - eddsa.InitNode.get_private_key("Fresh Key", "Bogus") - - def test_get_private_key_from_bytes_requires_bytes(self): - with pytest.raises(ValueError, match="Private key bytes are required"): - eddsa.InitNode.get_private_key("From Private Bytes", "Ed25519", None) - - def test_get_private_key_from_bytes_roundtrip(self): + def test_private_key_from_bytes_roundtrip(self): _, priv_key = eddsa.EdDSAPrivateKeyFormat().execute("Ed25519", "Raw", "Raw", "None", "") raw = priv_key.private_bytes_raw() - reloaded = eddsa.InitNode.get_private_key("From Private Bytes", "Ed25519", raw) + reloaded = eddsa.InitNode.get_private_format("Ed25519").from_private_bytes(raw) assert reloaded.private_bytes_raw() == raw @@ -65,62 +56,68 @@ class TestEdDSASignature: def _key(self): return eddsa.EdDSAPrivateKeyFormat().execute("Ed25519", "PEM", "PKCS8", "None", "")[1] - def test_sign_fresh_key(self): - signature = eddsa.EdDSASignature().execute("Fresh Key", "Ed25519", "Hello World", None, None)[0] - assert isinstance(signature, bytes) - assert len(signature) == 64 - def test_sign_from_private_bytes(self): _, priv_key = eddsa.EdDSAPrivateKeyFormat().execute("Ed25519", "Raw", "Raw", "None", "") - signature = eddsa.EdDSASignature().execute("From Private Bytes", "Ed25519", "Hello World", None, priv_key.private_bytes_raw())[0] + signature = eddsa.EdDSASignature().execute(False, "Ed25519", b"Hello World", None, priv_key.private_bytes_raw())[0] + assert isinstance(signature, bytes) assert len(signature) == 64 def test_sign_from_loaded_key(self): priv_key = self._key() - signature = eddsa.EdDSASignature().execute("From Loaded Key", "Ed25519", "Hello World", priv_key, None)[0] + signature = eddsa.EdDSASignature().execute(True, "Ed25519", b"Hello World", priv_key, None)[0] assert len(signature) == 64 def test_sign_and_verify_roundtrip(self): priv_key = self._key() - signature = eddsa.EdDSASignature().execute("From Loaded Key", "Ed25519", "Hello World", priv_key, None)[0] + signature = eddsa.EdDSASignature().execute(True, "Ed25519", b"Hello World", priv_key, None)[0] result = eddsa.EdDSAVerify().execute( "Ed25519", priv_key.public_key().public_bytes_raw(), None, b"Hello World", signature, - True, + False, ) assert result[0] is True def test_verify_rejects_tampered_message(self): priv_key = self._key() - signature = eddsa.EdDSASignature().execute("From Loaded Key", "Ed25519", "Hello World", priv_key, None)[0] + signature = eddsa.EdDSASignature().execute(True, "Ed25519", b"Hello World", priv_key, None)[0] result = eddsa.EdDSAVerify().execute( "Ed25519", priv_key.public_key().public_bytes_raw(), None, b"Hello World!", signature, - True, + False, ) assert result[0] is False def test_verify_from_loaded_key(self): priv_key = self._key() - signature = eddsa.EdDSASignature().execute("From Loaded Key", "Ed25519", "Hello World", priv_key, None)[0] - result = eddsa.EdDSAVerify().execute("Ed25519", None, priv_key.public_key(), b"Hello World", signature, False) + signature = eddsa.EdDSASignature().execute(True, "Ed25519", b"Hello World", priv_key, None)[0] + result = eddsa.EdDSAVerify().execute("Ed25519", None, priv_key.public_key(), b"Hello World", signature, True) assert result[0] is True + +class TestEdDSAPublicKeyFormat: def test_public_key_format_raw(self): - pub_bytes, pub_key = eddsa.EdDSAPublicKeyFormat().execute("Fresh Key", "Ed25519", "Raw", "Raw", None, None) + _, priv_key = eddsa.EdDSAPrivateKeyFormat().execute("Ed25519", "Raw", "Raw", "None", "") + pub_bytes, pub_key = eddsa.EdDSAPublicKeyFormat().execute(False, "Ed25519", "Raw", "Raw", None, priv_key.private_bytes_raw()) assert len(pub_bytes) == 32 assert pub_key.public_bytes_raw() == pub_bytes def test_public_key_format_pem(self): - pub_bytes, _ = eddsa.EdDSAPublicKeyFormat().execute("Fresh Key", "Ed25519", "PEM", "SubjectPublicKeyInfo", None, None) + _, priv_key = eddsa.EdDSAPrivateKeyFormat().execute("Ed25519", "Raw", "Raw", "None", "") + pub_bytes, _ = eddsa.EdDSAPublicKeyFormat().execute(False, "Ed25519", "PEM", "SubjectPublicKeyInfo", None, priv_key.private_bytes_raw()) assert pub_bytes.startswith(b"-----BEGIN PUBLIC KEY-----") + def test_public_key_format_from_loaded_key(self): + _, priv_key = eddsa.EdDSAPrivateKeyFormat().execute("Ed25519", "Raw", "Raw", "None", "") + pub_bytes, pub_key = eddsa.EdDSAPublicKeyFormat().execute(True, "Ed25519", "Raw", "Raw", priv_key, None) + assert len(pub_bytes) == 32 + assert pub_key.public_bytes_raw() == pub_bytes + # --------------------------------------------------------------------------- # Elliptic curve @@ -423,4 +420,3 @@ class TestRSANumbers: def test_restore_prime_factors(self): p, q = rsa.RSARestorePrime().execute(33, 17, 13) assert {p, q} == {3, 11} - diff --git a/tests/nodes/test_auth_encrypt.py b/tests/nodes/test_auth_encrypt.py index 45de5e5..185d14f 100644 --- a/tests/nodes/test_auth_encrypt.py +++ b/tests/nodes/test_auth_encrypt.py @@ -1,5 +1,6 @@ from src.cryptography_primitives import auth_encrypt + # Test suite for auth_encrypt.py class TestAuthEncrypt: def test_chacha20poly1305(self): diff --git a/tests/nodes/test_ciphers.py b/tests/nodes/test_ciphers.py index 56447ce..2eb632e 100644 --- a/tests/nodes/test_ciphers.py +++ b/tests/nodes/test_ciphers.py @@ -1,7 +1,26 @@ from src import ciphers + # Test suite for ciphers.py class TestCiphers: + def test_adfgx_cipher(self): + adfgx_cipher = ciphers.ADFGX() + text = "HelloWorld" + alphabet = "abcdefghiklmnopqrstuvwxyz" # 25 chars, no j + key = "KEY" + encrypted_text = adfgx_cipher.execute(text, alphabet, key, True, False) + decrypted_text = adfgx_cipher.execute(encrypted_text[0], alphabet, key, False, False) + assert decrypted_text[0] == text.lower() + + def test_adfgvx_cipher(self): + adfgx_cipher = ciphers.ADFGX() + text = "HelloWorld" + alphabet = "abcdefghiklmnopqrstuvwxyz" # 25 chars, no j + key = "KEY" + encrypted_text = adfgx_cipher.execute(text, alphabet, key, True, False) + decrypted_text = adfgx_cipher.execute(encrypted_text[0], alphabet, key, False, False) + assert decrypted_text[0] == text.lower() + def test_atbash_cipher(self): atbash_cipher = ciphers.Atbash() text = "Hello World" @@ -28,15 +47,6 @@ class TestCiphers: decrypted_text = vigenere_cipher.execute(encrypted_text[0], alphabet, key, False, False) assert decrypted_text[0] == text.lower().replace(" ", "") - def test_adfgx_cipher(self): - adfgx_cipher = ciphers.ADFGX() - text = "HelloWorld" - alphabet = "abcdefghiklmnopqrstuvwxyz" # 25 chars, no j - key = "KEY" - encrypted_text = adfgx_cipher.execute(text, alphabet, key, True, False) - decrypted_text = adfgx_cipher.execute(encrypted_text[0], alphabet, key, False, False) - assert decrypted_text[0] == text.lower() - def test_autokey_cipher(self): autokey_cipher = ciphers.Autokey() text = "Hello World" diff --git a/tests/nodes/test_error_correction.py b/tests/nodes/test_error_correction.py index aade14f..96838c1 100644 --- a/tests/nodes/test_error_correction.py +++ b/tests/nodes/test_error_correction.py @@ -168,4 +168,4 @@ class TestReedSolomonDecode: for ecc in [10, 14, 18, 20]: encoded = enc.execute(data, ecc) decoded = dec.execute(encoded[0], ecc) - assert decoded[0] == data \ No newline at end of file + assert decoded[0] == data diff --git a/tests/nodes/test_hashing.py b/tests/nodes/test_hashing.py index e1844c5..f5eb11a 100644 --- a/tests/nodes/test_hashing.py +++ b/tests/nodes/test_hashing.py @@ -1,5 +1,6 @@ from src.cryptography_primitives import hashing + # Test suite for hashing.py. Fixed to properly handle bytes class TestHashing: def test_sha256(self): diff --git a/tests/nodes/test_morse_code.py b/tests/nodes/test_morse_code.py index 4dbc1f9..ae38f04 100644 --- a/tests/nodes/test_morse_code.py +++ b/tests/nodes/test_morse_code.py @@ -1,6 +1,7 @@ import unicodedata from src import morse_code + # Test suite for morse_code.py class TestMorseCode: def setup_method(self): diff --git a/tests/nodes/test_rotor_machines.py b/tests/nodes/test_rotor_machines.py new file mode 100644 index 0000000..f824077 --- /dev/null +++ b/tests/nodes/test_rotor_machines.py @@ -0,0 +1,302 @@ +import pytest +from src.rotor_machines import RotorCipherRotor, RotorCipherPerm, RotorCipherMachine, RotorCipherPosition, RotorCipherEncrypt, RotorCipherPresetStepper, RotorCipherRuleStepper, RotorCipherMap, InitNode, RuleStepper, PresetStepper, Through, RotorCipherCircuitStepper + +# Setup machines and helpers for each major tests +A = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" +ROTORS = [("EKMFLGDQVZNTOWYHXUSPAIBRCJ", "Q"), # I + ("AJDKSIRUXBLHWTMCQGZNPYFVOE", "E"), # II + ("BDFHJLCPRTXVZNYEIWGAKMUSQO", "V")] # III +REFL_B = "YRUHQSLDPXNGOKMIEBFZCWVJAT" + +def machine(rings=(0, 0, 0), notches=True, stepper=None): + chain = None + for (w, n), r in zip(ROTORS, rings): + (chain,) = RotorCipherRotor().execute(A, w, n if notches else "", r, False, previous_rotor=chain) + (refl,) = RotorCipherPerm().execute(A, REFL_B) + return RotorCipherMachine().execute(chain, stepper or PresetStepper("Enigma", 3), reflector=refl)[0] + +def enc(m, start, text, preserve=False): + (p,) = RotorCipherPosition().execute(start) + out, final = RotorCipherEncrypt().execute(m, p, text, preserve) + return out, final.text + +def run(m, pos, text, decrypt=False): + out = [] + for c in text: + pos, y = m.press(pos, A.index(c), decrypt=decrypt) + out.append(A[y]) + return "".join(out) + +def aux_machine(base_rules="never\nnever\nnever\nnever", after=False, targets="E=2", reverse=""): + chain = None + for w, n in ROTORS: + (chain,) = RotorCipherRotor().execute(A, w, n, 0, False, previous_rotor=chain) + (aux,) = RotorCipherRotor().execute(A, ROTORS[0][0], "", 0, False) # a 1-rotor chain + (refl,) = RotorCipherPerm().execute(A, REFL_B) + (base,) = RotorCipherRuleStepper().execute(base_rules, reverse) + (st,) = RotorCipherCircuitStepper().execute(base, "3", "A", targets, after) + return RotorCipherMachine().execute(chain, st, reflector=refl, auxilliary_rotors=aux)[0] + +def one_key(m, start): + pos, _ = m.press(m.encode_pos(start), 0) + return m.decode_pos(pos) + +def one_rotor(wiring): + return RotorCipherRotor().execute(A, wiring, "", 0, False)[0][0] + +# --- published-style vector ------------------------------------------- +def test_enigma_i_known_vector(): + # Enigma I, I-II-III, UKW-B, rings AAA, start AAA. Widely cited; confirm + # against Crypto Museum's simulator if it ever fails. + assert enc(machine(), "AAA", "AAAAA")[0] == "BDZGO" + +# --- stepping ----------------------------------------------------------- +def test_double_step_sequence(): + m, pos, seen = machine(), (0, 3, 20), [] # ADU + for _ in range(3): + pos = m.stepper(m.rotors, pos) + seen.append("".join(A[p] for p in pos)) + assert seen == ["ADV", "AEW", "BFX"] + +def test_odometer_does_not_double_step(): + m = machine(stepper=PresetStepper("Odometer", 3)) + pos = m.stepper(m.rotors, (0, 3, 21)) # ADV: right at notch + assert "".join(A[p] for p in pos) == "AEW" + pos = m.stepper(m.rotors, pos) + assert "".join(A[p] for p in pos) == "AEX" # middle at notch, no carry + +def test_rule_string_matches_preset(): + m = machine() + rs = RuleStepper(RuleStepper.parse("1\n2|1\nalways")) + for pos in [(0, 3, 20), (0, 4, 22), (16, 4, 21), (0, 0, 0)]: + assert rs(m.rotors, pos) == m.stepper(m.rotors, pos) + +# --- cipher properties -------------------------------------------------- +def test_reciprocity(): + m = machine() + ct, _ = enc(m, "QEV", "THEQUICKBROWNFOXJUMPSOVERTHELAZYDOG") + assert enc(m, "QEV", ct)[0] == "THEQUICKBROWNFOXJUMPSOVERTHELAZYDOG" + +def test_no_self_encryption(): + # holds for Enigma-style reflector B (no fixed points) + m = machine() + text = "A" * 200 + assert "A" not in enc(m, "AAA", text)[0] + +def test_ring_offset_convention(): + # With no notches only the fast rotor steps, so ring r + position p must + # equal ring 0 + position (p - r). This is a consistency check of the + # pos - ring convention, NOT a published vector; add one from a real + # simulator with nonzero rings for external confirmation. + a = enc(machine(rings=(1, 2, 3), notches=False), "BDF", "HELLOWORLD") + b = enc(machine(rings=(0, 0, 0), notches=False), "ABC", "HELLOWORLD") + c = enc(machine(rings=(0, 0, 0), notches=False), "BDF", "HELLOWORLD") + print(a[0] == c[0]) + assert a[0] == b[0] + +def test_ring_single_rotor_external(): + ((r0,),) = RotorCipherRotor().execute(A, ROTORS[0][0], "Q", 0, False) + ((r1,),) = RotorCipherRotor().execute(A, ROTORS[0][0], "Q", 1, False) + assert A[Through(r0, 0).apply((0,), 0)] == "E" + assert A[Through(r1, 0).apply((0,), 0)] == "K" + +# --- modes -------------------------------------------------------------- +def test_skip_mode_drops_and_uppercases(): + m = machine() + out, _ = enc(m, "AAA", "Hello, World!") + assert out == enc(m, "AAA", "HELLOWORLD")[0] + +def test_preserve_mode_format(): + m = machine() + text = "Hello, World!" + out, _ = enc(m, "AAA", text, preserve=True) + assert len(out) == len(text) + for i, c in enumerate(text): + assert (out[i] == c) if not c.isalpha() else (out[i].isupper() == c.isupper()) + +def test_preserve_round_trip(): + m = machine() + text = "Hello, World! Attack at Dawn." + ct, _ = enc(m, "MCK", text, preserve=True) + assert enc(m, "MCK", ct, preserve=True)[0] == text + +def test_preserve_and_skip_agree_on_letters(): + m = machine() + p, _ = enc(m, "AAA", "Hi there", preserve=True) + s, _ = enc(m, "AAA", "Hi there") + assert "".join(c for c in p if c.isalpha()).upper() == s + +def test_chaining_continues_stream(): + m = machine() + whole, final = enc(m, "AAA", "ABCDEFGH") + first, mid = enc(m, "AAA", "ABCD") + second, end = enc(m, mid, "EFGH") + assert first + second == whole and end == final + +# --- validation --------------------------------------------------------- +def test_bad_wiring_rejected(): + with pytest.raises(ValueError): + RotorCipherRotor().execute(A, "AB" + A[2:-1], "Q", 0, False) # too short + with pytest.raises(ValueError): + RotorCipherRotor().execute(A, "A" * 26, "Q", 0, False) # repeats + +def test_ring_out_of_range(): + with pytest.raises(ValueError): + RotorCipherRotor().execute(A, ROTORS[0][0], "Q", 26, False) + +def test_non_involution_reflector_rejected(): + (chain,) = RotorCipherRotor().execute(A, ROTORS[0][0], "Q", 0, False) + (bad,) = RotorCipherPerm().execute(A, A[1:] + A[0]) # shift by one + with pytest.raises(ValueError): + RotorCipherMachine().execute(chain, RotorCipherPresetStepper().execute("Odometer", 1), reflector=bad) + +def test_alphabet_mismatch_across_rotors(): + (c1,) = RotorCipherRotor().execute(A, ROTORS[0][0], "Q", 0, False) + rev = A[::-1] + with pytest.raises(ValueError): + (c2,) = RotorCipherRotor().execute(rev, rev, "Q", 0, False, previous_rotor=c1) + RotorCipherMachine().execute(c2, RotorCipherPresetStepper().execute("Odometer", 2)) + +def test_positions_length_mismatch(): + m = machine() + with pytest.raises(ValueError): + enc(m, "AA", "HELLO") + +def test_positions_outside_rotor_alphabet(): + m = machine() + (p,) = RotorCipherPosition().execute("123") + with pytest.raises(ValueError): + RotorCipherEncrypt().execute(m, p, "HELLO", False) + +def test_preset_k_too_large(): + (chain,) = RotorCipherRotor().execute(A, ROTORS[0][0], "Q", 0, False) + with pytest.raises(ValueError): + RotorCipherMachine().execute(chain, PresetStepper("Enigma", 3)) + +# --- decryption -------------------------------------------------- + +def no_reflector_machine(): + chain = None + for w, n in ROTORS: + (chain,) = RotorCipherRotor().execute(A, w, n, 0, False, previous_rotor=chain) + (plug,) = RotorCipherPerm().execute(A, "BACDEFGHIJKLMNOPQRSTUVWXYZ") # swaps A,B + (entry,) = RotorCipherPerm().execute(A, A[1:] + A[0]) # not an involution, on purpose + return RotorCipherMachine().execute(chain, PresetStepper("Enigma", 3), + entry=entry, plugboard=plug)[0] + +def test_non_reflecting_round_trip(): + m, start = no_reflector_machine(), (16, 4, 21) # start near notches + text = "THEQUICKBROWNFOXJUMPSOVERTHELAZYDOG" * 3 + ct = run(m, start, text) + assert ct != text + assert run(m, start, ct, decrypt=True) == text + +def test_non_reflecting_is_not_self_inverse(): + m, start = no_reflector_machine(), (0, 0, 0) + ct = run(m, start, "HELLOWORLD") + assert run(m, start, ct) != "HELLOWORLD" # decrypting needs decrypt=True + +def test_reflecting_decrypt_equals_encrypt(): + m, start = machine(), (0, 3, 20) + text = "ATTACKATDAWN" + assert run(m, start, text) == run(m, start, text, decrypt=True) + +def test_positioned_reflector_matches_static_at_A(): + chain = None + for w, n in ROTORS: + (chain,) = RotorCipherRotor().execute(A, w, n, 0, False, previous_rotor=chain) + (rr,) = RotorCipherRotor().execute(A, REFL_B, "", 0, False) + m = RotorCipherMachine().execute(chain, PresetStepper("Enigma", 3), reflector_rotor=rr)[0] + assert run(m, (0, 0, 0, 0), "HELLOWORLD") == run(machine(), (0, 0, 0), "HELLOWORLD") + +def test_positioned_reflector_reciprocal_and_position_dependent(): + chain = None + for w, n in ROTORS: + (chain,) = RotorCipherRotor().execute(A, w, n, 0, False, previous_rotor=chain) + (rr,) = RotorCipherRotor().execute(A, REFL_B, "", 0, False) + m = RotorCipherMachine().execute(chain, PresetStepper("Enigma", 3), reflector_rotor=rr)[0] + a = run(m, (5, 0, 0, 0), "ATTACKATDAWN") + assert run(m, (5, 0, 0, 0), a) == "ATTACKATDAWN" + assert a != run(m, (0, 0, 0, 0), "ATTACKATDAWN") + +def test_mapping_stage_changes_alphabet(): + D = "0123456789" + (chain,) = RotorCipherRotor().execute(D, "5290817346", "", 0, False) + (ent,) = RotorCipherMap().execute(A, D, "0123456789" * 2 + "012345") + m = RotorCipherMachine().execute(chain, PresetStepper("Odometer", 1), entry=ent)[0] + assert (m.in_alphabet, m.out_alphabet) == (A, D) + pos, y = m.press((0,), A.index("C")) + assert 0 <= y < 10 + with pytest.raises(ValueError): + m.press((0,), 0, decrypt=True) + +# --- auxiliary rotors --------------------------------------------------------------------------------------- + +def test_circuit_depends_on_aux_position(): + m = aux_machine() + assert one_key(m, "AAAA") == "AABA" # aux A -> E -> target: slot 2 steps + assert one_key(m, "AAAB") == "AAAB" # aux B -> J -> no step + +def test_evaluate_after_base(): + rules = "never\nnever\nnever\nalways" # aux rotor steps every key + assert one_key(aux_machine(rules, after=False), "AAAA") == "AABB" # evaluated at A + assert one_key(aux_machine(rules, after=True), "AAAA") == "AAAB" # evaluated at B + +def test_aux_reflecting_round_trip(): + m = aux_machine("never\nnever\nnever\nalways") + start = m.encode_pos("AAAA") + ct = run(m, start, "THEQUICKBROWNFOXJUMPSOVERTHELAZYDOG") + assert run(m, start, ct) == "THEQUICKBROWNFOXJUMPSOVERTHELAZYDOG" + +def test_driven_slot_must_be_never(): + with pytest.raises(ValueError): + aux_machine("never\nnever\nalways\nnever") # slot 2 also driven by the circuit + +def test_circuit_slot_out_of_range(): + with pytest.raises(ValueError): + aux_machine(targets="E=9") + +def test_preset_stepper_rejected_with_aux(): + chain = None + for w, n in ROTORS: + (chain,) = RotorCipherRotor().execute(A, w, n, 0, False, previous_rotor=chain) + (aux,) = RotorCipherRotor().execute(A, ROTORS[0][0], "", 0, False) + with pytest.raises(ValueError): + RotorCipherMachine().execute(chain, PresetStepper("Enigma", 3), auxilliary_rotors=aux) + +# --- reversal --------------------------------------------------------------------------------- + +def test_parse_reverse(): + assert RuleStepper.parse_reverse("", 4) is None + assert RuleStepper.parse_reverse("0 2", 4) == (-1, 1, -1, 1) + assert RuleStepper.parse_reverse("1,3", 4) == (1, -1, 1, -1) + with pytest.raises(ValueError): + RuleStepper.parse_reverse("9", 4) + with pytest.raises(ValueError): + RuleStepper.parse_reverse("x", 4) + +def test_reverse_rule_stepper(): + rs = (one_rotor(ROTORS[0][0]), one_rotor(ROTORS[0][0])) + (st,) = RotorCipherRuleStepper().execute("always\nalways", "1") + assert st(rs, (1, 1)) == (2, 0) # slot 0 forward, slot 1 backward + assert st(rs, (0, 0)) == (1, 25) # backward wraps A -> Z + +def test_reverse_circuit_driven_slot(): + m = aux_machine(targets="E=2", reverse="2") + assert one_key(m, "AAAA") == "AAZA" + +def test_no_reverse_unchanged(): + # existing behavior: a blank reverse list gives steps=None + (st,) = RotorCipherRuleStepper().execute("1\n2|1\nalways", "") + assert st.steps is None + +# --- description --------------------------------------------------------- + +def test_describe_enigma(): + d = machine().describe() + assert "slot 0 first" in d and "2|1" in d and "always" in d + assert "rotor slot 2 forward" in d and "involution" in d + +def test_describe_aux_roles_and_circuit(): + d = aux_machine().describe() + assert "aux" in d and "Circuit:" in d and "circuit(E)" in d diff --git a/tests/nodes/test_spectrogram.py b/tests/nodes/test_spectrogram.py index 4f25596..0851c49 100644 --- a/tests/nodes/test_spectrogram.py +++ b/tests/nodes/test_spectrogram.py @@ -2,6 +2,7 @@ from src.spectrogram import SpectrogramEncoder, SpectrogramDecoder import torch import numpy as np + # Test suite for spectrogram.py class TestSpectrogramEncoder: def setup_method(self): @@ -52,10 +53,10 @@ class TestSpectrogramEncoder: # Test with default colormap result = self.encoder.color_mapper( - matrix, + matrix, "#000000", # low "#B73779", # mid - "#FCFDBF" # high + "#FCFDBF", # high ) # Check shape and type @@ -69,12 +70,14 @@ class TestSpectrogramEncoder: # Test crop_spectrogram_dimensions helper method def test_crop_spectrogram_dimensions_basic(self): # Create a test spectrogram with some active and inactive regions - db_spectrogram = np.array([ - [0.0, 0.0, 0.0, 0.0, 0.0], - [0.0, -20.0, -15.0, -10.0, 0.0], - [0.0, -30.0, -25.0, -20.0, 0.0], - [0.0, 0.0, 0.0, 0.0, 0.0], - ]) + db_spectrogram = np.array( + [ + [0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, -20.0, -15.0, -10.0, 0.0], + [0.0, -30.0, -25.0, -20.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0], + ] + ) # Test with active_range_db = 20 (should cut off values below -40) result = self.encoder.crop_spectrogram_dimensions(db_spectrogram, 20.0) @@ -86,10 +89,12 @@ class TestSpectrogramEncoder: def test_crop_spectrogram_dimensions_empty(self): # Test with all values below threshold - db_spectrogram = np.array([ - [-50.0, -60.0, -70.0], - [-55.0, -65.0, -75.0], - ]) + db_spectrogram = np.array( + [ + [-50.0, -60.0, -70.0], + [-55.0, -65.0, -75.0], + ] + ) result = self.encoder.crop_spectrogram_dimensions(db_spectrogram, 20.0) # Should return original when no active region found @@ -100,25 +105,11 @@ class TestSpectrogramEncoder: # Create a minimal mock audio input audio = { "waveform": torch.randn(1, 1, 22050), # 1 channel, 1 sec of audio at 22050 Hz - "sample_rate": 22050 + "sample_rate": 22050, } # Test with logarithmic scaling - result = self.encoder.execute( - audio=audio, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - db_norm=80.0, - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF", - cutoff_switch=False, - scaling_method="Logarithmic", - active_db=60.0, - mel_scale=False, - slaney_norm=False - ) + result = self.encoder.execute(audio=audio, window_size=25.0, hop_size=15.0, windowing_func="Hanning", db_norm=80.0, colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF", cutoff_switch=False, scaling_method="Logarithmic", active_db=60.0, mel_scale=False, slaney_norm=False) # Should return a tensor (spectrogram) spectrogram_tensor = result[0] @@ -127,27 +118,10 @@ class TestSpectrogramEncoder: def test_execute_linear_scaling(self): # Create a minimal mock audio input - audio = { - "waveform": torch.randn(1, 1, 22050), - "sample_rate": 22050 - } + audio = {"waveform": torch.randn(1, 1, 22050), "sample_rate": 22050} # Test with linear scaling - result = self.encoder.execute( - audio=audio, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - db_norm=80.0, - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF", - cutoff_switch=False, - scaling_method="Linear", - active_db=60.0, - mel_scale=False, - slaney_norm=False - ) + result = self.encoder.execute(audio=audio, window_size=25.0, hop_size=15.0, windowing_func="Hanning", db_norm=80.0, colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF", cutoff_switch=False, scaling_method="Linear", active_db=60.0, mel_scale=False, slaney_norm=False) # Should return a tensor spectrogram_tensor = result[0] @@ -155,27 +129,10 @@ class TestSpectrogramEncoder: def test_execute_quadratic_scaling(self): # Create a minimal mock audio input - audio = { - "waveform": torch.randn(1, 1, 22050), - "sample_rate": 22050 - } + audio = {"waveform": torch.randn(1, 1, 22050), "sample_rate": 22050} # Test with quadratic scaling - result = self.encoder.execute( - audio=audio, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - db_norm=80.0, - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF", - cutoff_switch=False, - scaling_method="Quadratic", - active_db=60.0, - mel_scale=False, - slaney_norm=False - ) + result = self.encoder.execute(audio=audio, window_size=25.0, hop_size=15.0, windowing_func="Hanning", db_norm=80.0, colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF", cutoff_switch=False, scaling_method="Quadratic", active_db=60.0, mel_scale=False, slaney_norm=False) # Should return a tensor spectrogram_tensor = result[0] @@ -183,27 +140,10 @@ class TestSpectrogramEncoder: def test_execute_logarithmic_mel_scaling(self): # Create a minimal mock audio input - audio = { - "waveform": torch.randn(1, 1, 22050), - "sample_rate": 22050 - } + audio = {"waveform": torch.randn(1, 1, 22050), "sample_rate": 22050} # Test with logarithmic mel scaling - result = self.encoder.execute( - audio=audio, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - db_norm=80.0, - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF", - cutoff_switch=False, - scaling_method="Logarithmic Mel", - active_db=60.0, - mel_scale=False, - slaney_norm=False - ) + result = self.encoder.execute(audio=audio, window_size=25.0, hop_size=15.0, windowing_func="Hanning", db_norm=80.0, colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF", cutoff_switch=False, scaling_method="Logarithmic Mel", active_db=60.0, mel_scale=False, slaney_norm=False) # Should return a tensor spectrogram_tensor = result[0] @@ -211,27 +151,10 @@ class TestSpectrogramEncoder: def test_execute_with_cutoff(self): # Create a minimal mock audio input - audio = { - "waveform": torch.randn(1, 1, 22050), - "sample_rate": 22050 - } + audio = {"waveform": torch.randn(1, 1, 22050), "sample_rate": 22050} # Test with cutoff enabled - result = self.encoder.execute( - audio=audio, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - db_norm=80.0, - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF", - cutoff_switch=True, - scaling_method="Logarithmic", - active_db=60.0, - mel_scale=False, - slaney_norm=False - ) + result = self.encoder.execute(audio=audio, window_size=25.0, hop_size=15.0, windowing_func="Hanning", db_norm=80.0, colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF", cutoff_switch=True, scaling_method="Logarithmic", active_db=60.0, mel_scale=False, slaney_norm=False) # Should return a tensor spectrogram_tensor = result[0] @@ -239,28 +162,11 @@ class TestSpectrogramEncoder: def test_execute_different_windowing_functions(self): # Create a minimal mock audio input - audio = { - "waveform": torch.randn(1, 1, 22050), - "sample_rate": 22050 - } + audio = {"waveform": torch.randn(1, 1, 22050), "sample_rate": 22050} # Test with different windowing functions for window_func in ["Hanning", "Bartlett", "Blackman", "Hamming"]: - result = self.encoder.execute( - audio=audio, - window_size=25.0, - hop_size=15.0, - windowing_func=window_func, - db_norm=80.0, - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF", - cutoff_switch=False, - scaling_method="Logarithmic", - active_db=60.0, - mel_scale=False, - slaney_norm=False - ) + result = self.encoder.execute(audio=audio, window_size=25.0, hop_size=15.0, windowing_func=window_func, db_norm=80.0, colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF", cutoff_switch=False, scaling_method="Logarithmic", active_db=60.0, mel_scale=False, slaney_norm=False) # Should return a tensor spectrogram_tensor = result[0] @@ -268,27 +174,10 @@ class TestSpectrogramEncoder: def test_execute_with_mel_scale(self): # Create a minimal mock audio input - audio = { - "waveform": torch.randn(1, 1, 22050), - "sample_rate": 22050 - } + audio = {"waveform": torch.randn(1, 1, 22050), "sample_rate": 22050} # Test with mel_scale enabled - result = self.encoder.execute( - audio=audio, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - db_norm=80.0, - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF", - cutoff_switch=False, - scaling_method="Logarithmic Mel", - active_db=60.0, - mel_scale=True, - slaney_norm=False - ) + result = self.encoder.execute(audio=audio, window_size=25.0, hop_size=15.0, windowing_func="Hanning", db_norm=80.0, colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF", cutoff_switch=False, scaling_method="Logarithmic Mel", active_db=60.0, mel_scale=True, slaney_norm=False) # Should return a tensor spectrogram_tensor = result[0] @@ -296,31 +185,16 @@ class TestSpectrogramEncoder: def test_execute_with_slaney_norm(self): # Create a minimal mock audio input - audio = { - "waveform": torch.randn(1, 1, 22050), - "sample_rate": 22050 - } + audio = {"waveform": torch.randn(1, 1, 22050), "sample_rate": 22050} # Test with slaney_norm enabled - result = self.encoder.execute( - audio=audio, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - db_norm=80.0, - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF", - cutoff_switch=False, - scaling_method="Logarithmic Mel", - active_db=60.0, - mel_scale=True, - slaney_norm=True - ) + result = self.encoder.execute(audio=audio, window_size=25.0, hop_size=15.0, windowing_func="Hanning", db_norm=80.0, colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF", cutoff_switch=False, scaling_method="Logarithmic Mel", active_db=60.0, mel_scale=True, slaney_norm=True) # Should return a tensor spectrogram_tensor = result[0] assert isinstance(spectrogram_tensor, torch.Tensor) + + class TestSpectrogramDecoder: def setup_method(self): self.decoder = SpectrogramDecoder() @@ -334,7 +208,7 @@ class TestSpectrogramDecoder: image, "#000000", # low "#B73779", # mid - "#FCFDBF" # high + "#FCFDBF", # high ) # Should return a normalized value between 0 and 1 @@ -350,7 +224,7 @@ class TestSpectrogramDecoder: image, "#FF0000", # red low "#00FF00", # green mid - "#0000FF" # blue high + "#0000FF", # blue high ) assert result.shape == (5, 5) @@ -363,22 +237,7 @@ class TestSpectrogramDecoder: image = torch.rand(513, 513, 3) # Height, Width, RGB result = self.decoder.execute( - image=image, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - autogain=False, - max_amp=1.0, - max_iter=32, - top_db=80.0, - sampling_rate=22050, - gl_momentum=0.99, - scaling_method="Linear", - mel_scale=False, - channel="luminance", - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF" + image=image, window_size=25.0, hop_size=15.0, windowing_func="Hanning", autogain=False, max_amp=1.0, max_iter=32, top_db=80.0, sampling_rate=22050, gl_momentum=0.99, scaling_method="Linear", mel_scale=False, channel="luminance", colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF" ) # Should return audio @@ -393,22 +252,7 @@ class TestSpectrogramDecoder: image = torch.rand(513, 513, 3) result = self.decoder.execute( - image=image, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - autogain=False, - max_amp=1.0, - max_iter=32, - top_db=80.0, - sampling_rate=22050, - gl_momentum=0.99, - scaling_method="Linear", - mel_scale=False, - channel="channel_red", - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF" + image=image, window_size=25.0, hop_size=15.0, windowing_func="Hanning", autogain=False, max_amp=1.0, max_iter=32, top_db=80.0, sampling_rate=22050, gl_momentum=0.99, scaling_method="Linear", mel_scale=False, channel="channel_red", colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF" ) # Should return audio @@ -421,22 +265,7 @@ class TestSpectrogramDecoder: image = torch.rand(513, 513, 3) result = self.decoder.execute( - image=image, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - autogain=False, - max_amp=1.0, - max_iter=32, - top_db=80.0, - sampling_rate=22050, - gl_momentum=0.99, - scaling_method="Linear", - mel_scale=False, - channel="channel_green", - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF" + image=image, window_size=25.0, hop_size=15.0, windowing_func="Hanning", autogain=False, max_amp=1.0, max_iter=32, top_db=80.0, sampling_rate=22050, gl_momentum=0.99, scaling_method="Linear", mel_scale=False, channel="channel_green", colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF" ) # Should return audio @@ -449,22 +278,7 @@ class TestSpectrogramDecoder: image = torch.rand(513, 513, 3) result = self.decoder.execute( - image=image, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - autogain=False, - max_amp=1.0, - max_iter=32, - top_db=80.0, - sampling_rate=22050, - gl_momentum=0.99, - scaling_method="Linear", - mel_scale=False, - channel="channel_blue", - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF" + image=image, window_size=25.0, hop_size=15.0, windowing_func="Hanning", autogain=False, max_amp=1.0, max_iter=32, top_db=80.0, sampling_rate=22050, gl_momentum=0.99, scaling_method="Linear", mel_scale=False, channel="channel_blue", colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF" ) # Should return audio @@ -477,22 +291,7 @@ class TestSpectrogramDecoder: image = torch.rand(513, 513, 3) result = self.decoder.execute( - image=image, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - autogain=False, - max_amp=1.0, - max_iter=32, - top_db=80.0, - sampling_rate=22050, - gl_momentum=0.99, - scaling_method="Linear", - mel_scale=False, - channel="colormap", - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF" + image=image, window_size=25.0, hop_size=15.0, windowing_func="Hanning", autogain=False, max_amp=1.0, max_iter=32, top_db=80.0, sampling_rate=22050, gl_momentum=0.99, scaling_method="Linear", mel_scale=False, channel="colormap", colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF" ) # Should return audio @@ -505,22 +304,7 @@ class TestSpectrogramDecoder: image = torch.rand(513, 513, 3) result = self.decoder.execute( - image=image, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - autogain=True, - max_amp=1.0, - max_iter=32, - top_db=80.0, - sampling_rate=22050, - gl_momentum=0.99, - scaling_method="Linear", - mel_scale=False, - channel="luminance", - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF" + image=image, window_size=25.0, hop_size=15.0, windowing_func="Hanning", autogain=True, max_amp=1.0, max_iter=32, top_db=80.0, sampling_rate=22050, gl_momentum=0.99, scaling_method="Linear", mel_scale=False, channel="luminance", colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF" ) # Should return audio @@ -533,22 +317,7 @@ class TestSpectrogramDecoder: image = torch.rand(513, 513, 3) result = self.decoder.execute( - image=image, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - autogain=False, - max_amp=1.0, - max_iter=32, - top_db=80.0, - sampling_rate=22050, - gl_momentum=0.99, - scaling_method="Logarithmic", - mel_scale=False, - channel="luminance", - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF" + image=image, window_size=25.0, hop_size=15.0, windowing_func="Hanning", autogain=False, max_amp=1.0, max_iter=32, top_db=80.0, sampling_rate=22050, gl_momentum=0.99, scaling_method="Logarithmic", mel_scale=False, channel="luminance", colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF" ) # Should return audio @@ -561,22 +330,7 @@ class TestSpectrogramDecoder: image = torch.rand(513, 513, 3) result = self.decoder.execute( - image=image, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - autogain=False, - max_amp=1.0, - max_iter=32, - top_db=80.0, - sampling_rate=22050, - gl_momentum=0.99, - scaling_method="Quadratic", - mel_scale=False, - channel="luminance", - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF" + image=image, window_size=25.0, hop_size=15.0, windowing_func="Hanning", autogain=False, max_amp=1.0, max_iter=32, top_db=80.0, sampling_rate=22050, gl_momentum=0.99, scaling_method="Quadratic", mel_scale=False, channel="luminance", colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF" ) # Should return audio @@ -589,22 +343,7 @@ class TestSpectrogramDecoder: image = torch.rand(513, 513, 3) result = self.decoder.execute( - image=image, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - autogain=True, - max_amp=1.0, - max_iter=32, - top_db=80.0, - sampling_rate=22050, - gl_momentum=0.99, - scaling_method="Logarithmic Mel", - mel_scale=False, - channel="luminance", - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF" + image=image, window_size=25.0, hop_size=15.0, windowing_func="Hanning", autogain=True, max_amp=1.0, max_iter=32, top_db=80.0, sampling_rate=22050, gl_momentum=0.99, scaling_method="Logarithmic Mel", mel_scale=False, channel="luminance", colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF" ) # Should return audio @@ -617,22 +356,7 @@ class TestSpectrogramDecoder: image = torch.rand(513, 513, 3) result = self.decoder.execute( - image=image, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - autogain=False, - max_amp=1.0, - max_iter=32, - top_db=80.0, - sampling_rate=22050, - gl_momentum=0.99, - scaling_method="Logarithmic Mel", - mel_scale=True, - channel="luminance", - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF" + image=image, window_size=25.0, hop_size=15.0, windowing_func="Hanning", autogain=False, max_amp=1.0, max_iter=32, top_db=80.0, sampling_rate=22050, gl_momentum=0.99, scaling_method="Logarithmic Mel", mel_scale=True, channel="luminance", colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF" ) # Should return audio @@ -645,22 +369,7 @@ class TestSpectrogramDecoder: image = torch.rand(513, 513, 3) result = self.decoder.execute( - image=image, - window_size=25.0, - hop_size=15.0, - windowing_func="Hanning", - autogain=True, - max_amp=1.0, - max_iter=32, - top_db=80.0, - sampling_rate=22050, - gl_momentum=0.99, - scaling_method="Logarithmic Mel", - mel_scale=True, - channel="luminance", - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF" + image=image, window_size=25.0, hop_size=15.0, windowing_func="Hanning", autogain=True, max_amp=1.0, max_iter=32, top_db=80.0, sampling_rate=22050, gl_momentum=0.99, scaling_method="Logarithmic Mel", mel_scale=True, channel="luminance", colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF" ) # Should return audio @@ -675,22 +384,7 @@ class TestSpectrogramDecoder: # Test with different windowing functions for window_func in ["Hanning", "Bartlett", "Blackman", "Hamming"]: result = self.decoder.execute( - image=image, - window_size=25.0, - hop_size=15.0, - windowing_func=window_func, - autogain=False, - max_amp=1.0, - max_iter=32, - top_db=80.0, - sampling_rate=22050, - gl_momentum=0.99, - scaling_method="Linear", - mel_scale=False, - channel="luminance", - colormap_low="#000000", - colormap_mid="#B73779", - colormap_high="#FCFDBF" + image=image, window_size=25.0, hop_size=15.0, windowing_func=window_func, autogain=False, max_amp=1.0, max_iter=32, top_db=80.0, sampling_rate=22050, gl_momentum=0.99, scaling_method="Linear", mel_scale=False, channel="luminance", colormap_low="#000000", colormap_mid="#B73779", colormap_high="#FCFDBF" ) # Should return audio @@ -720,7 +414,7 @@ class TestSpectrogramDecoder: channel="luminance", colormap_low="#000000", colormap_mid="#B73779", - colormap_high="#FCFDBF" + colormap_high="#FCFDBF", ) # Should return audio @@ -750,7 +444,7 @@ class TestSpectrogramDecoder: channel="luminance", colormap_low="#000000", colormap_mid="#B73779", - colormap_high="#FCFDBF" + colormap_high="#FCFDBF", ) # Should return audio diff --git a/tests/nodes/test_utils.py b/tests/nodes/test_utils.py index 94a49dd..7baa1a5 100644 --- a/tests/nodes/test_utils.py +++ b/tests/nodes/test_utils.py @@ -2,6 +2,7 @@ from src import utils # Test suite for utils.py + class TestUtils: # Converters