nodes/AzelusLightvale-ComfyUI-ARG-Toolkit
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
NexusLightvale
2026-10-24 17:38:49 +07:00
parent fa9bc7b30e
commit a26db01ac3
34 changed files with 1720 additions and 655 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+10 -6
View File
@@ -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.
__init__.py
+2 -8
View File
@@ -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
src/ciphers.py
+3 -3
View File
@@ -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)
src/cryptography_primitives/asymmetric_encryption/eddsa.py
+59 -71
View File
@@ -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
src/cryptography_primitives/asymmetric_encryption/elliptic_curve.py
+1
View File
@@ -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
src/cryptography_primitives/asymmetric_encryption/mldsa.py
+241
View File
@@ -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",
}
src/cryptography_primitives/asymmetric_encryption/rsa.py
+151 -71
View File
@@ -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,
src/cryptography_primitives/asymmetric_encryption/serialization.py
+23 -4
View File
@@ -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):
src/cryptography_primitives/asymmetric_encryption/xdh.py
+1
View File
@@ -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
src/cryptography_primitives/auth_encrypt.py
+3 -3
View File
@@ -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,)
src/cryptography_primitives/const_time.py
+2
View File
@@ -1,3 +1,5 @@
# Constant time compare, the cryptography way
from cryptography.hazmat.primitives import constant_time
src/cryptography_primitives/hashing.py
+2
View File
@@ -1,3 +1,5 @@
# Message digests/hashes
from cryptography.hazmat.primitives import hashes
src/cryptography_primitives/kdf.py
+13 -33
View File
@@ -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,)
src/cryptography_primitives/key_wrapper.py
+3
View File
@@ -1,5 +1,8 @@
# Key wrapping
from cryptography.hazmat.primitives import keywrap
class WrapKeyNodes:
def __init__(self):
pass
src/cryptography_primitives/symm_padding.py
+3
View File
@@ -1,5 +1,8 @@
# Symmetric padding for block ciphers
from cryptography.hazmat.primitives import padding
class PaddingNode:
CATEGORY = "ARG Toolkit/Cryptography/Modern/Padding"
src/cryptography_primitives/symmetrical_encrypt.py
+4 -2
View File
@@ -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
src/debugging_nodes.py
+3 -3
View File
@@ -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 = {
src/error_correction.py
+2 -1
View File
@@ -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",)
src/fernet.py
+2 -1
View File
@@ -1,4 +1,5 @@
# Fernet encryption, the basics of basics
# Fernet encryption, the basics of basics in modern cryptography
from cryptography import fernet
src/morse_code.py
+1 -1
View File
@@ -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
src/rotor_machines.py
+719
View File
@@ -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",
}
src/spectrogram.py
+49 -28
View File
@@ -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}
src/steganography.py
+4 -3
View File
@@ -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":
src/utils.py
+17 -21
View File
@@ -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",)
tests/conftest.py
+1 -1
View File
@@ -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__), "..")))
tests/nodes/test_asymmetric_encryption.py
+24 -28
View File
@@ -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}
tests/nodes/test_auth_encrypt.py
+1
View File
@@ -1,5 +1,6 @@
from src.cryptography_primitives import auth_encrypt
# Test suite for auth_encrypt.py
class TestAuthEncrypt:
def test_chacha20poly1305(self):
tests/nodes/test_ciphers.py
+19 -9
View File
@@ -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"
tests/nodes/test_error_correction.py
+1 -1
View File
@@ -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
assert decoded[0] == data
tests/nodes/test_hashing.py
+1
View File
@@ -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):
tests/nodes/test_morse_code.py
+1
View File
@@ -1,6 +1,7 @@
import unicodedata
from src import morse_code
# Test suite for morse_code.py
class TestMorseCode:
def setup_method(self):
tests/nodes/test_rotor_machines.py
+302
View File
@@ -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
tests/nodes/test_spectrogram.py
+51 -357
View File
@@ -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
tests/nodes/test_utils.py
+1
View File
@@ -2,6 +2,7 @@ from src import utils
# Test suite for utils.py
class TestUtils:
# Converters