randnum.py 2.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495
  1. # Copyright 2011 Sybren A. Stüvel <sybren@stuvel.eu>
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # https://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. """Functions for generating random numbers."""
  15. # Source inspired by code by Yesudeep Mangalapilly <yesudeep@gmail.com>
  16. import os
  17. import struct
  18. from rsa import common, transform
  19. def read_random_bits(nbits: int) -> bytes:
  20. """Reads 'nbits' random bits.
  21. If nbits isn't a whole number of bytes, an extra byte will be appended with
  22. only the lower bits set.
  23. """
  24. nbytes, rbits = divmod(nbits, 8)
  25. # Get the random bytes
  26. randomdata = os.urandom(nbytes)
  27. # Add the remaining random bits
  28. if rbits > 0:
  29. randomvalue = ord(os.urandom(1))
  30. randomvalue >>= 8 - rbits
  31. randomdata = struct.pack("B", randomvalue) + randomdata
  32. return randomdata
  33. def read_random_int(nbits: int) -> int:
  34. """Reads a random integer of approximately nbits bits."""
  35. randomdata = read_random_bits(nbits)
  36. value = transform.bytes2int(randomdata)
  37. # Ensure that the number is large enough to just fill out the required
  38. # number of bits.
  39. value |= 1 << (nbits - 1)
  40. return value
  41. def read_random_odd_int(nbits: int) -> int:
  42. """Reads a random odd integer of approximately nbits bits.
  43. >>> read_random_odd_int(512) & 1
  44. 1
  45. """
  46. value = read_random_int(nbits)
  47. # Make sure it's odd
  48. return value | 1
  49. def randint(maxvalue: int) -> int:
  50. """Returns a random integer x with 1 <= x <= maxvalue
  51. May take a very long time in specific situations. If maxvalue needs N bits
  52. to store, the closer maxvalue is to (2 ** N) - 1, the faster this function
  53. is.
  54. """
  55. bit_size = common.bit_size(maxvalue)
  56. tries = 0
  57. while True:
  58. value = read_random_int(bit_size)
  59. if value <= maxvalue:
  60. break
  61. if tries % 10 == 0 and tries:
  62. # After a lot of tries to get the right number of bits but still
  63. # smaller than maxvalue, decrease the number of bits by 1. That'll
  64. # dramatically increase the chances to get a large enough number.
  65. bit_size -= 1
  66. tries += 1
  67. return value