diff --git a/apps/cli/blob.py b/apps/cli/blob.py index d2bf390..6041050 100644 --- a/apps/cli/blob.py +++ b/apps/cli/blob.py @@ -1,3 +1,4 @@ +import re from hashlib import sha256 from binascii import hexlify, unhexlify from cryptography.hazmat.primitives.ciphers.aead import AESGCM @@ -8,6 +9,9 @@ from apps.cli import SUPPORTED_HASH_FUNCTIONS, SUPPORTED_CIPHERS class Blob: def __init__(self, data, cipher, hash_function): + if type(data) is not str or re.search(r'^[0-9A-Fa-f]+$', data) is None: + raise ValueError("Non-Hex character found in txid.") + self.data = data self.cipher = cipher self.hash_function = hash_function @@ -23,6 +27,12 @@ class Blob: SUPPORTED_CIPHERS)) def encrypt(self, tx_id): + if len(tx_id) != 64: + raise ValueError("txid does not matches the expected size (32-byte / 64 hex chars).") + + elif re.search(r'^[0-9A-Fa-f]+$', tx_id) is None: + raise ValueError("Non-Hex character found in txid.") + # Transaction to be encrypted # FIXME: The blob data should contain more things that just the transaction. Leaving like this for now. tx = unhexlify(self.data) diff --git a/test/unit/test_blob.py b/test/unit/test_blob.py new file mode 100644 index 0000000..fc95450 --- /dev/null +++ b/test/unit/test_blob.py @@ -0,0 +1,94 @@ +from os import urandom + +from pisa import logging +from apps.cli.blob import Blob +from pisa.conf import SUPPORTED_CIPHERS, SUPPORTED_HASH_FUNCTIONS + + +def test_init_blob(): + data = urandom(64).hex() + + # Fixed (valid) hash function, try different valid ciphers + hash_function = SUPPORTED_HASH_FUNCTIONS[0] + for cipher in SUPPORTED_CIPHERS: + cipher_cases = [cipher, cipher.lower(), cipher.capitalize()] + + for case in cipher_cases: + blob = Blob(data, case, hash_function) + assert(blob.data == data and blob.cipher == case and blob.hash_function == hash_function) + + # Fixed (valid) cipher, try different valid hash functions + cipher = SUPPORTED_CIPHERS[0] + for hash_function in SUPPORTED_HASH_FUNCTIONS: + hash_function_cases = [hash_function, hash_function.lower(), hash_function.capitalize()] + + for case in hash_function_cases: + blob = Blob(data, cipher, case) + assert(blob.data == data and blob.cipher == cipher and blob.hash_function == case) + + # Invalid data + data = urandom(64) + cipher = SUPPORTED_CIPHERS[0] + hash_function = SUPPORTED_HASH_FUNCTIONS[0] + + try: + Blob(data, cipher, hash_function) + assert False, "Able to create blob with wrong data" + + except ValueError: + assert True + + # Invalid cipher + data = urandom(64).hex() + cipher = "A" * 10 + hash_function = SUPPORTED_HASH_FUNCTIONS[0] + + try: + Blob(data, cipher, hash_function) + assert False, "Able to create blob with wrong data" + + except ValueError: + assert True + + # Invalid hash function + data = urandom(64).hex() + cipher = SUPPORTED_CIPHERS[0] + hash_function = "A" * 10 + + try: + Blob(data, cipher, hash_function) + assert(False, "Able to create blob with wrong data") + + except ValueError: + assert True + + +def test_encrypt(): + # Valid data, valid key + data = urandom(64).hex() + blob = Blob(data, SUPPORTED_CIPHERS[0], SUPPORTED_HASH_FUNCTIONS[0]) + key = urandom(32).hex() + + encrypted_blob = blob.encrypt(key) + + # Invalid key (note that encrypt cannot be called with invalid data since that's checked when the Blob is created) + invalid_key = urandom(32) + + try: + blob.encrypt(invalid_key) + assert (False, "Able to create encrypt with invalid key") + + except ValueError: + assert True + + # Check that two encryptions of the same data have the same result + encrypted_blob2 = blob.encrypt(key) + + assert(encrypted_blob == encrypted_blob2 and id(encrypted_blob) != id(encrypted_blob2)) + + +logging.getLogger().disabled = True + +test_init_blob() +test_encrypt() +