Merge pull request #248 from dreid/hash-context-from-backend

Get a HashContext from the backend like we do a CipherContext

2 files changed, 48 insertions, 34 deletions

diff --git a/cryptography/hazmat/bindings/openssl/backend.py b/cryptography/hazmat/bindings/openssl/backend.py
index 1cb886dc..235a8e30 100644
--- a/cryptography/hazmat/bindings/openssl/backend.py
+++ b/cryptography/hazmat/bindings/openssl/backend.py
@@ -269,47 +269,62 @@ class Ciphers(object):
                               _CipherContext._DECRYPT)
 
 
-class Hashes(object):
-    def __init__(self, backend):
-        super(Hashes, self).__init__()
+@interfaces.register(interfaces.HashContext)
+class _HashContext(object):
+    def __init__(self, backend, algorithm, ctx=None):
+        self.algorithm = algorithm
+
         self._backend = backend
 
-    def supported(self, hash_cls):
-        return (self._backend.ffi.NULL !=
-                self._backend.lib.EVP_get_digestbyname(
-                    hash_cls.name.encode("ascii")))
+        if ctx is None:
+            ctx = self._backend.lib.EVP_MD_CTX_create()
+            ctx = self._backend.ffi.gc(ctx,
+                                       self._backend.lib.EVP_MD_CTX_destroy)
+            evp_md = self._backend.lib.EVP_get_digestbyname(
+                algorithm.name.encode("ascii"))
+            assert evp_md != self._backend.ffi.NULL
+            res = self._backend.lib.EVP_DigestInit_ex(ctx, evp_md,
+                                                      self._backend.ffi.NULL)
+            assert res != 0
 
-    def create_ctx(self, hashobject):
-        ctx = self._backend.lib.EVP_MD_CTX_create()
-        ctx = self._backend.ffi.gc(ctx, self._backend.lib.EVP_MD_CTX_destroy)
-        evp_md = self._backend.lib.EVP_get_digestbyname(
-            hashobject.name.encode("ascii"))
-        assert evp_md != self._backend.ffi.NULL
-        res = self._backend.lib.EVP_DigestInit_ex(ctx, evp_md,
-                                                  self._backend.ffi.NULL)
+        self._ctx = ctx
+
+    def copy(self):
+        copied_ctx = self._backend.lib.EVP_MD_CTX_create()
+        copied_ctx = self._backend.ffi.gc(copied_ctx,
+                                          self._backend.lib.EVP_MD_CTX_destroy)
+        res = self._backend.lib.EVP_MD_CTX_copy_ex(copied_ctx, self._ctx)
         assert res != 0
-        return ctx
+        return _HashContext(self._backend, self.algorithm, ctx=copied_ctx)
 
-    def update_ctx(self, ctx, data):
-        res = self._backend.lib.EVP_DigestUpdate(ctx, data, len(data))
+    def update(self, data):
+        res = self._backend.lib.EVP_DigestUpdate(self._ctx, data, len(data))
         assert res != 0
 
-    def finalize_ctx(self, ctx, digest_size):
-        buf = self._backend.ffi.new("unsigned char[]", digest_size)
-        res = self._backend.lib.EVP_DigestFinal_ex(ctx, buf,
+    def finalize(self):
+        buf = self._backend.ffi.new("unsigned char[]",
+                                    self.algorithm.digest_size)
+        res = self._backend.lib.EVP_DigestFinal_ex(self._ctx, buf,
                                                    self._backend.ffi.NULL)
         assert res != 0
-        res = self._backend.lib.EVP_MD_CTX_cleanup(ctx)
+        res = self._backend.lib.EVP_MD_CTX_cleanup(self._ctx)
         assert res == 1
-        return self._backend.ffi.buffer(buf)[:digest_size]
+        return self._backend.ffi.buffer(buf)[:self.algorithm.digest_size]
 
-    def copy_ctx(self, ctx):
-        copied_ctx = self._backend.lib.EVP_MD_CTX_create()
-        copied_ctx = self._backend.ffi.gc(copied_ctx,
-                                          self._backend.lib.EVP_MD_CTX_destroy)
-        res = self._backend.lib.EVP_MD_CTX_copy_ex(copied_ctx, ctx)
-        assert res != 0
-        return copied_ctx
+
+class Hashes(object):
+    def __init__(self, backend):
+        super(Hashes, self).__init__()
+        self._backend = backend
+
+    def supported(self, algorithm):
+        digest = self._backend.lib.EVP_get_digestbyname(
+            algorithm.name.encode("ascii")
+        )
+        return digest != self._backend.ffi.NULL
+
+    def create_ctx(self, algorithm):
+        return _HashContext(self._backend, algorithm)
 
 
 class HMACs(object):
diff --git a/cryptography/hazmat/primitives/hashes.py b/cryptography/hazmat/primitives/hashes.py
index c7748118..c14d0437 100644
--- a/cryptography/hazmat/primitives/hashes.py
+++ b/cryptography/hazmat/primitives/hashes.py
@@ -39,15 +39,14 @@ class Hash(object):
     def update(self, data):
         if isinstance(data, six.text_type):
             raise TypeError("Unicode-objects must be encoded before hashing")
-        self._backend.hashes.update_ctx(self._ctx, data)
+        self._ctx.update(data)
 
     def copy(self):
         return self.__class__(self.algorithm, backend=self._backend,
-                              ctx=self._backend.hashes.copy_ctx(self._ctx))
+                              ctx=self._ctx.copy())
 
     def finalize(self):
-        return self._backend.hashes.finalize_ctx(self._ctx,
-                                                 self.algorithm.digest_size)
+        return self._ctx.finalize()
 
 
 @interfaces.register(interfaces.HashAlgorithm)