use native methods for number to bytes conversion

Author: tomato42

tlslite/utils/compat.py

@@ -16,12 +16,21 @@
 if sys.version_info >= (3,0):
 
     def compat26Str(x): return x
-    
-    # Python 3 requires bytes instead of bytearrays for HMAC   
-    
+
+    # Python 3.3 requires bytes instead of bytearrays for HMAC
     # So, python 2.6 requires strings, python 3 requires 'bytes',
-    # and python 2.7 can handle bytearrays...     
-    def compatHMAC(x): return bytes(x)
+    # and python 2.7 and 3.5 can handle bytearrays...
+    # pylint: disable=invalid-name
+    # we need to keep compatHMAC and `x` for API compatibility
+    if sys.version_info < (3, 4):
+        def compatHMAC(x):
+            """Convert bytes-like input to format acceptable for HMAC."""
+            return bytes(x)
+    else:
+        def compatHMAC(x):
+            """Convert bytes-like input to format acceptable for HMAC."""
+            return x
+    # pylint: enable=invalid-name
 
     def compatAscii2Bytes(val):
         """Convert ASCII string to bytes."""
@@ -80,7 +89,24 @@ def remove_whitespace(text):
         """Removes all whitespace from passed in string"""
         return re.sub(r"\s+", "", text, flags=re.UNICODE)
 
+    # pylint: disable=invalid-name
+    # pylint is stupid here and deson't notice it's a function, not
+    # constant
     bytes_to_int = int.from_bytes
+    # pylint: enable=invalid-name
+
+    def bit_length(val):
+        """Return number of bits necessary to represent an integer."""
+        return val.bit_length()
+
+    def int_to_bytes(val, length=None, byteorder="big"):
+        """Return number converted to bytes"""
+        if length is None:
+            length = byte_length(val)
+        # for gmpy we need to convert back to native int
+        if type(val) != int:
+            val = int(val)
+        return bytearray(val.to_bytes(length=length, byteorder=byteorder))
 
 else:
     # Python 2.6 requires strings instead of bytearrays in a couple places,
@@ -94,13 +120,23 @@ def compat26Str(x): return str(x)
         def remove_whitespace(text):
             """Removes all whitespace from passed in string"""
             return re.sub(r"\s+", "", text)
+
+        def bit_length(val):
+            """Return number of bits necessary to represent an integer."""
+            if val == 0:
+                return 0
+            return len(bin(val))-2
     else:
         def compat26Str(x): return x
 
         def remove_whitespace(text):
             """Removes all whitespace from passed in string"""
             return re.sub(r"\s+", "", text, flags=re.UNICODE)
 
+        def bit_length(val):
+            """Return number of bits necessary to represent an integer."""
+            return val.bit_length()
+
     def compatAscii2Bytes(val):
         """Convert ASCII string to bytes."""
         return val
@@ -159,6 +195,25 @@ def bytes_to_int(val, byteorder):
             return int(b2a_hex(val[::-1]), 16)
         raise ValueError("Only 'big' and 'little' endian supported")
 
+    def int_to_bytes(val, length=None, byteorder="big"):
+        """Return number converted to bytes"""
+        if length is None:
+            length = byte_length(val)
+        if byteorder == "big":
+            return bytearray((val >> i) & 0xff
+                             for i in reversed(range(0, length*8, 8)))
+        if byteorder == "little":
+            return bytearray((val >> i) & 0xff
+                             for i in range(0, length*8, 8))
+        raise ValueError("Only 'big' or 'little' endian supported")
+
+
+def byte_length(val):
+    """Return number of bytes necessary to represent an integer."""
+    length = bit_length(val)
+    return (length + 7) // 8
+
+
 try:
     # Fedora and Red Hat Enterprise Linux versions have small curves removed
     getattr(ecdsa, 'NIST192p')

tlslite/utils/cryptomath.py

@@ -13,9 +13,9 @@
 import math
 import base64
 import binascii
-import sys
 
-from .compat import compat26Str, compatHMAC, compatLong, bytes_to_int
+from .compat import compat26Str, compatHMAC, compatLong, \
+        bytes_to_int, int_to_bytes, bit_length, byte_length
 from .codec import Writer
 
 from . import tlshashlib as hashlib
@@ -215,16 +215,14 @@ def numberToByteArray(n, howManyBytes=None, endian="big"):
     not be larger.  The returned bytearray will contain a big- or little-endian
     encoding of the input integer (n). Big endian encoding is used by default.
     """
-    if howManyBytes == None:
-        howManyBytes = numBytes(n)
-    if endian == "big":
-        return bytearray((n >> i) & 0xff
-                         for i in reversed(range(0, howManyBytes*8, 8)))
-    elif endian == "little":
-        return bytearray((n >> i) & 0xff
-                         for i in range(0, howManyBytes*8, 8))
-    else:
-        raise ValueError("Only 'big' and 'little' endian supported")
+    if howManyBytes is not None:
+        length = byte_length(n)
+        if howManyBytes < length:
+            ret = int_to_bytes(n, length, endian)
+            if endian == "big":
+                return ret[length-howManyBytes:length]
+            return ret[:howManyBytes]
+    return int_to_bytes(n, howManyBytes, endian)
 
 
 def mpiToNumber(mpi):
@@ -255,23 +253,16 @@ def numberToMPI(n):
 # Misc. Utility Functions
 # **************************************************************************
 
-def numBits(n):
-    """Return number of bits necessary to represent the integer in binary"""
-    if n==0:
-        return 0
-    if sys.version_info < (2, 7):
-        # bit_length() was introduced in 2.7, and it is an order of magnitude
-        # faster than the below code
-        return len(bin(n))-2
-    else:
-        return n.bit_length()
 
-def numBytes(n):
-    """Return number of bytes necessary to represent the integer in bytes"""
-    if n==0:
-        return 0
-    bits = numBits(n)
-    return (bits + 7) // 8
+# pylint: disable=invalid-name
+# pylint recognises them as constants, not function names, also
+# we can't change their names without API change
+numBits = bit_length
+
+
+numBytes = byte_length
+# pylint: enable=invalid-name
+
 
 # **************************************************************************
 # Big Number Math