codekingpro/portable-devtools

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
15"""Functions for generating random numbers."""
16
17# Source inspired by code by Yesudeep Mangalapilly <yesudeep@gmail.com>
18
19import os
20import struct
21
22from rsa import common, transform
23
24
25def read_random_bits(nbits: int) -> bytes:
26    """Reads 'nbits' random bits.
27
28    If nbits isn't a whole number of bytes, an extra byte will be appended with
29    only the lower bits set.
30    """
31
32    nbytes, rbits = divmod(nbits, 8)
33
34    # Get the random bytes
35    randomdata = os.urandom(nbytes)
36
37    # Add the remaining random bits
38    if rbits > 0:
39        randomvalue = ord(os.urandom(1))
40        randomvalue >>= 8 - rbits
41        randomdata = struct.pack("B", randomvalue) + randomdata
42
43    return randomdata
44
45
46def read_random_int(nbits: int) -> int:
47    """Reads a random integer of approximately nbits bits."""
48
49    randomdata = read_random_bits(nbits)
50    value = transform.bytes2int(randomdata)
51
52    # Ensure that the number is large enough to just fill out the required
53    # number of bits.
54    value |= 1 << (nbits - 1)
55
56    return value
57
58
59def read_random_odd_int(nbits: int) -> int:
60    """Reads a random odd integer of approximately nbits bits.
61
62    >>> read_random_odd_int(512) & 1
63    1
64    """
65
66    value = read_random_int(nbits)
67
68    # Make sure it's odd
69    return value | 1
70
71
72def randint(maxvalue: int) -> int:
73    """Returns a random integer x with 1 <= x <= maxvalue
74
75    May take a very long time in specific situations. If maxvalue needs N bits
76    to store, the closer maxvalue is to (2 ** N) - 1, the faster this function
77    is.
78    """
79
80    bit_size = common.bit_size(maxvalue)
81
82    tries = 0
83    while True:
84        value = read_random_int(bit_size)
85        if value <= maxvalue:
86            break
87
88        if tries % 10 == 0 and tries:
89            # After a lot of tries to get the right number of bits but still
90            # smaller than maxvalue, decrease the number of bits by 1. That'll
91            # dramatically increase the chances to get a large enough number.
92            bit_size -= 1
93        tries += 1
94
95    return value
96