Harden malformed storage payload handling

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."""

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"