fixup! Added checks for valid PublicKeys in encryption/decryption functions

Author: threema-lenny

tests/test_base.py

@@ -1,7 +1,8 @@
-import pytest
-from packaging.version import Version
 import libnacl
 import libnacl.public
+import pytest
+from packaging.version import Version
+
 from threema.gateway import (
     e2e,
     key,
@@ -26,6 +27,21 @@ def test_incorrect_ciphertext(self):
             e2e._pk_decrypt(key_pair, nonce, data_encrypted + b'0')
         assert 'decrypt' in str(exc_info.value)
 
+    def test_invalid_non_contributory_public_key(self):
+        all_zero_public_key = libnacl.public.PublicKey(
+            bytes(libnacl.crypto_box_PUBLICKEYBYTES))
+        key_pair = key.Key.generate_secret_key, all_zero_public_key
+        data_in = b"meow"
+        nonce = b"0" * 24
+
+        with pytest.raises(libnacl.CryptError) as exc_info:
+            e2e._pk_encrypt(key_pair, data_in, nonce=nonce)
+        assert "Invalid public key (non-contributory)" in str(exc_info.value)
+
+        with pytest.raises(libnacl.CryptError) as exc_info:
+            e2e._pk_decrypt(key_pair, nonce, bytes(5))
+        assert "Invalid public key (non-contributory)" in str(exc_info.value)
+
     def test_valid(self):
         key_pair = key.Key.generate_pair()
         data_in = b'meow'
@@ -34,25 +50,12 @@ def test_valid(self):
         data_out = e2e._pk_decrypt(key_pair, nonce, data_encrypted)
         assert data_in == data_out
 
-    def test_invalid_pk(self):
-        all_zero_pk = libnacl.public.PublicKey(bytes(libnacl.crypto_box_PUBLICKEYBYTES))
-        key_pair = key.Key.generate_secret_key, all_zero_pk
-        data_in = b'meow'
-        nonce = b'0' * 24
-        with pytest.raises(libnacl.CryptError) as exc_info:
-            e2e._pk_encrypt(key_pair, data_in, nonce=nonce)
-        assert 'Invalid public key' in str(exc_info.value)
-
-        with pytest.raises(libnacl.CryptError) as exc_info:
-            e2e._pk_decrypt(key_pair, nonce, bytes(5))
-        assert 'Invalid public key' in str(exc_info.value)
-
-    @pytest.mark.skipif(Version(libnacl.sodium_version_string().decode("ascii")) < Version("1.0.7"),
-                        reason="no zero-result check on X25519 in this version of libsodium")
-    def test_contributory(self):
+    @pytest.mark.skipif(
+        Version(libnacl.sodium_version_string().decode("ascii")) < Version("1.0.7"),
+        reason="no zero-result check on X25519 in this libnacl backend")
+    def test_libsodium_non_contributory_public_key(self):
         all_zero_pk = libnacl.public.PublicKey(bytes(libnacl.crypto_box_PUBLICKEYBYTES))
         alice_sk, _ = key.Key.generate_pair()
         with pytest.raises(libnacl.CryptError) as exc_info:
-            alice_box = libnacl.public.Box(alice_sk, all_zero_pk)
+            libnacl.public.Box(alice_sk, all_zero_pk)
         assert 'Unable to compute shared key' in str(exc_info)
-

threema/gateway/e2e.py

@@ -82,8 +82,7 @@ def _pk_encrypt(key_pair: Tuple[Key, Key], data: bytes, nonce: Optional[bytes] =
     """
     # Assemble and encrypt the payload
     private, public = key_pair
-    if not Key.is_valid_public_key(public):
-        raise libnacl.CryptError("Invalid public key")
+    Key.ensure_valid_public_key(public)
     box = libnacl.public.Box(sk=private, pk=public)
     return box.encrypt(data, nonce=nonce, pack_nonce=False)
 
@@ -105,8 +104,7 @@ def _pk_decrypt(key_pair: Tuple[Key, Key], nonce: bytes, data: bytes):
     """
     # Decrypt payload
     private, public = key_pair
-    if not Key.is_valid_public_key(public):
-        raise libnacl.CryptError("Invalid public key")
+    Key.ensure_valid_public_key(public)
     box = libnacl.public.Box(sk=private, pk=public)
     return box.decrypt(data, nonce=nonce)
 

threema/gateway/key.py

@@ -42,6 +42,10 @@ def hash(message, hash_type):
         return hmac.new(HMAC.keys[hash_type], message.encode('ascii'), hashlib.sha256)
 
 
+_NON_CONTRIBUTORY_PUBLIC_KEY = libnacl.public.PublicKey(
+    bytes(libnacl.crypto_box_PUBLICKEYBYTES))
+
+
 class Key:
     """
     Encode or decode a key.
@@ -156,21 +160,21 @@ def derive_public(private_key):
         return libnacl.public.PublicKey(private_key.pk)
 
     @staticmethod
-    def is_valid_public_key(public_key):
+    def ensure_valid_public_key(public_key):
         """
-        Check if the public key is valid.
+        Ensure a public key is valid.
 
         Arguments:
             - `public_key`: An instance of
               :class:`libnacl.public.PublicKey`.
 
-        Return True if the key is a valid public key, False if not.
+        Raises :class:`libnacl.CryptError` if the public key is invalid.
         """
         if not isinstance(public_key, libnacl.public.PublicKey):
-            return False
+            raise libnacl.CryptError("Invalid public key")
         if len(public_key.pk) != libnacl.crypto_box_PUBLICKEYBYTES:
-            return False
+            raise libnacl.CryptError("Invalid public key")
 
-        # Reject all-zero public keys
-        zero_public_key = libnacl.public.PublicKey(bytes(libnacl.crypto_box_PUBLICKEYBYTES))
-        return public_key != zero_public_key
+        # Reject all-zero public keys (mon-contributory)
+        if public_key == _NON_CONTRIBUTORY_PUBLIC_KEY:
+            raise libnacl.CryptError("Invalid public key (non-contributory)")