#! /usr/bin/python
"""
**Thread Scheduling Entropy Pseudo Random Number Generator and
Random Pool CSPRNGs**
"""

import time
from aes.aes import aes
from sha256.sha256 import sha256
from cPickle import *
from threading import Thread, Event, Lock
from whrandom import whrandom
import sys


class threadgate:
    """
    The thread gatekeeper.

    Accumulates ready threads until all can be raced at once.
    Creates intentional race conditions and taps the resulting
    entropy.

    Internal Use only.
    """

    def __init__(self, threadamount):
        self.counter = self.threadamount = threadamount
        self.lock = Lock()
        self.readyEvent = Event()
        self.goEvent = Event()
        self.finishedEvent = Event()

    def ready(self):
        self.lock.acquire()
        self.counter-=1
        if self.counter == 0:
            self.readyEvent.set()
            self.counter = self.threadamount
        self.lock.release()
        self.goEvent.wait()

    def untilReady(self):
        self.readyEvent.wait()

    def go(self):
        self.goEvent.set()

    def untilDone(self):
        self.finishedEvent.wait()

    def finished(self):
        self.lock.acquire()
        self.counter-=1
        if self.counter == 0:
            self.finishedEvent.set()
        self.lock.release()

def rAppend(list, char, racecontrol):
    """
    Appends the thread's identity (the character assigned to it) into a shared list
    """
    racecontrol.ready() # blocks until master thread calls go
    list.append(char)
    racecontrol.finished()

def CreateList(emptylist, listsize, racecontrol, random):
    """
    Internal use only. Don't call this.
    """
    childgate = threadgate(listsize)
    for i in range(listsize):
        thread = Thread(target=rAppend, args=(emptylist, chr(random.randint(0, 255)), childgate))
        thread.start()

    childgate.untilReady()
    racecontrol.ready()
    childgate.go()
    childgate.untilDone()
    racecontrol.finished()

def CreateBuffers(bufferamount, buffersize, random):
    """
    Internal use only. Don't call this.
    """
    creategate = threadgate(bufferamount)
    lists = []
    for i in range(bufferamount):
        lists.append([])
        thread = Thread(target=CreateList, args=(lists[i], buffersize, creategate, random))
        thread.start()
    creategate.untilReady()
    creategate.go()
    creategate.untilDone()
    return lists

def ShuffleMix(list1, list2, racecontrol, random):
    """
    Internal use only. Don't call this.
    """
    childgate = threadgate(len(list1)+len(list2))
    for i in range(len(list1)+len(list2)):
        if i%2:
            list, dest = list1, list2
        else:
            list, dest = list2, list1
        char = list.pop(random.randint(0, len(list)-1)) # again
        thread = Thread(target=rAppend, args=(dest, char, childgate))
        thread.start()
    childgate.untilReady()
    racecontrol.ready()
    childgate.go()
    childgate.untilDone()
    racecontrol.finished()

def ShuffleBuffers(lists, random):
    """
    Internal use only. Don't call this.
    """
    shufflegate = threadgate(len(lists)/2)
    indexes = range(len(lists))
    for i in range(len(lists)/2):    # steps of two
        list1 = lists[indexes.pop(random.randint(0, len(indexes)-1))]   # pseudo random list
        list2 = lists[indexes.pop(random.randint(0, len(indexes)-1))]   # pseudo random list
        thread = Thread(target=ShuffleMix, args=(list1, list2, shufflegate, random))
        thread.start()
    shufflegate.untilReady()
    shufflegate.go()
    shufflegate.untilDone()

def Taint(list, index, racecontrol, random):
    """
    The race condition is the call to the shared random number generator.

    Very processor intensive. Very entropic.

    Internal use only. Don't call this.
    """
    racecontrol.ready() # blocks until master thread calls go
    # Trying to guess the distribution of random numbers
    # is rendered much more difficult with the while loop
    while random.randint(0, 1):
        list[index] = chr(ord(list[index])^random.randint(0, 255))

    racecontrol.finished()

