    {-# LANGUAGE FlexibleInstances, TypeSynonymInstances, BangPatterns, ScopedTypeVariables #-}
    {-|
     Maintainer: Thomas.DuBuisson@gmail.com
     Stability: beta
     Portability: portable

    This module is the convenience interface for the DRBG (NIST standardized
    number-theoretically secure random number generator).  Everything is setup
    for using the "crypto-api" 'CryptoRandomGen' type class.

    To instantiate the base types of 'HmacDRBG', 'HashDRBG', or 'CtrDRBG' just use
    the 'CryptoRandomGen' primitives of 'newGen' or 'newGenIO'.

    For example, to seed a new generator with the system secure random
    ('System.Entropy') and generate some bytes (stepping the generator along
    the way) one would do:

    @
        gen <- newGenIO :: IO CtrDRBG
        let Right (randomBytes, newGen) = genBytes 1024 gen
    @

    or the same thing with your own entropy (throwing exceptions instead of dealing
    with 'Either' this time):

    @
        let gen = throwLeft (newGen entropy)
            (bytes,gen') = throwLeft (genBytes 1024 gen)
        in ...
    @

    There are several modifiers that allow you to compose generators together, producing
    generators with modified security, reseed, and performance properties.  'GenXor'
    will xor the random bytes of two generators.  'GenBuffered' will spark off work
    to generate several megabytes of random data and keep that data buffered for
    quick use.  'GenAutoReseed' will use one generator to automatically reseed
    another after every 32 kilobytes of requested randoms.

    Most likely you will want to automatically reseed using system randomness
    (via lazy IO).  Thus, you are left with a generator that is random, not
    pseudo random but without the dangerously unsafe IO found in some other
    RNGs:

    @
        import Crypto.Random.DRBG hiding (genBytes)
        import Crypto.Classes.Exceptions (genBytes)

        -- An AES CTR generator that automatically reseeds.
        getCtrGen :: IO (GenAutoReseed CtrDRBG SystemEntropy)
        getCtrGen = newGenAutoReseedIO

        f = do g1 <- getCtrGen
               let (bytes, g2) = getBytes 1024 g1
               g bytes g2
    @

    For a complex example, here is a generator that buffers several megabytes of
    random values which are an Xor of AES with a SHA384 hash that are each reseeded
    every 32kb with the output of a SHA512 HMAC generator (not to claim this has
    any enhanced security properties, but just to show the composition can be
    nested).

    @
        gen <- newGenIO :: IO (GenBuffered (GenAutoReseed (GenXor AesCntDRBG (HashDRBGWith SHA384)) HmacDRBG))
    @

     -}

    module Crypto.Random.DRBG
            (
            -- * Basic Hash-based Generators
              HmacDRBG, HashDRBG, CtrDRBG
            , HmacDRBGWith, HashDRBGWith, CtrDRBGWith
            -- * CryptoRandomGen Transformers
            , GenXor(..)
            , GenBuffered
            , GenAutoReseed
            -- * AutoReseed generator construction with custom reseed interval
            , newGenAutoReseed, newGenAutoReseedIO
            -- * Helper Re-exports
            , module Crypto.Random
            , module Crypto.Types
            ) where

    import qualified Crypto.Random.DRBG.HMAC as M
    import qualified Crypto.Random.DRBG.Hash as H
    import qualified Crypto.Random.DRBG.CTR  as CTR
    import Crypto.Util
    import Crypto.Classes
    import Crypto.Random
    import Crypto.Hash.CryptoAPI
    import Crypto.Cipher.AES128 (AESKey128)
    import Crypto.Types
    import System.Entropy
    import qualified Data.ByteString as B
    import Data.Tagged
    import Control.Parallel
    import Control.Monad.Error () -- Either instance
    import Data.Word

    instance H.SeedLength SHA512 where
            seedlen = Tagged 888

    instance H.SeedLength SHA384 where
            seedlen = Tagged  888

    instance H.SeedLength SHA256 where
            seedlen = Tagged 440

    instance H.SeedLength SHA224 where
            seedlen = Tagged 440

    instance H.SeedLength SHA1 where
            seedlen = Tagged 440

    -- |The HMAC DRBG state (of kind * -> *) allowing selection
    -- of the underlying hash algorithm (SHA1, SHA224 ... SHA512)
    type HmacDRBGWith = M.State

    -- |The Hash DRBG state (of kind * -> *) allowing selection
    -- of the underlying hash algorithm.
    type HashDRBGWith = H.State

    -- |The AES CTR DRBG state (of kind * -> *) allowing selection
    -- of the underlying cipher algorithm.
    type CtrDRBGWith = CTR.State

    -- |An alias for an HMAC DRBG generator using SHA512.
    type HmacDRBG = M.State SHA512

    -- |An Alias for a Hash DRBG generator using SHA512.
    --
    -- As of 1July2014 this remains the fastest cryptographic RNG on hackage
    -- that has been ran against known answer tests.
    type HashDRBG = H.State SHA512

    -- |The recommended generator which uses AES-128 in counter mode.
    --
    -- This is an alias for a Counter DRBG generator using AES 128.
    type CtrDRBG = CTR.State AESKey128

    -- |@newGenAutoReseed bs i@ creates a new 'GenAutoReseed' with a custom interval
    -- of @i@ bytes using the provided entropy in @bs@.
    --
    -- This is for extremely long running uses of 'CryptoRandomGen' instances
    -- that can't explicitly reseed as often as a single underlying generator
    -- would need (usually every 2^48 bytes).
    --
    -- For example:
    --
    -- @
    -- newGenAutoReseedIO (2^48) :: IO (Either GenError (GenAutoReseed HashDRBG HashDRBG))
    -- @
    --
    -- Will last for @2^48 * 2^41@ bytes of randomly generated data.  That's
    -- 2^49 terabytes of random values (128 byte reseeds every 2^48 bytes generated).
    newGenAutoReseed :: (CryptoRandomGen a, CryptoRandomGen b) => B.ByteString -> Word64 -> Either GenError (GenAutoReseed a b)
    newGenAutoReseed bs rsInterval=
            let (b1,b2) = B.splitAt (genSeedLength `for` fromRight g1) bs
                g1 = newGen b1
                g2 = newGen b2
                fromRight ~(Right x) = x
            in case (g1, g2) of
                    (Right a, Right b) -> Right $ GenAutoReseed rsInterval 0 a b
                    (Left e, _) -> Left e
                    (_, Left e) -> Left e

    -- |@newGenAutoReseedIO i@ creates a new 'GenAutoReseed' with a custom
    -- interval of @i@ bytes, using the system random number generator as a seed.
    --
    -- See 'newGenAutoReseed'.
    newGenAutoReseedIO :: (CryptoRandomGen a, CryptoRandomGen b) => Word64 -> IO (GenAutoReseed a b)
    newGenAutoReseedIO i   = do
            g1 <- newGenIO
            g2 <- newGenIO
            return $ GenAutoReseed i 0 g1 g2

    instance CryptoRandomGen HmacDRBG where
            newGen bs =
                    let res = M.instantiate bs B.empty B.empty
                    in if B.length bs < genSeedLength `for` res
                            then Left NotEnoughEntropy
                            else Right res
            genSeedLength = Tagged (512 `div` 8)
            genBytes req g =
                    let res = M.generate g (req * 8) B.empty
                    in case res of
                            Nothing -> Left NeedReseed
                            Just (r,s) -> Right (r, s)
            genBytesWithEntropy req ai g =
                    let res = M.generate g (req * 8) ai
                    in case res of
                            Nothing -> Left NeedReseed
                            Just (r,s) -> Right (r, s)
            reseed ent g =
                    let res = M.reseed g ent B.empty
                    in if B.length ent < genSeedLength `for` res
                            then Left NotEnoughEntropy
                            else Right res

            reseedInfo s = InXCalls (M.counter s)
            reseedPeriod _ = InXCalls M.reseedInterval

    instance CryptoRandomGen HashDRBG where
            newGen bs =
                    let res = H.instantiate bs B.empty B.empty
                    in if B.length bs < genSeedLength `for` res
                            then Left NotEnoughEntropy
                            else Right res
            genSeedLength = Tagged $ 512 `div` 8
            genBytes req g =
                    let res = H.generate g (req * 8) B.empty
                    in case res of
                            Nothing -> Left NeedReseed
                            Just (r,s) -> Right (r, s)
            genBytesWithEntropy req ai g =
                    let res = H.generate g (req * 8) ai
                    in case res of
                            Nothing -> Left NeedReseed
                            Just (r,s) -> Right (r, s)
            reseed ent g =
                    let res = H.reseed g ent B.empty
                    in if B.length ent < genSeedLength `for` res
                            then Left NotEnoughEntropy
                            else Right res
            reseedInfo s = InXCalls (H.counter s)
            reseedPeriod _ = InXCalls H.reseedInterval

    helper1 :: Tagged (GenAutoReseed a b) Int -> a
    helper1 = const undefined

    helper2 :: Tagged (GenAutoReseed a b) Int -> b
    helper2 = const undefined

    -- |@g :: GenAutoReseed a b@ is a generator of type a that gets
    -- automatically reseeded by generator b upon every 32kB generated.
    --
    -- @reseed g ent@ will reseed both the component generators by
    -- breaking ent up into two parts determined by the genSeedLength of each generator.
    --
    -- @genBytes@ will generate the requested bytes with generator @a@ and reseed @a@
    -- using generator @b@ if there has been 32KB of generated data since the last reseed.
    -- Note a request for > 32KB of data will be filled in one request to generator @a@ before
    -- @a@ is reseeded by @b@.
    --
    -- @genBytesWithEntropy@ is lifted into the same call for generator @a@, but
    -- it will still reseed from generator @b@ if the limit is hit.
    --
    -- Reseed interval: If generator @a@ needs a @genSeedLength a = a'@ and generator B
    -- needs reseeded every @2^b@ bytes then a @GenAutoReseed a b@ will need reseeded every
    -- @2^15 * (2^b / a')@ bytes.  For the common values of @a' = 128@ and @2^b = 2^48@ this
    -- means reseeding every 2^56 byte.  For the example numbers this translates to
    -- about 200 years of continually generating random values at a rate of 10MB/s.
    data GenAutoReseed a b = GenAutoReseed
            { garInterval  :: {-# UNPACK #-} !Word64
            , garCounter   :: {-# UNPACK #-} !Word64
            , garPrimeGen  :: !a
            , garBackupGen :: !b
            }

    genBytesAutoReseed :: (CryptoRandomGen a, CryptoRandomGen b)
                       => ByteLength
                       -> GenAutoReseed a b
                       -> Either GenError (B.ByteString, GenAutoReseed a b)
    genBytesAutoReseed req gar@(GenAutoReseed rs cnt a b) =
        case genBytes req a of
           Left NeedReseed -> do
                   (a',b') <- a `reseedWith` b
                   (res, aNew) <- genBytes req a'
                   return (res,GenAutoReseed rs 0 aNew b')
           Left RequestedTooManyBytes -> do
             let reqSmall = 1 + (req `div` 2)
             (s1, gar1) <- genBytes reqSmall gar
             (s2, gar2) <- genBytes reqSmall gar1
             return (B.take req (B.append s1 s2), gar2)
           Left err -> Left err
           Right (res,aNew) -> do
             gNew <- if (cnt + fromIntegral req) > rs
                       then do
                         (a',b') <- a `reseedWith` b
                         return (GenAutoReseed rs 0 a' b')
                       else return $ GenAutoReseed rs (cnt + fromIntegral req) aNew b
             return (res, gNew)

    reseedWith :: (CryptoRandomGen a, CryptoRandomGen b)
             => a -> b -> Either GenError (a,b)
    reseedWith x y = do
         (ent,y2) <- genBytes (genSeedLength `for` x) y
         x2 <- reseed ent x
         return (x2,y2)

    genBytesWithEntropyAutoReseed
            :: (CryptoRandomGen a, CryptoRandomGen b)
            => ByteLength
            -> B.ByteString
            -> GenAutoReseed a b
            -> Either GenError (B.ByteString, GenAutoReseed a b)
    genBytesWithEntropyAutoReseed req entropy gar@(GenAutoReseed rs cnt a b) =
       case genBytesWithEntropy req entropy a of
         Left NeedReseed -> do
                 (a',b') <- a `reseedWith` b
                 (res, aNew) <- genBytesWithEntropy req entropy a'
                 return (res, GenAutoReseed rs 0 aNew b')
         Left RequestedTooManyBytes -> do
           let reqSmall = 1 + (req `div` 2)
           (s1, gar1) <- genBytesWithEntropy reqSmall entropy gar
           (s2, gar2) <- genBytes reqSmall gar1
           return (B.take req (B.append s1 s2), gar2)
         Left err -> Left err
         Right (res,aNew) -> do
           gNew <- if (cnt + fromIntegral req) > rs
                     then do
                       (a',b') <- a `reseedWith` b
                       return (GenAutoReseed rs 0 a' b')
                     else return $ GenAutoReseed rs (cnt + fromIntegral req) aNew b
           return (res, gNew)

    instance (CryptoRandomGen a, CryptoRandomGen b) => CryptoRandomGen (GenAutoReseed a b) where
            {-# SPECIALIZE instance CryptoRandomGen (GenAutoReseed HmacDRBG HmacDRBG) #-}
            {-# SPECIALIZE instance CryptoRandomGen (GenAutoReseed HashDRBG HashDRBG) #-}
            {-# SPECIALIZE instance CryptoRandomGen (GenAutoReseed HashDRBG HmacDRBG) #-}
            {-# SPECIALIZE instance CryptoRandomGen (GenAutoReseed HmacDRBG HashDRBG) #-}
            newGen bs = newGenAutoReseed bs (2^(19::Int))
            newGenIO  = newGenAutoReseedIO (2^(19::Int))
            genSeedLength =
               let a = helper1 res
                   b = helper2 res
                   res = Tagged $ genSeedLength `for` a + genSeedLength `for` b
               in res
            genBytes = genBytesAutoReseed
            genBytesWithEntropy = genBytesWithEntropyAutoReseed
            reseed ent gen@(GenAutoReseed rs _ a b)
              | genSeedLength `for` gen > B.length ent = Left NotEnoughEntropy
              | otherwise = do
                    let (e1,e2) = B.splitAt (genSeedLength `for` a) ent
                    a' <- reseed e1 a
                    b' <- if B.length e2 /= 0
                            then reseed e2 b
                            else return b
                    return $ GenAutoReseed rs 0 a' b'
            reseedPeriod ~(GenAutoReseed rs _ ag bg) =
              case (reseedPeriod ag, reseedPeriod bg) of
                (Never, _) -> Never
                (_, Never) -> Never
                (NotSoon, _) -> NotSoon
                (_, NotSoon) -> NotSoon
                (_, InXCalls b) ->
                     if fromIntegral rs * fromIntegral b >
                        fromIntegral (maxBound `asTypeOf` b)
                         then NotSoon
                         else InXBytes (rs * b)
                (_, InXBytes b) ->
                     let s = genSeedLength `for` ag
                         nr = if s <= 0 then 1 else (b `div` fromIntegral s) - 1
                     in InXBytes $ rs * nr
            reseedInfo (GenAutoReseed rs x ag bg) =
              -- Attempt to provide a lower bound on the next reseed
              case (reseedInfo ag, reseedInfo bg) of
                 (NotSoon, _)    -> NotSoon
                 (_, NotSoon)    -> NotSoon
                 (Never, _)      -> Never
                 (_, Never)      -> Never
                 (_, InXBytes b) ->
                     let s = genSeedLength `for` ag
                         nr = if s <= 0 then 1 else (b `div` fromIntegral s) - 1
                     in InXBytes $ rs - x + rs * nr
                 (_, InXCalls b) ->
                     if fromIntegral rs * fromIntegral b >
                        fromIntegral (maxBound `asTypeOf` b)
                         then NotSoon
                         else InXBytes (rs - x + rs * b)

    -- |@g :: GenXor a b@ generates bytes with sub-generators a and b
    -- and exclusive-or's the outputs to produce the resulting bytes.
    data GenXor a b = GenXor !a !b

    helperXor1 :: Tagged (GenXor a b) c -> a
    helperXor1 = const undefined
    helperXor2 :: Tagged (GenXor a b) c -> b
    helperXor2 = const undefined

    instance (CryptoRandomGen a, CryptoRandomGen b) => CryptoRandomGen (GenXor a b) where
            {-# SPECIALIZE instance CryptoRandomGen (GenXor HmacDRBG HmacDRBG) #-}
            {-# SPECIALIZE instance CryptoRandomGen (GenXor HashDRBG HmacDRBG) #-}
            {-# SPECIALIZE instance CryptoRandomGen (GenXor HmacDRBG HashDRBG) #-}
            {-# SPECIALIZE instance CryptoRandomGen (GenXor HashDRBG HashDRBG) #-}
            newGen bs = do
                    let g1 = newGen b1
                        g2 = newGen b2
                        (b1,b2) = B.splitAt (genSeedLength `for` fromRight g1) bs
                        fromRight (Right x) = x
                    a <- g1
                    b <- g2
                    return (GenXor a b)
            newGenIO = do
                    a <- newGenIO
                    b <- newGenIO
                    return (GenXor a b)
            genSeedLength =
                    let a = helperXor1 res
                        b = helperXor2 res
                        res = Tagged $ (genSeedLength `for` a) + (genSeedLength `for` b)
                    in res
            genBytes req (GenXor a b) = do
                    (r1, a') <- genBytes req a
                    (r2, b') <- genBytes req b
                    return (zwp' r1 r2, GenXor a' b')
            genBytesWithEntropy req ent (GenXor a b) = do
                    (r1, a') <- genBytesWithEntropy req ent a
                    (r2, b') <- genBytesWithEntropy req ent b
                    return (zwp' r1 r2, GenXor a' b')
            reseed ent (GenXor a b) = do
                    let (b1, b2) = B.splitAt (genSeedLength `for` a) ent
                    a' <- reseed b1 a
                    b' <- reseed b2 b
                    return (GenXor a' b')
            reseedPeriod ~(GenXor a b) = min (reseedPeriod a) (reseedPeriod b)
            reseedInfo   ~(GenXor a b) = min (reseedInfo a) (reseedInfo b)

    -- |@g :: GenBuffered a@ is a generator of type @a@ that attempts to
    -- maintain a buffer of random values size >= 1MB and <= 5MB at any time.
    data GenBuffered g = GenBuffered Int Int (Either (GenError, g) (B.ByteString, g)) {-# UNPACK #-} !B.ByteString

    bufferMinDef, bufferMaxDef :: Int
    bufferMinDef = 2^20
    bufferMaxDef = 2^22

    newGenBuffered :: (CryptoRandomGen g) => Int -> Int -> B.ByteString -> Either GenError (GenBuffered g)
    newGenBuffered min max bs = do
            g <- newGen bs
            (rs,g') <- genBytes min g
            let new = wrapErr (genBytes min g') g'
            rs `par` return (GenBuffered min max new rs)

    newGenBufferedIO :: CryptoRandomGen g => Int -> Int -> IO (GenBuffered g)
    newGenBufferedIO min max = do
            g <- newGenIO
            let (Right !gBuf) = do
                    (rs,g') <- genBytes min g
                    let new = wrapErr (genBytes min g') g'
                    rs `par` return (GenBuffered min max new rs)
            return gBuf

    instance (CryptoRandomGen g) => CryptoRandomGen (GenBuffered g) where
            {-# SPECIALIZE instance CryptoRandomGen (GenBuffered HmacDRBG) #-}
            {-# SPECIALIZE instance CryptoRandomGen (GenBuffered HashDRBG) #-}
            {-# SPECIALIZE instance CryptoRandomGen (GenBuffered CtrDRBG) #-}
            newGen = newGenBuffered bufferMinDef bufferMaxDef
            newGenIO = newGenBufferedIO bufferMinDef bufferMaxDef
            genSeedLength =
                    let a = help res
                        res = Tagged $ genSeedLength `for` a
                    in res
              where
              help :: Tagged (GenBuffered g) c -> g
              help = const undefined
            genBytes req (GenBuffered min max g bs)
                    | remSize >= min =  Right (B.take req bs, GenBuffered min max g (B.drop req bs))
                    | B.length bs < min =
                            case g of
                                    Left (err,_)  -> Left err
                                    Right g   -> Left (GenErrorOther "Buffering generator failed to buffer properly - unknown reason")
                    | req > B.length bs =
                        case g of
                          Left (err,_) -> Left err
                          Right (bo,g2)     ->
                            case genBytes req g2 of
                              Left err -> Left err
                              Right (b,g3) ->
                                Right (b, GenBuffered min max (Right (bo,g3)) bs)
                    | remSize < min =
                            case g of
                                    Left (err,_) -> Left err
                                    Right (rnd, gen) ->
                                            let new | B.length rnd > 0 = wrapErr (genBytes (max - (remSize + B.length rnd)) gen) gen
                                                    | otherwise = Right (B.empty,gen)
                                                (rs,rem) = B.splitAt req bs
                                            in eval new `par` Right (rs, GenBuffered min max new (B.append rem rnd))
                    | otherwise = Left $ GenErrorOther "Buffering generator hit an impossible case.  Please inform the Haskell crypto-api maintainer"
              where
              remSize = B.length bs - req
            genBytesWithEntropy req ent g = reseed ent g >>= \gen -> genBytes req gen
            reseed ent (GenBuffered min max g bs) = do
                    let (rs, g') =
                          case g of
                            Left (_,g') -> (B.empty, g')
                            Right (rs, g') -> (rs, g')
                    g'' <- reseed ent g'
                    let new = wrapErr (genBytes (min - B.length bs') g'') g''
                        bs' = B.take max (B.append bs rs)
                    return (GenBuffered min max new bs')
            reseedPeriod ~(GenBuffered _ _ g _) = reseedPeriod . either snd snd $ g
            reseedInfo ~(GenBuffered _ _ g _) = reseedInfo . either snd snd $ g

    wrapErr :: Either x y -> g -> Either (x,g) y
    wrapErr (Left x) g = Left (x,g)
    wrapErr (Right r) _ = Right r

    -- |Force evaluation for use by GenBuffered.
    eval :: Either x (B.ByteString, g) -> Either x (B.ByteString, g)
    eval (Left x) = Left x
    eval (Right (g,bs)) = bs `seq` (g `seq` Right (g, bs))

    instance BlockCipher x => CryptoRandomGen (CtrDRBGWith x) where
      newGen bytes =
          case CTR.instantiate bytes B.empty of
            Nothing -> Left NotEnoughEntropy
            Just st -> Right st

      newGenIO = do
          let k = undefined :: x -- Types weren't unifying with proxy, bug?
              b = (keyLength `for` k) + (blockSize `for` k) + 7
              e =  "Unable to generate enough entropy to instantiate CTR DRBG"
          kd <- getEntropy (b `div` 8)
          case CTR.instantiate kd B.empty of
            Nothing -> error e
            Just st -> return st

      genSeedLength =
            let rt :: Tagged x Int -> Tagged x Int -> Tagged (CtrDRBGWith x) Int
                rt x y = Tagged $
                    let k = unTagged x
                        b = unTagged y
                    in k + b
            in rt keyLengthBytes blockSizeBytes

      -- Generates req bytes (even if that is over the 2^16 bytes allowed by
      -- the NIST standard!)
      genBytes req st =
          case CTR.generate st req B.empty of
                Nothing       -> Left NeedReseed
                Just (bs,new) -> Right (bs,new)

      reseed ent st  =
        case CTR.reseed st ent B.empty of
          Nothing -> Left NeedReseed
          Just s  -> Right s

      reseedPeriod _ = InXCalls CTR.reseedInterval

      reseedInfo st  = InXCalls (CTR.reseedInterval - CTR.getCounter st)