From fcfa5f12a4cd482973fdf03000af4a26a360dc2e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Damir=20Jeli=C4=87?= <poljar@termina.org.uk>
Date: Tue, 2 Apr 2019 12:56:06 +0200
Subject: [PATCH] python: Expose the sha256() function in the utilities.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Damir Jelić <poljar@termina.org.uk>
---
 python/olm/__init__.py     |  2 +-
 python/olm/utility.py      | 57 ++++++++++++++++++++++++++++++++------
 python/tests/utils_test.py | 29 +++++++++++++++++++
 3 files changed, 78 insertions(+), 10 deletions(-)
 create mode 100644 python/tests/utils_test.py

diff --git a/python/olm/__init__.py b/python/olm/__init__.py
index fe3e199..9ac45b0 100644
--- a/python/olm/__init__.py
+++ b/python/olm/__init__.py
@@ -21,7 +21,7 @@ Olm Python bindings
 |  © Copyright 2015-2017 by OpenMarket Ltd
 |  © Copyright 2018 by Damir Jelić
 """
-from .utility import ed25519_verify, OlmVerifyError
+from .utility import ed25519_verify, OlmVerifyError, OlmHashError, sha256
 from .account import Account, OlmAccountError
 from .session import (
     Session,
diff --git a/python/olm/utility.py b/python/olm/utility.py
index 121ff63..10d5ab4 100644
--- a/python/olm/utility.py
+++ b/python/olm/utility.py
@@ -32,6 +32,7 @@ Examples:
 
 # pylint: disable=redefined-builtin,unused-import
 from typing import AnyStr, Type
+from future.utils import bytes_to_native_str
 
 # pylint: disable=no-name-in-module
 from _libolm import ffi, lib  # type: ignore
@@ -49,6 +50,10 @@ class OlmVerifyError(Exception):
     """libolm signature verification exception."""
 
 
+class OlmHashError(Exception):
+    """libolm hash calculation exception."""
+
+
 class _Utility(object):
     # pylint: disable=too-few-public-methods
     """libolm Utility class."""
@@ -64,12 +69,12 @@ class _Utility(object):
         track_for_finalization(cls, cls._utility, _clear_utility)
 
     @classmethod
-    def _check_error(cls, ret):
-        # type: (int) -> None
+    def _check_error(cls, ret, error_class):
+        # type: (int, Type) -> None
         if ret != lib.olm_error():
             return
 
-        raise OlmVerifyError("{}".format(
+        raise error_class("{}".format(
             ffi.string(lib.olm_utility_last_error(
                 cls._utility)).decode("utf-8")))
 
@@ -84,18 +89,41 @@ class _Utility(object):
         byte_signature = to_bytearray(signature)
 
         try:
-            cls._check_error(
-                lib.olm_ed25519_verify(cls._utility, byte_key, len(byte_key),
-                                       ffi.from_buffer(byte_message),
-                                       len(byte_message),
-                                       ffi.from_buffer(byte_signature),
-                                       len(byte_signature)))
+            ret = lib.olm_ed25519_verify(
+                cls._utility,
+                byte_key,
+                len(byte_key),
+                ffi.from_buffer(byte_message),
+                len(byte_message),
+                ffi.from_buffer(byte_signature),
+                len(byte_signature)
+            )
+
+            cls._check_error(ret, OlmVerifyError)
+
         finally:
             # clear out copies of the message, which may be a plaintext
             if byte_message is not message:
                 for i in range(0, len(byte_message)):
                     byte_message[i] = 0
 
+    @classmethod
+    def _sha256(cls, input):
+        # type: (Type[_Utility], AnyStr) -> str
+        if not cls._utility:
+            cls._allocate()
+
+        byte_input = to_bytes(input)
+        hash_length = lib.olm_sha256_length(cls._utility)
+        hash = ffi.new("char[]", hash_length)
+
+        ret = lib.olm_sha256(cls._utility, byte_input, len(byte_input),
+                             hash, hash_length)
+
+        cls._check_error(ret, OlmHashError)
+
+        return bytes_to_native_str(ffi.unpack(hash, hash_length))
+
 
 def ed25519_verify(key, message, signature):
     # type: (AnyStr, AnyStr, AnyStr) -> None
@@ -109,3 +137,14 @@ def ed25519_verify(key, message, signature):
         signature(bytes): The message signature.
     """
     return _Utility._ed25519_verify(key, message, signature)
+
+
+def sha256(input_string):
+    # type: (AnyStr) -> str
+    """Calculate the SHA-256 hash of the input and encodes it as base64.
+
+    Args:
+        input_string(str): The input for which the hash will be calculated.
+
+    """
+    return _Utility._sha256(input_string)
diff --git a/python/tests/utils_test.py b/python/tests/utils_test.py
new file mode 100644
index 0000000..a552d12
--- /dev/null
+++ b/python/tests/utils_test.py
@@ -0,0 +1,29 @@
+import base64
+import hashlib
+
+from future.utils import bytes_to_native_str
+from hypothesis import given
+from hypothesis.strategies import text
+
+from olm import sha256
+from olm._compat import to_bytes
+
+
+class TestClass(object):
+    @given(text(), text())
+    def test_sha256(self, input1, input2):
+        first_hash = sha256(input1)
+        second_hash = sha256(input2)
+
+        hashlib_hash = base64.b64encode(
+            hashlib.sha256(to_bytes(input1)).digest()
+        )
+
+        hashlib_hash = bytes_to_native_str(hashlib_hash[:-1])
+
+        if input1 == input2:
+            assert first_hash == second_hash
+        else:
+            assert first_hash != second_hash
+
+        assert hashlib_hash == first_hash