def TaintList(list, racecontrol, random):
    """
    Internal use only. Don't call this.
    """
    childgate = threadgate(len(list))
    indexes = range(len(list))
    for i in range(len(list)):
        index = indexes.pop(random.randint(0, len(indexes)-1))
        thread = Thread(target=Taint, args=(list, index, childgate, random))
        thread.start()
    childgate.untilReady()
    racecontrol.ready()
    childgate.go()
    childgate.untilDone()
    racecontrol.finished()

def TaintBuffers(lists, random):
    """
    Internal use only. Don't call this.
    """
    taintgate = threadgate(len(lists))
    for eachlist in lists:
        thread = Thread(target=TaintList, args=(eachlist, taintgate, random))
        thread.start()
    taintgate.untilReady()
    taintgate.go()
    taintgate.untilDone()


def trandom(bufferAmount=4, bufferSize=8, rounds=3, x=0, y=0, z=0):
    """
    The only external interface to the TSEPRNG. Call this function to
    get random data.

    Keep in mind that this is resource intensive random number generation.
    You should only be calling this function for cryptographically secure
    entropy. It is highly customizable.

    *bufferAmount* sets the number of buffers to be raced over.
    The more, the merrier. Just watch the resources. This value should be
    proportional to your paranoia.

    *bufferSize* is the size in bytes of the buffer to randomize. Make
    sure the buffer size is a multiple of 2.

    *rounds* specify how many iterations of the algorithm to conduct. This
    value should be proportional to your paranoia. Cranking this will greatly
    increase both the time required to generate a number and the overall
    unpredictability of the result.

    *x*, *y*, *z* are seed initializers to whrandom. Although
    the TSEPRNG will produce numbers that are intractable to the seed used, you can
    set these by hand it you wish. The default (0,0,0) will result in a time-based
    seed.

    The total number of threads that participate in the race is equal to
    *bufferAmount* * *bufferSize*.

    The function will return a string of length *bufferAmount* *
    *bufferSize*.

    This function defaults to 4 buffers of 8 bytes over 3 rounds
    with seeds of 0,0,0.
    """
    if bufferAmount%2 != 0:
        raise ValueError("If the Buffer amount is a multiple of two (2), more entropy results.")
    random = whrandom()
    random.seed(x, y, z)
    buffers = CreateBuffers(bufferAmount, bufferSize, random)
    for i in range(rounds):
        TaintBuffers(buffers, random)
        ShuffleBuffers(buffers, random)
    TaintBuffers(buffers, random)          # One more taint for good luck


    # Build a string from entropic data
    outstring = ""
    for i in range(bufferAmount):
        for j in range(bufferSize):
            outstring += buffers[i][j]

    return outstring



class EntropyPool:
    """
    Cryptographically strong random number generation.

    To be cryptographically strong, it must be difficult to determine the
    RNG's output, whether in the future or the past.  This is done
    by using encryption algorithms to "stir" the random data.

    Originally AMK's but modified by Bryan Mongeau to use sha256 and
    AES as part of the CryptKit suite.

    Entropy estimates are no longer kept due to their imperfect
    estimation of the real entropy accumulated.
    """
    def __init__(self, numbytes = 256):

        """
        Numbytes is the size of the random pool.
        The random pool is initially constructed with
        succesive time hashes. /dev/urandom existence
        is tested and used if available. Finally a call
        to stir() ensures at least somewhat entropic data.
        For real crypto applications, add entropic events
        directly by calling addBytes()
        """

        self.bytes=numbytes
        self.__counter, self.__pos = 0, 0
        # New coolness
        self.cipher = aes()
        self.hash = sha256()

        # Construct a random pool from succesive hashed time values
        self._randpool = ""
        while len(self._randpool) < self.bytes:
            self.hash.update(dumps(time.time()))
            self._randpool+=self._randpool+self.hash.digest()

        # Reset the hashing context
        self.hash.digest(1)

        # Linux supports a /dev/urandom device; soon other OSes will, too.
        # We'll grab some randomness from it.
        try:
          self.useURandom()
          self.canUseURandom=1
        except IOError, (num, msg):
          if num!=2:
            raise IOError, (num, msg)
          else:
            self.canUseURandom=0
            # If the file wasn't found, ignore the error

        self.stir()

    def useURandom(self):
      """
      If /dev/urandom is present on a given platform,
      we would be stupid not to use it. Grabs bytes and
      adds them to the pool.
      """
      f=open('/dev/urandom')
      data=f.read(self.bytes)
      f.close()
      self.addBytes(data)


    def stir(self):
        """
        Critical function for ensuring the entropic
        quality of the random pool in a Yarrow-esque
        fashion. Proceeds as follows:

        - Adds the time to the random pool.

        - Taps /dev/urandom if available and adds
          it to the random pool.

        - Creates a stir "token". The token is a hash
          of the random pool and the two running counters.
          Adds the token to the random pool.

        - Encrypts the entire pool with AES using the token
          as the encryption key. Uses CBC mode to chop things
          up with the first 16 bytes of the token as the
          initialization vector.

        Calling this function is relatively cheap on the
        cycles but remember that for real security, you
        must add other entropic events using addBytes().
        """

        # Stir in the time
        self.addBytes(dumps(time.time()))

        # Use /dev/urandom if available
        if self.canUseURandom:
          self.useURandom()

        # Stir in the counters
        self.hash.update(self._randpool+str(self.__counter)+str(self.__pos))
        token = self.hash.digest(1)
        self.addBytes(token)

        # Encrypt the randpool
        self.cipher.setKey( token )
        self.cipher.setCipherMode("CBC",token[:16])
        self._randpool = self.cipher.encrypt(self._randpool)[:self.bytes]



    def getBytes(self, N=32):
        "Return N bytes of random data, 32 being the default."
        randStr=''
        while len(randStr) < N:
          self.stir()
          self.hash.update(self._randpool)
          randStr+=self.hash.digest()

        # Clear the hashing context
        self.hash.digest(1)
        return randStr[:N]

    def getInt(self):
        "Return a random integer."
        self.stir()
        self.hash.update(self._randpool)
        randstr=self.hash.digest()[:4]
        x=0
        for y in range(len(randstr)):
          x+=ord(randstr[y])
          x=x<<8
        return x

    def addBytes(self, s):
        "XORs the contents of the string S into the random pool"
        i = self.__pos
        for j in range(len(s)):
          c = chr( ord(self._randpool[i]) ^ ord(s[j]) )
          self._randpool=self._randpool[:i]+c+self._randpool[i+1:]
          i=(i+1) % self.bytes

        self.__pos=i
        self.__counter+=1


class CSPRNG(EntropyPool):
  """
  Cryptographically Secure Pseudo Random Number Generation.

  Based on RandomPool, it uses the thread scheduling entropy
  pseudo random number generator (TSEPRNG) to add entropic
  bytes to the pool.
  """
  def __init__(self, tStirNow=0):
    """
    Constructs a tgCSPRNG. Set tStirNow to true if you wish to
    run a threading stir right off the bat using the default
    settings.
    """
    EntropyPool.__init__(self)
    if tStirNow:
      self.tStir()

  def tStir(self, bufferAmount=4, bufferSize=16, rounds=3, x=0, y=0, z=0):
    """
    Uses thread scheduling entropy to stir really entropic bytes into
    the random pool. By default, it will produce and stir in 64 bytes.
    This takes about 0.6 seconds on my 800.

    For information about tweaking the thread parameters, see the
    documentation for the trandom function. And experiment :)

    tStir is very resource intensive. Use carefully and only if you
    require cryptographically strong entropy. Otherwise use whrandom! It
    is good practice to call tStir to ensure the pool is really entropic,
    then simply call stir() subsequently.
    """

    self.addBytes( trandom(bufferAmount,bufferSize,rounds,x,y,z) )
    self.stir()

