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
35import hashlib
36
37
38class HashContainer(object):
39    def __init__(self, hash_constructor):
40        self.new = hash_constructor
41        self.digest_size = hash_constructor().digest_size
42
43
44sha1_module = HashContainer(hashlib.sha1)
45sha256_module = HashContainer(hashlib.sha256)
46
47
48def hmacSha1(key, text):
49    if isinstance(key, str):
50        key = bytes(key, encoding="utf-8")
51    if isinstance(text, str):
52        text = bytes(text, encoding="utf-8")
53    return hmac.new(key, text, sha1_module).digest()
54
55
56def sha1(s):
57    if isinstance(s, str):
58        s = bytes(s, encoding="utf-8")
59    return sha1_module.new(s).digest()
60
61
62def hmacSha256(key, text):
63    if isinstance(key, str):
64        key = bytes(key, encoding="utf-8")
65    if isinstance(text, str):
66        text = bytes(text, encoding="utf-8")
67    return hmac.new(key, text, sha256_module).digest()
68
69
70def sha256(s):
71    if isinstance(s, str):
72        s = bytes(s, encoding="utf-8")
73    return sha256_module.new(s).digest()
74
75
76SHA256_AVAILABLE = True
77
78try:
79    from Crypto.Util.number import long_to_bytes, bytes_to_long
80except ImportError:
81    # In the case where we don't have pycrypto installed, define substitute
82    # functionality.
83
84    import pickle
85
86    def longToBinary(l):
87        if l == 0:
88            return b'\x00'
89        b = bytearray(pickle.encode_long(l))
90        b.reverse()
91        return bytes(b)
92
93    def binaryToLong(s):
94        if isinstance(s, str):
95            s = s.encode("utf-8")
96        b = bytearray(s)
97        b.reverse()
98        return pickle.decode_long(bytes(b))
99else:
100    # We have pycrypto, so wrap its functions instead.
101
102    def longToBinary(l):
103        if l < 0:
104            raise ValueError('This function only supports positive integers')
105
106        bytestring = long_to_bytes(l)
107        if bytestring[0] > 127:
108            return b'\x00' + bytestring
109        else:
110            return bytestring
111
112    def binaryToLong(bytestring):
113        if not bytestring:
114            raise ValueError('Empty string passed to strToLong')
115
116        if bytestring[0] > 127:
117            raise ValueError('This function only supports positive integers')
118
119        return bytes_to_long(bytestring)
120
121
122# A cryptographically safe source of random bytes
123getBytes = os.urandom
124
125# A randrange function that works for longs
126randrange = random.randrange
127
128
129def longToBase64(l):
130    return toBase64(longToBinary(l))
131
132
133def base64ToLong(s):
134    return binaryToLong(fromBase64(s))
135
136
137def randomString(length, chrs=None):
138    """Produce a string of length random bytes, chosen from chrs."""
139    if chrs is None:
140        return getBytes(length)
141    else:
142        n = len(chrs)
143        return ''.join([chrs[randrange(n)] for _ in range(length)])
144
145
146def const_eq(s1, s2):
147    if len(s1) != len(s2):
148        return False
149
150    result = True
151    for i in range(len(s1)):
152        result = result and (s1[i] == s2[i])
153
154    return result
155