Merge pull request #249 from dreid/hmac-context-from-backend

Get a HashContext from the hmac backend like we do a CipherContext

3 files changed, 55 insertions, 34 deletions

diff --git a/cryptography/hazmat/bindings/openssl/backend.py b/cryptography/hazmat/bindings/openssl/backend.py
index 71b94abe..4760f9b7 100644
--- a/cryptography/hazmat/bindings/openssl/backend.py
+++ b/cryptography/hazmat/bindings/openssl/backend.py
@@ -324,43 +324,64 @@ class Hashes(object):
         return _HashContext(self._backend, algorithm)
 
 
-class HMACs(object):
-    def __init__(self, backend):
+@interfaces.register(interfaces.HashContext)
+class _HMACContext(object):
+    def __init__(self, backend, key, algorithm, ctx=None):
+        self.algorithm = algorithm
         self._backend = backend
 
-    def create_ctx(self, key, hash_cls):
-        ctx = self._backend.ffi.new("HMAC_CTX *")
-        self._backend.lib.HMAC_CTX_init(ctx)
-        ctx = self._backend.ffi.gc(ctx, self._backend.lib.HMAC_CTX_cleanup)
-        evp_md = self._backend.lib.EVP_get_digestbyname(
-            hash_cls.name.encode('ascii'))
-        assert evp_md != self._backend.ffi.NULL
-        res = self._backend.lib.Cryptography_HMAC_Init_ex(
-            ctx, key, len(key), evp_md, self._backend.ffi.NULL
+        if ctx is None:
+            ctx = self._backend.ffi.new("HMAC_CTX *")
+            self._backend.lib.HMAC_CTX_init(ctx)
+            ctx = self._backend.ffi.gc(ctx, self._backend.lib.HMAC_CTX_cleanup)
+            evp_md = self._backend.lib.EVP_get_digestbyname(
+                algorithm.name.encode('ascii'))
+            assert evp_md != self._backend.ffi.NULL
+            res = self._backend.lib.Cryptography_HMAC_Init_ex(
+                ctx, key, len(key), evp_md, self._backend.ffi.NULL
+            )
+            assert res != 0
+
+        self._ctx = ctx
+        self._key = key
+
+    def copy(self):
+        copied_ctx = self._backend.ffi.new("HMAC_CTX *")
+        self._backend.lib.HMAC_CTX_init(copied_ctx)
+        copied_ctx = self._backend.ffi.gc(
+            copied_ctx, self._backend.lib.HMAC_CTX_cleanup
+        )
+        res = self._backend.lib.Cryptography_HMAC_CTX_copy(
+            copied_ctx, self._ctx
         )
         assert res != 0
-        return ctx
+        return _HMACContext(
+            self._backend, self._key, self.algorithm, ctx=copied_ctx
+        )
 
-    def update_ctx(self, ctx, data):
-        res = self._backend.lib.Cryptography_HMAC_Update(ctx, data, len(data))
+    def update(self, data):
+        res = self._backend.lib.Cryptography_HMAC_Update(
+            self._ctx, data, len(data)
+        )
         assert res != 0
 
-    def finalize_ctx(self, ctx, digest_size):
-        buf = self._backend.ffi.new("unsigned char[]", digest_size)
-        buflen = self._backend.ffi.new("unsigned int *", digest_size)
-        res = self._backend.lib.Cryptography_HMAC_Final(ctx, buf, buflen)
+    def finalize(self):
+        buf = self._backend.ffi.new("unsigned char[]",
+                                    self.algorithm.digest_size)
+        buflen = self._backend.ffi.new("unsigned int *",
+                                       self.algorithm.digest_size)
+        res = self._backend.lib.Cryptography_HMAC_Final(self._ctx, buf, buflen)
         assert res != 0
-        self._backend.lib.HMAC_CTX_cleanup(ctx)
-        return self._backend.ffi.buffer(buf)[:digest_size]
+        self._backend.lib.HMAC_CTX_cleanup(self._ctx)
+        return self._backend.ffi.buffer(buf)[:self.algorithm.digest_size]
 
-    def copy_ctx(self, ctx):
-        copied_ctx = self._backend.ffi.new("HMAC_CTX *")
-        self._backend.lib.HMAC_CTX_init(copied_ctx)
-        copied_ctx = self._backend.ffi.gc(copied_ctx,
-                                          self._backend.lib.HMAC_CTX_cleanup)
-        res = self._backend.lib.Cryptography_HMAC_CTX_copy(copied_ctx, ctx)
-        assert res != 0
-        return copied_ctx
+
+class HMACs(object):
+    def __init__(self, backend):
+        self._backend = backend
+
+    def create_ctx(self, key, algorithm):
+        return _HMACContext(self._backend, key, algorithm)
 
 
 backend = Backend()
diff --git a/cryptography/hazmat/primitives/hmac.py b/cryptography/hazmat/primitives/hmac.py
index ed2dd54a..0f1b4fac 100644
--- a/cryptography/hazmat/primitives/hmac.py
+++ b/cryptography/hazmat/primitives/hmac.py
@@ -39,16 +39,15 @@ class HMAC(object):
     def update(self, msg):
         if isinstance(msg, six.text_type):
             raise TypeError("Unicode-objects must be encoded before hashing")
-        self._backend.hmacs.update_ctx(self._ctx, msg)
+        self._ctx.update(msg)
 
     def copy(self):
         return HMAC(
             self._key,
             self.algorithm,
             backend=self._backend,
-            ctx=self._backend.hmacs.copy_ctx(self._ctx)
+            ctx=self._ctx.copy()
         )
 
     def finalize(self):
-        return self._backend.hmacs.finalize_ctx(self._ctx,
-                                                self.algorithm.digest_size)
+        return self._ctx.finalize()
diff --git a/tests/hazmat/primitives/test_hmac.py b/tests/hazmat/primitives/test_hmac.py
index a44838cf..0f627a10 100644
--- a/tests/hazmat/primitives/test_hmac.py
+++ b/tests/hazmat/primitives/test_hmac.py
@@ -37,9 +37,10 @@ class TestHMAC(object):
             h.update(six.u("\u00FC"))
 
     def test_copy_backend_object(self):
-        pretend_hmac = pretend.stub(copy_ctx=lambda a: True)
+        pretend_hmac = pretend.stub()
         pretend_backend = pretend.stub(hmacs=pretend_hmac)
-        pretend_ctx = pretend.stub()
+        copied_ctx = pretend.stub()
+        pretend_ctx = pretend.stub(copy=lambda: copied_ctx)
         h = hmac.HMAC(b"key", hashes.SHA1(), backend=pretend_backend,
                       ctx=pretend_ctx)
         assert h._backend is pretend_backend