python: Add Short Authentication String bindings.

This patch adds bindings to the SAS part of the Olm library contained in
the sas.h header file.

Signed-off-by: Damir Jelić <poljar@termina.org.uk>

5 files changed, 315 insertions, 2 deletions

diff --git a/python/Makefile b/python/Makefile
index 5da703a..6283fb5 100644
--- a/python/Makefile
+++ b/python/Makefile
@@ -6,10 +6,13 @@ include/olm/olm.h: ../include/olm/olm.h ../include/olm/inbound_group_session.h .
 # add memset to the header so that we can use it to clear buffers
 	echo 'void *memset(void *s, int c, size_t n);' >> include/olm/olm.h
 
-olm-python2: include/olm/olm.h
+include/olm/sas.h: include/olm/olm.h ../include/olm/sas.h
+	$(CPP) -I dummy -I ../include ../include/olm/sas.h -o include/olm/sas.h
+
+olm-python2: include/olm/olm.h include/olm/sas.h
 	DEVELOP=$(DEVELOP) python2 setup.py build
 
-olm-python3: include/olm/olm.h
+olm-python3: include/olm/olm.h include/olm/sas.h
 	DEVELOP=$(DEVELOP) python3 setup.py build
 
 install: install-python2 install-python3
diff --git a/python/olm/__init__.py b/python/olm/__init__.py
index 015c4f8..fe3e199 100644
--- a/python/olm/__init__.py
+++ b/python/olm/__init__.py
@@ -36,3 +36,4 @@ from .group_session import (
     OutboundGroupSession,
     OlmGroupSessionError
 )
