random.py 5.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139
  1. # -*- coding: utf-8 -*-
  2. #
  3. # Random/random.py : Strong alternative for the standard 'random' module
  4. #
  5. # Written in 2008 by Dwayne C. Litzenberger <dlitz@dlitz.net>
  6. #
  7. # ===================================================================
  8. # The contents of this file are dedicated to the public domain. To
  9. # the extent that dedication to the public domain is not available,
  10. # everyone is granted a worldwide, perpetual, royalty-free,
  11. # non-exclusive license to exercise all rights associated with the
  12. # contents of this file for any purpose whatsoever.
  13. # No rights are reserved.
  14. #
  15. # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  16. # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
  17. # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  18. # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
  19. # BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
  20. # ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
  21. # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  22. # SOFTWARE.
  23. # ===================================================================
  24. __all__ = ['StrongRandom', 'getrandbits', 'randrange', 'randint', 'choice', 'shuffle', 'sample']
  25. from Cryptodome import Random
  26. from Cryptodome.Util.py3compat import is_native_int
  27. class StrongRandom(object):
  28. def __init__(self, rng=None, randfunc=None):
  29. if randfunc is None and rng is None:
  30. self._randfunc = None
  31. elif randfunc is not None and rng is None:
  32. self._randfunc = randfunc
  33. elif randfunc is None and rng is not None:
  34. self._randfunc = rng.read
  35. else:
  36. raise ValueError("Cannot specify both 'rng' and 'randfunc'")
  37. def getrandbits(self, k):
  38. """Return an integer with k random bits."""
  39. if self._randfunc is None:
  40. self._randfunc = Random.new().read
  41. mask = (1 << k) - 1
  42. return mask & bytes_to_long(self._randfunc(ceil_div(k, 8)))
  43. def randrange(self, *args):
  44. """randrange([start,] stop[, step]):
  45. Return a randomly-selected element from range(start, stop, step)."""
  46. if len(args) == 3:
  47. (start, stop, step) = args
  48. elif len(args) == 2:
  49. (start, stop) = args
  50. step = 1
  51. elif len(args) == 1:
  52. (stop,) = args
  53. start = 0
  54. step = 1
  55. else:
  56. raise TypeError("randrange expected at most 3 arguments, got %d" % (len(args),))
  57. if (not is_native_int(start) or not is_native_int(stop) or not
  58. is_native_int(step)):
  59. raise TypeError("randrange requires integer arguments")
  60. if step == 0:
  61. raise ValueError("randrange step argument must not be zero")
  62. num_choices = ceil_div(stop - start, step)
  63. if num_choices < 0:
  64. num_choices = 0
  65. if num_choices < 1:
  66. raise ValueError("empty range for randrange(%r, %r, %r)" % (start, stop, step))
  67. # Pick a random number in the range of possible numbers
  68. r = num_choices
  69. while r >= num_choices:
  70. r = self.getrandbits(size(num_choices))
  71. return start + (step * r)
  72. def randint(self, a, b):
  73. """Return a random integer N such that a <= N <= b."""
  74. if not is_native_int(a) or not is_native_int(b):
  75. raise TypeError("randint requires integer arguments")
  76. N = self.randrange(a, b+1)
  77. assert a <= N <= b
  78. return N
  79. def choice(self, seq):
  80. """Return a random element from a (non-empty) sequence.
  81. If the seqence is empty, raises IndexError.
  82. """
  83. if len(seq) == 0:
  84. raise IndexError("empty sequence")
  85. return seq[self.randrange(len(seq))]
  86. def shuffle(self, x):
  87. """Shuffle the sequence in place."""
  88. # Fisher-Yates shuffle. O(n)
  89. # See http://en.wikipedia.org/wiki/Fisher-Yates_shuffle
  90. # Working backwards from the end of the array, we choose a random item
  91. # from the remaining items until all items have been chosen.
  92. for i in range(len(x)-1, 0, -1): # iterate from len(x)-1 downto 1
  93. j = self.randrange(0, i+1) # choose random j such that 0 <= j <= i
  94. x[i], x[j] = x[j], x[i] # exchange x[i] and x[j]
  95. def sample(self, population, k):
  96. """Return a k-length list of unique elements chosen from the population sequence."""
  97. num_choices = len(population)
  98. if k > num_choices:
  99. raise ValueError("sample larger than population")
  100. retval = []
  101. selected = {} # we emulate a set using a dict here
  102. for i in range(k):
  103. r = None
  104. while r is None or r in selected:
  105. r = self.randrange(num_choices)
  106. retval.append(population[r])
  107. selected[r] = 1
  108. return retval
  109. _r = StrongRandom()
  110. getrandbits = _r.getrandbits
  111. randrange = _r.randrange
  112. randint = _r.randint
  113. choice = _r.choice
  114. shuffle = _r.shuffle
  115. sample = _r.sample
  116. # These are at the bottom to avoid problems with recursive imports
  117. from Cryptodome.Util.number import ceil_div, bytes_to_long, long_to_bytes, size
  118. # vim:set ts=4 sw=4 sts=4 expandtab: