Merge pull request #86 from matrix-org/add_python_pk_signing

add python bindings for PK signing

1 files changed, 110 insertions, 2 deletions

diff --git a/python/olm/pk.py b/python/olm/pk.py
index b67d5a4..193aba5 100644
--- a/python/olm/pk.py
+++ b/python/olm/pk.py
@@ -17,7 +17,8 @@
 
 This module contains bindings to the PK part of the Olm library.
 It contains two classes PkDecryption and PkEncryption that are used to
-establish an encrypted communication channel using public key encryption.
+establish an encrypted communication channel using public key encryption,
+as well as a class PkSigning that is used to sign a message.
 
 Examples:
     >>> decryption = PkDecryption()
@@ -25,16 +26,22 @@ Examples:
     >>> plaintext = "It's a secret to everybody."
     >>> message = encryption.encrypt(plaintext)
     >>> decrypted_plaintext = decryption.decrypt(message)
+    >>> seed = PkSigning.generate_seed()
+    >>> signing = PkSigning(seed)
+    >>> signature = signing.sign(plaintext)
+    >>> ed25519_verify(signing.public_key, plaintext, signature)
 
 """
 
 from builtins import super
 from typing import AnyStr, Type
+
 from future.utils import bytes_to_native_str
 
 from _libolm import ffi, lib  # type: ignore
-from ._finalize import track_for_finalization
+
 from ._compat import URANDOM, to_bytearray
+from ._finalize import track_for_finalization
 
 
 class PkEncryptionError(Exception):
@@ -45,6 +52,10 @@ class PkDecryptionError(Exception):
     """libolm Pk decryption exception."""
 
 
+class PkSigningError(Exception):
+    """libolm Pk signing exception."""
+
+
 def _clear_pk_encryption(pk_struct):
     lib.olm_clear_pk_encryption(pk_struct)
 
@@ -344,3 +355,100 @@ class PkDecryption(object):
         lib.memset(plaintext_buffer, 0, max_plaintext_length)
 
         return bytes_to_native_str(plaintext)
+
+
+def _clear_pk_signing(pk_struct):
+    lib.olm_clear_pk_signing(pk_struct)
+
+
+class PkSigning(object):
+    """PkSigning class.
+
+    Signs messages using public key cryptography.
+
+    Attributes:
+        public_key (str): The public key of the PkSigning object, can be
+            shared and used to verify using Utility.ed25519_verify.
+
+    """
+
+    def __init__(self, seed):
+        # type: (bytes) -> None
+        """Create a new signing object.
+
+        Args:
+            seed(bytes): the seed to use as the private key for signing.  The
+                seed must have the same length as the seeds generated by
+                PkSigning.generate_seed().
+        """
+        if not seed:
+            raise ValueError("seed can't be empty")
+
+        self._buf = ffi.new("char[]", lib.olm_pk_signing_size())
+        self._pk_signing = lib.olm_pk_signing(self._buf)
+        track_for_finalization(self, self._pk_signing, _clear_pk_signing)
+
+        seed_buffer = ffi.new("char[]", seed)
+
+        pubkey_length = lib.olm_pk_signing_public_key_length()
+        pubkey_buffer = ffi.new("char[]", pubkey_length)
+
+        ret = lib.olm_pk_signing_key_from_seed(
+            self._pk_signing,
+            pubkey_buffer, pubkey_length,
+            seed_buffer, len(seed)
+        )
+
+        # zero out copies of the seed
+        lib.memset(seed_buffer, 0, len(seed))
+
+        self._check_error(ret)
+
+        self.public_key = bytes_to_native_str(
+            ffi.unpack(pubkey_buffer, pubkey_length)
+        )
+
+    def _check_error(self, ret):
+        # type: (int) -> None
+        if ret != lib.olm_error():
+            return
+
+        last_error = bytes_to_native_str(
+            ffi.string(lib.olm_pk_signing_last_error(self._pk_signing)))
+
+        raise PkSigningError(last_error)
+
+    @classmethod
+    def generate_seed(cls):
+        # type: () -> bytes
+        """Generate a random seed.
+        """
+        random_length = lib.olm_pk_signing_seed_length()
+        random = URANDOM(random_length)
+
+        return random
+
+    def sign(self, message):
+        # type: (AnyStr) -> str
+        """Sign a message
+
+        Returns the signature.
+        Raises PkSigningError on failure.
+
+        Args:
+            message(str): the message to sign.
+        """
+        bytes_message = to_bytearray(message)
+
+        signature_length = lib.olm_pk_signature_length()
+        signature_buffer = ffi.new("char[]", signature_length)
+
+        ret = lib.olm_pk_sign(
+            self._pk_signing,
+            ffi.from_buffer(bytes_message), len(bytes_message),
+            signature_buffer, signature_length)
+        self._check_error(ret)
+
+        return bytes_to_native_str(
+            ffi.unpack(signature_buffer, signature_length)
+        )