randnum.py 2.6 KB

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