# Pseudorandom number generator with key erasure / forward secrecy /
# backtracking resistance using ChaCha20.

# Copyright (c) 2015, 2016 Taylor R. Campbell
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.

import chacha
import os
import struct
import threading

class PRNG(object):
    _const32 = chacha.const32
    _zero = [0,0,0,0, 0,0,0,0]

    def __init__(self, seed):
        self._key = list(struct.unpack('<IIIIIIII', seed)) + [0,0,0,0, 0,0,0,0]

    def random32(self):
        key = self._key
        chacha.core(20, key, self._zero, key, self._const32)
        r = key[8]; key[8] = 0
        return r

    def random64(self):
        key = self._key
        chacha.core(20, key, self._zero, key, self._const32)
        rl = key[8]; key[8] = 0
        rh = key[9]; key[9] = 0
        return (rh << 32) | rl

    def random_uniform(self, n):
        assert 0 < n
        nbits = n.bit_length()
        nwords = (nbits + 31)//32
        l = ((1 << nbits) - n) % n
        while True:
            r = 0
            for i in range(nwords):
                r <<= 32
                # Suboptimal, but you're unlikely to ask for >1 word.
                r |= self.random32()
            if r < l:
                continue
            return (r % n)

    def random_bytearray(self, buf, start, end):
        assert end <= len(buf)
        assert start <= end
        key = self._key
        zero = self._zero
        const32 = self._const32
        nbytes = end - start
        nblocks = nbytes//32
        if nblocks < 2:
            for i in range(nblocks):
                chacha.core(20, key, zero, key, const32)
                buf[start + 32*i : start + 32*(i + 1)] = \
                    struct.pack('<IIIIIIII', *key[8:16])
                key[8:16] = zero
            nextra = nbytes - 32*nblocks
            if 0 < nextra:
                chacha.core(20, key, zero, key, const32)
                blkstart = start + 32*nblocks
                buf[blkstart : blkstart + nextra] = \
                    struct.pack('<IIIIIIII', *key[8:16])[0 : nextra]
                key[8:16] = zero
        else:
            chacha.core(20, key, self._zero, key, const32)
            subkey = key[8:16]; key[8:16] = zero
            ctr = [0,0,0,0]
            out = [0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0]
            for i in range(nblocks//2):
                chacha.core(20, out, ctr, subkey, const32)
                # Overflows after 2^64 blocks.  Not worth carrying.
                t = ctr[0] + 1; ctr[0] = t & 0xffffffff; ctr[1] += t >> 32
                buf[start + 64*i : start + 64*(i + 1)] = \
                    struct.pack('<IIIIIIIIIIIIIIII', *out)
            nextra = nbytes - 64*(nblocks//2)
            if 0 < nextra:
                chacha.core(20, out, ctr, subkey, const32)
                blkstart = start + 64*(nblocks//2)
                buf[blkstart : blkstart + nextra] \
                    = struct.pack('<IIIIIIIIIIIIIIII', *out)[0 : nextra]
            subkey[0:8] = out[0:8] = out[8:16] = zero

    def random_bytes(self, n):
        buf = bytearray(n)
        self.random_bytearray(buf, 0, n)
        return bytes(buf)

def prng(seed):
    return PRNG(seed)

if prng(bytes(bytearray([0] * 32))).random64() != 0x8d4857517c5941da:
    raise Exception('prng self-test failed')

nondeterministic_prng_local = threading.local()
nondeterministic_prng_local.prng = None

def nondeterministic_prng():
    if nondeterministic_prng_local.prng == None:
        nondeterministic_prng_local.prng = prng(os.urandom(32))
    return nondeterministic_prng_local.prng

def nondeterministic_random32():
    return nondeterministic_prng().random32()
def nondeterministic_random64():
    return nondeterministic_prng().random64()
def nondeterministic_random_uniform(n):
    return nondeterministic_prng().random_uniform(n)
def nondeterministic_random_bytearray(buf, start, end):
    return nondeterministic_prng().random_bytearray(buf, start, end)
def nondeterministic_random_bytes(n):
    return nondeterministic_prng().random_bytes(n)

nondeterministic_random32()
nondeterministic_random64()
nondeterministic_random_uniform(123)
nondeterministic_random_bytearray(bytearray([0] * 16), 3, 8)
nondeterministic_random_bytearray(bytearray([0] * 80), 3, 67)
nondeterministic_random_bytearray(bytearray([0] * 160), 3, 132)
nondeterministic_random_bytearray(bytearray([0] * 1024), 3, 1023)
nondeterministic_random_bytearray(bytearray([0] * 1024), 3, 963)
nondeterministic_random_bytes(42)