Better calculation of bit size

1 files changed, 29 insertions, 2 deletions

diff --git a/rsa/common.py b/rsa/common.py
index 7c852b6..ecb42c9 100644
--- a/rsa/common.py
+++ b/rsa/common.py
@@ -8,10 +8,17 @@ def bit_size(number):
     >>> bit_size(1023)
     10
     >>> bit_size(1024)
-    10
+    11
     >>> bit_size(1025)
     11
 
+    >>> bit_size(1 << 1024)
+    1025
+    >>> bit_size((1 << 1024) + 1)
+    1025
+    >>> bit_size((1 << 1024) - 1)
+    1024
+
     '''
 
     if number < 0:
@@ -20,12 +27,32 @@ def bit_size(number):
     if number == 0:
         return 1
     
-    return int(math.ceil(math.log(number, 2)))
+    # This works, even with very large numbers. When using math.log(number, 2),
+    # you'll get rounding errors and it'll fail.
+    bits = 0
+    while number:
+        bits += 1
+        number >>= 1
+
+    return bits
+
 
 def byte_size(number):
     """Returns the number of bytes required to hold a specific long number.
     
     The number of bytes is rounded up.
+
+    >>> byte_size(1 << 1023)
+    128
+    >>> byte_size((1 << 1024) - 1)
+    128
+    >>> byte_size(1 << 1024)
+    129
     """
 
     return int(math.ceil(bit_size(number) / 8.0))
+
+if __name__ == '__main__':
+    import doctest
+    doctest.testmod()
+