+from .sas import Sas, OlmSasError
diff --git a/python/olm/sas.py b/python/olm/sas.py
new file mode 100644
index 0000000..c7cccb2
--- /dev/null
+++ b/python/olm/sas.py
@@ -0,0 +1,223 @@
+# -*- coding: utf-8 -*-
+# libolm python bindings
+# Copyright © 2015-2017 OpenMarket Ltd
+# Copyright © 2019 Damir Jelić <poljar@termina.org.uk>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""libolm SAS module.
+
+This module contains functions to perform key verification using the Short
+Authentication String (SAS) method.
+
+Examples:
+    >>> sas = Sas()
+    >>> bob_key = "3p7bfXt9wbTTW2HC7OQ1Nz+DQ8hbeGdNrfx+FG+IK08"
+    >>> message = "Hello world!"
+    >>> extra_info = "MAC"
+    >>> sas_alice.set_their_pubkey(bob_key)
+    >>> sas_alice.calculate_mac(message, extra_info)
+    >>> sas_alice.generate_bytes(extra_info, 5)
+
+"""
+
+from functools import wraps
+from builtins import bytes
+from typing import Optional
+
+from future.utils import bytes_to_native_str
+
+from _libolm import ffi, lib
+
+from ._compat import URANDOM, to_bytes, to_bytearray
+from ._finalize import track_for_finalization
+
+
+def other_pubkey_set(func):
+    """Ensure that the other pubkey is added to the Sas object."""
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        if not self.other_key_set:
+            raise OlmSasError("The other public key isn't set.")
+        return func(self, *args, **kwargs)
+    return wrapper
+
+
+def _clear_sas(sas):
+    # type: (ffi.cdata) -> None
+    lib.olm_clear_sas(sas)
+
+
+class OlmSasError(Exception):
+    """libolm Sas error exception."""
+
+
+class Sas(object):
+    """libolm Short Authenticaton String (SAS) class."""
+
+    def __init__(self, other_users_pubkey=None):
+        # type: (Optional[str]) -> None
+        """Create a new SAS object.
+
+        Args:
+            other_users_pubkey(str, optional): The other users public key, this
+                key is necesary to generate bytes for the authentication string
+                as well as to calculate the MAC.
+
+        Attributes:
+            other_key_set (bool): A boolean flag that tracks if we set the
+                other users public key for this SAS object.
+
+        Raises OlmSasError on failure.
+
+        """
+        self._buf = ffi.new("char[]", lib.olm_sas_size())
+        self._sas = lib.olm_sas(self._buf)
+        self.other_key_set = False
+        track_for_finalization(self, self._sas, _clear_sas)
+
+        random_length = lib.olm_create_sas_random_length(self._sas)
+        random = URANDOM(random_length)
+
+        self._create_sas(random, random_length)
+
+        if other_users_pubkey:
+            self.set_their_pubkey(other_users_pubkey)
+
+    def _create_sas(self, buffer, buffer_length):
+        self._check_error(
+            lib.olm_create_sas(
+                self._sas,
+                ffi.from_buffer(buffer),
+                buffer_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_sas_last_error(self._sas))))
+
+        raise OlmSasError(last_error)
+
+    @property
+    def pubkey(self):
+        # type: () -> str
+        """Get the public key for the SAS object.
+
+        This returns the public key of the SAS object that can then be shared
+        with another user to perform the authentication process.
+
+        Raises OlmSasError on failure.
+
+        """
+        pubkey_length = lib.olm_sas_pubkey_length(self._sas)
+        pubkey_buffer = ffi.new("char[]", pubkey_length)
+
+        self._check_error(
+            lib.olm_sas_get_pubkey(self._sas, pubkey_buffer, pubkey_length)
+        )
+
+        return bytes_to_native_str(ffi.unpack(pubkey_buffer, pubkey_length))
+
+    def set_their_pubkey(self, key):
+        # type: (str) -> None
+        """Set the public key of the other user.
+
+        This sets the public key of the other user, it needs to be set before
+        bytes can be generated for the authentication string and a MAC can be
+        calculated.
+
+        Args:
+            key (str): The other users public key.
+
+        Raises OlmSasError on failure.
+
+        """
+        byte_key = to_bytearray(key)
+
+        self._check_error(
+            lib.olm_sas_set_their_key(
+                self._sas,
+                ffi.from_buffer(byte_key),
+                len(byte_key)
+            )
+        )
+        self.other_key_set = True
+
+    @other_pubkey_set
+    def generate_bytes(self, extra_info, length):
+        # type: (str, int) -> bytes
+        """Generate bytes to use for the short authentication string.
+
+        Args:
+            extra_info (str): Extra information to mix in when generating the
+                bytes.
+            length (int): The number of bytes to generate.
+
+        Raises OlmSasError if the other users persons public key isn't set or
+        an internal Olm error happens.
+
+        """
+        if length < 1:
+            raise ValueError("The length needs to be a positive integer value")
+
+        byte_info = to_bytearray(extra_info)
+        out_buffer = ffi.new("char[]", length)
+
+        self._check_error(
+            lib.olm_sas_generate_bytes(
+                self._sas,
+                ffi.from_buffer(byte_info),
+                len(byte_info),
+                out_buffer,
+                length
+            )
+        )
+
+        return ffi.unpack(out_buffer, length)
+
+    @other_pubkey_set
+    def calculate_mac(self, message, extra_info):
+        # type: (str, str) -> str
+        """Generate a message authentication code based on the shared secret.
+
+        Args:
+            message (str): The message to produce the authentication code for.
+            extra_info (str): Extra information to mix in when generating the
+                MAC
+
+        Raises OlmSasError on failure.
+
+        """
+        byte_message = to_bytes(message)
+        byte_info = to_bytes(extra_info)
+
+        mac_length = lib.olm_sas_mac_length(self._sas)
+        mac_buffer = ffi.new("char[]", mac_length)
+
+        self._check_error(
+            lib.olm_sas_calculate_mac(
+                self._sas,
+                ffi.from_buffer(byte_message),
+                len(byte_message),
+                ffi.from_buffer(byte_info),
+                len(byte_info),
+                mac_buffer,
+                mac_length
+            )
+        )
+        return bytes_to_native_str(ffi.unpack(mac_buffer, mac_length))
diff --git a/python/olm_build.py b/python/olm_build.py
index ee836d8..1c610a1 100644
--- a/python/olm_build.py
+++ b/python/olm_build.py
@@ -39,6 +39,7 @@ ffibuilder.set_source(
         #include <olm/olm.h>
         #include <olm/inbound_group_session.h>
         #include <olm/outbound_group_session.h>
+        #include <olm/sas.h>
     """,
     libraries=["olm"],
     extra_compile_args=compile_args,
