diff --git a/bitcoin_safe_lib/storage.py b/bitcoin_safe_lib/storage.py index 8f0d36f..80fa980 100644 --- a/bitcoin_safe_lib/storage.py +++ b/bitcoin_safe_lib/storage.py @@ -37,8 +37,10 @@ # from https://stackoverflow.com/questions/2490334/simple-way-to-encode-a-string-according-to-a-password import secrets from abc import abstractmethod +from base64 import b64decode from base64 import urlsafe_b64decode as b64d from base64 import urlsafe_b64encode as b64e +from binascii import Error as BinasciiError from collections.abc import Callable, Iterable from pathlib import Path from typing import Any, ClassVar, Protocol, TypeAlias, TypeGuard, TypeVar @@ -69,6 +71,15 @@ def from_dump(cls, dct: dict[str, Any], class_kwargs: ClassKwargs | None = None) logger = logging.getLogger(__name__) +MIN_ENCRYPTED_PAYLOAD_SIZE = 16 + 4 + 1 +MAX_ENCRYPT_ITERATIONS = 1_000_000 +ASCII_WHITESPACE = b" \t\n\r\f\v" + + +class StorageDecryptError(Exception): + """Raised when an encrypted storage payload is malformed.""" + + def varnames(method: Callable) -> Iterable[str]: """Varnames.""" return method.__code__.co_varnames[: method.__code__.co_argcount] @@ -85,6 +96,37 @@ def filtered_for_init(d: dict, cls: type[Any]) -> dict: class Encrypt: + @staticmethod + def _invalid_encrypted_payload() -> StorageDecryptError: + return StorageDecryptError("Invalid encrypted payload") + + @classmethod + def _parse_encrypted_payload(cls, token: bytes) -> tuple[bytes, int, bytes]: + try: + decoded = b64decode(token.strip(ASCII_WHITESPACE), altchars=b"-_", validate=True) + except (BinasciiError, ValueError) as e: + raise cls._invalid_encrypted_payload() from e + + if len(decoded) < MIN_ENCRYPTED_PAYLOAD_SIZE: + raise cls._invalid_encrypted_payload() + + salt = decoded[:16] + iteration_bytes = decoded[16:20] + encrypted_message = decoded[20:] + iterations = int.from_bytes(iteration_bytes, "big") + if not 1 <= iterations <= MAX_ENCRYPT_ITERATIONS: + raise cls._invalid_encrypted_payload() + + return salt, iterations, b64e(encrypted_message) + + @classmethod + def is_encrypted_payload(cls, token: bytes) -> bool: + try: + cls._parse_encrypted_payload(token) + except StorageDecryptError: + return False + return True + def _derive_key(self, password: bytes, salt: bytes, iterations: int) -> bytes: """Derive a secret key from a given password and salt.""" kdf = PBKDF2HMAC( @@ -111,11 +153,7 @@ def password_encrypt(self, message: bytes, password: str, iterations: int = 100_ def password_decrypt(self, token: bytes, password: str) -> bytes: """Password decrypt.""" - decoded = b64d(token) - salt, iter, token = decoded[:16], decoded[16:20], b64e(decoded[20:]) - iterations = int.from_bytes(iter, "big") - if iterations > 1e6: - raise Exception("Error in decrypting") + salt, iterations, token = self._parse_encrypted_payload(token) key = self._derive_key(password.encode(), salt, iterations) return Fernet(key).decrypt(token) @@ -145,15 +183,10 @@ def has_password(cls, filename: str) -> bool: # - carriage rtn (\r) # - form feed (\f) # - vertical tab (\v) - strip_char = b" \t\n\r\f\v" - try: - if token.lstrip(strip_char).startswith(b"{") and token.rstrip(strip_char).endswith(b"}"): - return False - except Exception as e: - logger.exception(str(e)) - return True - - return True + stripped_token = token.strip(ASCII_WHITESPACE) + if stripped_token.startswith(b"{") and stripped_token.endswith(b"}"): + return False + return Encrypt.is_encrypted_payload(token) def load(self, filename: str, password: str | None = None) -> str: """Load.""" diff --git a/tests/test_storage.py b/tests/test_storage.py index 2203e7b..91f0237 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -1,6 +1,10 @@ import logging +from base64 import urlsafe_b64encode as b64e -from bitcoin_safe_lib.storage import SaveAllClass +import pytest +from cryptography.fernet import InvalidToken + +from bitcoin_safe_lib.storage import Encrypt, SaveAllClass, Storage, StorageDecryptError class ExampleSaveable(SaveAllClass): @@ -27,3 +31,77 @@ def test_from_dumps_prefers_dct_values_over_class_kwargs(caplog) -> None: assert "Duplicate deserialization keys" in caplog.text assert "ExampleSaveable" in caplog.text assert "value" in caplog.text + + +def make_payload(iterations: int, encrypted_message: bytes = b"x") -> bytes: + return b64e((b"s" * 16) + iterations.to_bytes(4, "big") + encrypted_message) + + +def test_password_encrypt_decrypt_round_trip() -> None: + encrypt = Encrypt() + + token = encrypt.password_encrypt(b"secret", "pw") + + assert encrypt.password_decrypt(token, "pw") == b"secret" + + +@pytest.mark.parametrize("wrapper", [b" %b ", b"\t%b\t", b"\n%b\n", b" \t\n%b\n\t "]) +def test_password_decrypt_strips_outer_whitespace(wrapper: bytes) -> None: + encrypt = Encrypt() + token = encrypt.password_encrypt(b"secret", "pw") + + wrapped_token = wrapper % token + + assert encrypt.password_decrypt(wrapped_token, "pw") == b"secret" + + +@pytest.mark.parametrize( + ("token", "label"), + [ + (b"", "empty"), + (b'{"foo": 1', "truncated json"), + (b"AAAA", "short base64"), + (b"!!!", "invalid base64"), + (make_payload(0), "zero iterations"), + (make_payload(1_000_001), "too many iterations"), + ], +) +def test_password_decrypt_rejects_malformed_payloads(token: bytes, label: str) -> None: + encrypt = Encrypt() + + with pytest.raises(StorageDecryptError, match="Invalid encrypted payload"): + encrypt.password_decrypt(token, "pw") + + +def test_password_decrypt_wrong_password_raises_invalid_token() -> None: + encrypt = Encrypt() + token = encrypt.password_encrypt(b"secret", "pw") + + with pytest.raises(InvalidToken): + encrypt.password_decrypt(token, "wrong") + + +def test_has_password_detects_plaintext_and_encrypted_files(tmp_path) -> None: + plaintext = tmp_path / "plain.wallet" + truncated = tmp_path / "truncated.wallet" + encrypted = tmp_path / "encrypted.wallet" + encrypted_with_whitespace = tmp_path / "encrypted_with_whitespace.wallet" + + plaintext.write_text('{"foo": 1}') + truncated.write_bytes(b'{"foo": 1') + encrypted_token = Encrypt().password_encrypt(b"secret", "pw") + encrypted.write_bytes(encrypted_token) + encrypted_with_whitespace.write_bytes(b" \t\n" + encrypted_token + b"\n\t ") + + assert not Storage.has_password(str(plaintext)) + assert not Storage.has_password(str(truncated)) + assert Storage.has_password(str(encrypted)) + assert Storage.has_password(str(encrypted_with_whitespace)) + + +def test_load_decrypts_encrypted_file_wrapped_in_whitespace(tmp_path) -> None: + encrypted = tmp_path / "encrypted_with_whitespace.wallet" + encrypted_token = Encrypt().password_encrypt(b"secret", "pw") + encrypted.write_bytes(b" \t\n" + encrypted_token + b"\n\t ") + + assert Storage().load(str(encrypted), password="pw") == "secret"