From 5fe88ea0500c6e418492f4b166c0d4a24e9632cc Mon Sep 17 00:00:00 2001
From: Paul Kehrer <paul.l.kehrer@gmail.com>
Date: Sat, 19 Jan 2019 00:52:43 -0600
Subject: shake128/256 support (#4611)

* shake128/256 support

* remove block_size

* doc an exception

* change how we detect XOF by adding _xof attribute

* interface!

* review feedback
---
 src/cryptography/hazmat/backends/openssl/hashes.py | 25 ++++++++++---
 src/cryptography/hazmat/primitives/hashes.py       | 41 ++++++++++++++++++++++
 2 files changed, 61 insertions(+), 5 deletions(-)

(limited to 'src')

diff --git a/src/cryptography/hazmat/backends/openssl/hashes.py b/src/cryptography/hazmat/backends/openssl/hashes.py
index e9a50705..549fa2bf 100644
--- a/src/cryptography/hazmat/backends/openssl/hashes.py
+++ b/src/cryptography/hazmat/backends/openssl/hashes.py
@@ -54,10 +54,25 @@ class _HashContext(object):
         self._backend.openssl_assert(res != 0)
 
     def finalize(self):
+        if isinstance(self.algorithm, hashes.ExtendableOutputFunction):
+            # extendable output functions use a different finalize
+            return self._finalize_xof()
+        else:
+            buf = self._backend._ffi.new("unsigned char[]",
+                                         self._backend._lib.EVP_MAX_MD_SIZE)
+            outlen = self._backend._ffi.new("unsigned int *")
+            res = self._backend._lib.EVP_DigestFinal_ex(self._ctx, buf, outlen)
+            self._backend.openssl_assert(res != 0)
+            self._backend.openssl_assert(
+                outlen[0] == self.algorithm.digest_size
+            )
+            return self._backend._ffi.buffer(buf)[:outlen[0]]
+
+    def _finalize_xof(self):
         buf = self._backend._ffi.new("unsigned char[]",
-                                     self._backend._lib.EVP_MAX_MD_SIZE)
-        outlen = self._backend._ffi.new("unsigned int *")
-        res = self._backend._lib.EVP_DigestFinal_ex(self._ctx, buf, outlen)
+                                     self.algorithm.digest_size)
+        res = self._backend._lib.EVP_DigestFinalXOF(
+            self._ctx, buf, self.algorithm.digest_size
+        )
         self._backend.openssl_assert(res != 0)
-        self._backend.openssl_assert(outlen[0] == self.algorithm.digest_size)
-        return self._backend._ffi.buffer(buf)[:outlen[0]]
+        return self._backend._ffi.buffer(buf)[:self.algorithm.digest_size]
diff --git a/src/cryptography/hazmat/primitives/hashes.py b/src/cryptography/hazmat/primitives/hashes.py
index 0d6e47fb..9be2b600 100644
--- a/src/cryptography/hazmat/primitives/hashes.py
+++ b/src/cryptography/hazmat/primitives/hashes.py
@@ -57,6 +57,13 @@ class HashContext(object):
         """
 
 
+@six.add_metaclass(abc.ABCMeta)
+class ExtendableOutputFunction(object):
+    """
+    An interface for extendable output functions.
+    """
+
+
 @utils.register_interface(HashContext)
 class Hash(object):
     def __init__(self, algorithm, backend, ctx=None):
@@ -173,6 +180,40 @@ class SHA3_512(object):  # noqa: N801
     digest_size = 64
 
 
+@utils.register_interface(HashAlgorithm)
+@utils.register_interface(ExtendableOutputFunction)
+class SHAKE128(object):
+    name = "shake128"
+
+    def __init__(self, digest_size):
+        if not isinstance(digest_size, six.integer_types):
+            raise TypeError("digest_size must be an integer")
+
+        if digest_size < 1:
+            raise ValueError("digest_size must be a positive integer")
+
+        self._digest_size = digest_size
+
+    digest_size = utils.read_only_property("_digest_size")
+
+
+@utils.register_interface(HashAlgorithm)
+@utils.register_interface(ExtendableOutputFunction)
+class SHAKE256(object):
+    name = "shake256"
+
+    def __init__(self, digest_size):
+        if not isinstance(digest_size, six.integer_types):
+            raise TypeError("digest_size must be an integer")
+
+        if digest_size < 1:
+            raise ValueError("digest_size must be a positive integer")
+
+        self._digest_size = digest_size
+
+    digest_size = utils.read_only_property("_digest_size")
+
+
 @utils.register_interface(HashAlgorithm)
 class MD5(object):
     name = "md5"
-- 
cgit v1.2.3