use built-in HMAC, if it allows MD5

Author: tomato42

tlslite/utils/tlshmac.py

@@ -17,64 +17,72 @@
 except ImportError:
     __all__ = ["new", "HMAC"]
 
+try:
+    from hmac import HMAC, new
+    # if we can calculate HMAC on MD5, then use the built-in HMAC
+    # implementation
+    _val = HMAC(b'some key', b'msg', 'md5')
+    _val.digest()
+    del _val
+except Exception:
+    # fallback only when MD5 doesn't work
+    class HMAC(object):
+        """Hacked version of HMAC that works in FIPS mode even with MD5."""
 
-class HMAC(object):
-    """Hacked version of HMAC that works in FIPS mode even with MD5."""
+        def __init__(self, key, msg=None, digestmod=None):
+            """
+            Initialise the HMAC and hash first portion of data.
 
-    def __init__(self, key, msg=None, digestmod=None):
-        """
-        Initialise the HMAC and hash first portion of data.
+            msg: data to hash
+            digestmod: name of hash or object that be used as a hash and be cloned
+            """
+            self.key = key
+            if digestmod is None:
+                digestmod = 'md5'
+            if callable(digestmod):
+                digestmod = digestmod()
+            if not hasattr(digestmod, 'digest_size'):
+                digestmod = tlshashlib.new(digestmod)
+            self.block_size = digestmod.block_size
+            self.digest_size = digestmod.digest_size
+            self.digestmod = digestmod
+            if len(key) > self.block_size:
+                k_hash = digestmod.copy()
+                k_hash.update(compatHMAC(key))
+                key = k_hash.digest()
+            if len(key) < self.block_size:
+                key = key + b'\x00' * (self.block_size - len(key))
+            key = bytearray(key)
+            ipad = bytearray(b'\x36' * self.block_size)
+            opad = bytearray(b'\x5c' * self.block_size)
+            i_key = bytearray(i ^ j for i, j in zip(key, ipad))
+            self._o_key = bytearray(i ^ j for i, j in zip(key, opad))
+            self._context = digestmod.copy()
+            self._context.update(compatHMAC(i_key))
+            if msg:
+                self._context.update(compatHMAC(msg))
 
-        msg: data to hash
-        digestmod: name of hash or object that be used as a hash and be cloned
-        """
-        self.key = key
-        if digestmod is None:
-            digestmod = 'md5'
-        if callable(digestmod):
-            digestmod = digestmod()
-        if not hasattr(digestmod, 'digest_size'):
-            digestmod = tlshashlib.new(digestmod)
-        self.block_size = digestmod.block_size
-        self.digest_size = digestmod.digest_size
-        self.digestmod = digestmod
-        if len(key) > self.block_size:
-            k_hash = digestmod.copy()
-            k_hash.update(compatHMAC(key))
-            key = k_hash.digest()
-        if len(key) < self.block_size:
-            key = key + b'\x00' * (self.block_size - len(key))
-        key = bytearray(key)
-        ipad = bytearray(b'\x36' * self.block_size)
-        opad = bytearray(b'\x5c' * self.block_size)
-        i_key = bytearray(i ^ j for i, j in zip(key, ipad))
-        self._o_key = bytearray(i ^ j for i, j in zip(key, opad))
-        self._context = digestmod.copy()
-        self._context.update(compatHMAC(i_key))
-        if msg:
+        def update(self, msg):
             self._context.update(compatHMAC(msg))
 
-    def update(self, msg):
-        self._context.update(compatHMAC(msg))
-
-    def digest(self):
-        i_digest = self._context.digest()
-        o_hash = self.digestmod.copy()
-        o_hash.update(compatHMAC(self._o_key))
-        o_hash.update(compatHMAC(i_digest))
-        return o_hash.digest()
+        def digest(self):
+            i_digest = self._context.digest()
+            o_hash = self.digestmod.copy()
+            o_hash.update(compatHMAC(self._o_key))
+            o_hash.update(compatHMAC(i_digest))
+            return o_hash.digest()
 
-    def copy(self):
-        new = HMAC.__new__(HMAC)
-        new.key = self.key
-        new.digestmod = self.digestmod
-        new.block_size = self.block_size
-        new.digest_size = self.digest_size
-        new._o_key = self._o_key
-        new._context = self._context.copy()
-        return new
+        def copy(self):
+            new = HMAC.__new__(HMAC)
+            new.key = self.key
+            new.digestmod = self.digestmod
+            new.block_size = self.block_size
+            new.digest_size = self.digest_size
+            new._o_key = self._o_key
+            new._context = self._context.copy()
+            return new
 
 
-def new(*args, **kwargs):
-    """General constructor that works in FIPS mode."""
-    return HMAC(*args, **kwargs)
+    def new(*args, **kwargs):
+        """General constructor that works in FIPS mode."""
+        return HMAC(*args, **kwargs)