Added checks for valid PublicKeys in encryption/decryption functions

.gitignore

@@ -74,3 +74,6 @@ gen
 
 # MyPy Cache
 .mypy_cache/
+
+# OS specific
+.DS_Store

tests/test_base.py

@@ -1,6 +1,7 @@
-import libnacl
 import pytest
-
+from packaging.version import Version
+import libnacl
+import libnacl.public
 from threema.gateway import (
     e2e,
     key,
@@ -32,3 +33,26 @@ def test_valid(self):
         _, data_encrypted = e2e._pk_encrypt(key_pair, data_in, nonce=nonce)
         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):
+        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)
+        assert 'Unable to compute shared key' in str(exc_info)
+

threema/gateway/e2e.py

@@ -82,6 +82,8 @@ 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")
     box = libnacl.public.Box(sk=private, pk=public)
     return box.encrypt(data, nonce=nonce, pack_nonce=False)
 
@@ -103,6 +105,8 @@ 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")
     box = libnacl.public.Box(sk=private, pk=public)
     return box.decrypt(data, nonce=nonce)
 

threema/gateway/key.py

@@ -154,3 +154,23 @@ def derive_public(private_key):
         Return the :class:`libnacl.public.PublicKey` instance.
         """
         return libnacl.public.PublicKey(private_key.pk)
+
+    @staticmethod
+    def is_valid_public_key(public_key):
+        """
+        Check if the 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.
+        """
+        if not isinstance(public_key, libnacl.public.PublicKey):
+            return False
+        if len(public_key.pk) != libnacl.crypto_box_PUBLICKEYBYTES:
+            return False
+
+        # Reject all-zero public keys
+        zero_public_key = libnacl.public.PublicKey(bytes(libnacl.crypto_box_PUBLICKEYBYTES))
+        return public_key != zero_public_key