@@ -47,5 +48,8 @@ ffibuilder.set_source(
 with open(os.path.join(PATH, "include/olm/olm.h")) as f:
     ffibuilder.cdef(f.read(), override=True)
 
+with open(os.path.join(PATH, "include/olm/sas.h")) as f:
+    ffibuilder.cdef(f.read(), override=True)
+
 if __name__ == "__main__":
     ffibuilder.compile(verbose=True)
diff --git a/python/tests/sas_test.py b/python/tests/sas_test.py
new file mode 100644
index 0000000..05013b7
--- /dev/null
+++ b/python/tests/sas_test.py
@@ -0,0 +1,82 @@
+from builtins import bytes
+
+import pytest
+
+from olm import OlmSasError, Sas
+
+MESSAGE = "Test message"
+EXTRA_INFO = "extra_info"
+
+
+class TestClass(object):
+    def test_sas_creation(self):
+        sas = Sas()
+        assert sas.pubkey
+
+    def test_other_key_setting(self):
+        sas_alice = Sas()
+        sas_bob = Sas()
+
+        assert not sas_alice.other_key_set
+        sas_alice.set_their_pubkey(sas_bob.pubkey)
+        assert sas_alice.other_key_set
+
+    def test_bytes_generating(self):
+        sas_alice = Sas()
+        sas_bob = Sas(sas_alice.pubkey)
+
+        assert sas_bob.other_key_set
+
+        with pytest.raises(OlmSasError):
+            sas_alice.generate_bytes(EXTRA_INFO, 5)
+
+        sas_alice.set_their_pubkey(sas_bob.pubkey)
+
+        with pytest.raises(ValueError):
+            sas_alice.generate_bytes(EXTRA_INFO, 0)
+
+        alice_bytes = sas_alice.generate_bytes(EXTRA_INFO, 5)
+        bob_bytes = sas_bob.generate_bytes(EXTRA_INFO, 5)
+
+        assert alice_bytes == bob_bytes
+
+    def test_mac_generating(self):
+        sas_alice = Sas()
+        sas_bob = Sas()
+
+        with pytest.raises(OlmSasError):
+            sas_alice.calculate_mac(MESSAGE, EXTRA_INFO)
+
+        sas_alice.set_their_pubkey(sas_bob.pubkey)
+        sas_bob.set_their_pubkey(sas_alice.pubkey)
+
+        alice_mac = sas_alice.calculate_mac(MESSAGE, EXTRA_INFO)
+        bob_mac = sas_bob.calculate_mac(MESSAGE, EXTRA_INFO)
+
+        assert alice_mac == bob_mac
+
+    def test_cross_language_mac(self):
+        """Test MAC generating with a predefined key pair.
+
+        This test imports a private and public key from the C test and checks
+        if we are getting the same MAC that the C code calculated.
+        """
+        alice_private = [
+            0x77, 0x07, 0x6D, 0x0A, 0x73, 0x18, 0xA5, 0x7D,
+            0x3C, 0x16, 0xC1, 0x72, 0x51, 0xB2, 0x66, 0x45,
+            0xDF, 0x4C, 0x2F, 0x87, 0xEB, 0xC0, 0x99, 0x2A,
+            0xB1, 0x77, 0xFB, 0xA5, 0x1D, 0xB9, 0x2C, 0x2A
+        ]
+
+        bob_key = "3p7bfXt9wbTTW2HC7OQ1Nz+DQ8hbeGdNrfx+FG+IK08"
+        message = "Hello world!"
+        extra_info = "MAC"
+        expected_mac = "2nSMTXM+TStTU3RUVTNSVVZUTlNWVlpVVGxOV1ZscFY"
+
+        sas_alice = Sas()
+        sas_alice._create_sas(bytes(alice_private), 32)
+        sas_alice.set_their_pubkey(bob_key)
+
+        alice_mac = sas_alice.calculate_mac(message, extra_info)
+
+        assert alice_mac == expected_mac