1"""Module containing a cryptographic-quality source of randomness and
2other cryptographically useful functionality
3
4Python 2.4 needs no external support for this module, nor does Python
52.3 on a system with /dev/urandom.
6
7Other configurations will need a quality source of random bytes and
8access to a function that will convert binary strings to long
9integers. This module will work with the Python Cryptography Toolkit
10(pycrypto) if it is present. pycrypto can be found with a search
11engine, but is currently found at:
12
13http://www.amk.ca/python/code/crypto
14"""
15
16__all__ = [
17    'base64ToLong',
18    'binaryToLong',
19    'hmacSha1',
20    'hmacSha256',
21    'longToBase64',
22    'longToBinary',
23    'randomString',
24    'randrange',
25    'sha1',
26    'sha256',
27    ]
28
29import hmac
30import os
31import random
32
33from openid.oidutil import toBase64, fromBase64
34
35try:
36    import hashlib
37except ImportError:
38    import sha as sha1_module
39
40    try:
41        from Crypto.Hash import SHA256 as sha256_module
42    except ImportError:
43        sha256_module = None
44
45else:
46    class HashContainer(object):
47        def __init__(self, hash_constructor):
48            self.new = hash_constructor
49            self.digest_size = hash_constructor().digest_size
50
51    sha1_module = HashContainer(hashlib.sha1)
52    sha256_module = HashContainer(hashlib.sha256)
53
54def hmacSha1(key, text):
55    return hmac.new(key, text, sha1_module).digest()
56
57def sha1(s):
58    return sha1_module.new(s).digest()
59
60if sha256_module is not None:
61    def hmacSha256(key, text):
62        return hmac.new(key, text, sha256_module).digest()
63
64    def sha256(s):
65        return sha256_module.new(s).digest()
66
67    SHA256_AVAILABLE = True
68
69else:
70    _no_sha256 = NotImplementedError(
71        'Use Python 2.5, install pycrypto or install hashlib to use SHA256')
72
73    def hmacSha256(unused_key, unused_text):
74        raise _no_sha256
75
76    def sha256(s):
77        raise _no_sha256
78
79    SHA256_AVAILABLE = False
80
81try:
82    from Crypto.Util.number import long_to_bytes, bytes_to_long
83except ImportError:
84    import pickle
85    try:
86        # Check Python compatiblity by raising an exception on import
87        # if the needed functionality is not present. Present in
88        # Python >= 2.3
89        pickle.encode_long
90        pickle.decode_long
91    except AttributeError:
92        raise ImportError(
93            'No functionality for serializing long integers found')
94
95    # Present in Python >= 2.4
96    try:
97        reversed
98    except NameError:
99        def reversed(seq):
100            return map(seq.__getitem__, xrange(len(seq) - 1, -1, -1))
101
102    def longToBinary(l):
103        if l == 0:
104            return '\x00'
105
106        return ''.join(reversed(pickle.encode_long(l)))
107
108    def binaryToLong(s):
109        return pickle.decode_long(''.join(reversed(s)))
110else:
111    # We have pycrypto
112
113    def longToBinary(l):
114        if l < 0:
115            raise ValueError('This function only supports positive integers')
116
117        bytes = long_to_bytes(l)
118        if ord(bytes[0]) > 127:
119            return '\x00' + bytes
120        else:
121            return bytes
122
123    def binaryToLong(bytes):
124        if not bytes:
125            raise ValueError('Empty string passed to strToLong')
126
127        if ord(bytes[0]) > 127:
128            raise ValueError('This function only supports positive integers')
129
130        return bytes_to_long(bytes)
131
132# A cryptographically safe source of random bytes
133try:
134    getBytes = os.urandom
135except AttributeError:
136    try:
137        from Crypto.Util.randpool import RandomPool
138    except ImportError:
139        # Fall back on /dev/urandom, if present. It would be nice to
140        # have Windows equivalent here, but for now, require pycrypto
141        # on Windows.
142        try:
143            _urandom = file('/dev/urandom', 'rb')
144        except IOError:
145            raise ImportError('No adequate source of randomness found!')
146        else:
147            def getBytes(n):
148                bytes = []
149                while n:
150                    chunk = _urandom.read(n)
151                    n -= len(chunk)
152                    bytes.append(chunk)
153                    assert n >= 0
154                return ''.join(bytes)
155    else:
156        _pool = RandomPool()
157        def getBytes(n, pool=_pool):
158            if pool.entropy < n:
159                pool.randomize()
160            return pool.get_bytes(n)
161
162# A randrange function that works for longs
163try:
164    randrange = random.SystemRandom().randrange
165except AttributeError:
166    # In Python 2.2's random.Random, randrange does not support
167    # numbers larger than sys.maxint for randrange. For simplicity,
168    # use this implementation for any Python that does not have
169    # random.SystemRandom
170    from math import log, ceil
171
172    _duplicate_cache = {}
173    def randrange(start, stop=None, step=1):
174        if stop is None:
175            stop = start
176            start = 0
177
178        r = (stop - start) // step
179        try:
180            (duplicate, nbytes) = _duplicate_cache[r]
181        except KeyError:
182            rbytes = longToBinary(r)
183            if rbytes[0] == '\x00':
184                nbytes = len(rbytes) - 1
185            else:
186                nbytes = len(rbytes)
187
188            mxrand = (256 ** nbytes)
189
190            # If we get a number less than this, then it is in the
191            # duplicated range.
192            duplicate = mxrand % r
193
194            if len(_duplicate_cache) > 10:
195                _duplicate_cache.clear()
196
197            _duplicate_cache[r] = (duplicate, nbytes)
198
199        while 1:
200            bytes = '\x00' + getBytes(nbytes)
201            n = binaryToLong(bytes)
202            # Keep looping if this value is in the low duplicated range
203            if n >= duplicate:
204                break
205
206        return start + (n % r) * step
207
208def longToBase64(l):
209    return toBase64(longToBinary(l))
210
211def base64ToLong(s):
212    return binaryToLong(fromBase64(s))
213
214def randomString(length, chrs=None):
215    """Produce a string of length random bytes, chosen from chrs."""
216    if chrs is None:
217        return getBytes(length)
218    else:
219        n = len(chrs)
220        return ''.join([chrs[randrange(n)] for _ in xrange(length)])
221