-- |
-- Module      : Crypto.Random.AESCtr
-- License     : BSD-style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : stable
-- Portability : unknown
--
-- this CPRNG is an AES based counter system.
--
-- the internal size of fields are: 16 bytes IV, 16 bytes counter, 32 bytes key
--
-- each block are generated the following way:
--   aes (IV `xor` counter) -> 16 bytes output
--
{-# LANGUAGE CPP #-}
{-# LANGUAGE BangPatterns #-}
module Crypto.Random.AESCtr
    ( AESRNG
    , make
    , makeSystem
    ) where

import Crypto.Random
import Crypto.Random.AESCtr.Internal
import Control.Arrow (second)

import Data.ByteString (ByteString)
import qualified Data.ByteString as B

import Data.Byteable
import Data.Bits (xor, (.&.))

-- | AES Counter mode Pseudo random generator.
--
-- Provide a very good Cryptographic pseudo random generator
-- that create pseudo random output based an AES cipher
-- used in counter mode, initialized from random key, random IV
-- and random nonce.
--
-- This CPRG uses 64 bytes of pure entropy to create its random state.
--
-- By default, this generator will automatically reseed after generating
-- 1 megabyte of data.
data AESRNG = AESRNG { AESRNG -> RNG
aesrngState     :: !RNG
                     , AESRNG -> EntropyPool
aesrngEntropy   :: EntropyPool
                     , AESRNG -> Int
aesrngThreshold :: !Int -- ^ in number of generated block
                     , AESRNG -> ByteString
aesrngCache     :: !ByteString }

instance Show AESRNG where
    show :: AESRNG -> String
show AESRNG
_ = String
"aesrng[..]"

makeFrom :: EntropyPool -> B.ByteString -> AESRNG
makeFrom :: EntropyPool -> ByteString -> AESRNG
makeFrom EntropyPool
entPool ByteString
b = AESRNG :: RNG -> EntropyPool -> Int -> ByteString -> AESRNG
AESRNG
    { aesrngState :: RNG
aesrngState        = RNG
rng
    , aesrngEntropy :: EntropyPool
aesrngEntropy      = EntropyPool
entPool
    , aesrngThreshold :: Int
aesrngThreshold    = Int
1024 -- in blocks generated, so 1mb
    , aesrngCache :: ByteString
aesrngCache        = ByteString
B.empty }
  where rng :: RNG
rng = ByteString -> RNG
makeRNG ByteString
b

-- | make an AES RNG from an EntropyPool.
--
-- use `makeSystem` to not have to deal with the entropy pool.
make :: EntropyPool -> AESRNG
make :: EntropyPool -> AESRNG
make EntropyPool
entPool = EntropyPool -> ByteString -> AESRNG
makeFrom EntropyPool
entPool ByteString
b
  where !b :: ByteString
b = SecureMem -> ByteString
forall a. Byteable a => a -> ByteString
toBytes (SecureMem -> ByteString) -> SecureMem -> ByteString
forall a b. (a -> b) -> a -> b
$ Int -> EntropyPool -> SecureMem
grabEntropy Int
64 EntropyPool
entPool

-- | Initialize a new AES RNG using the system entropy.
-- {-# DEPRECATED makeSystem "use cprgCreate with an entropy pool" #-}
makeSystem :: IO AESRNG
makeSystem :: IO AESRNG
makeSystem = EntropyPool -> AESRNG
make (EntropyPool -> AESRNG) -> IO EntropyPool -> IO AESRNG
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` IO EntropyPool
createEntropyPool

-- | get a Random number of bytes from the RNG.
-- it generate randomness by block of chunkSize bytes and will returns
-- a block bigger or equal to the size requested.
genRandomBytesState :: RNG -> Int -> (ByteString, RNG)
genRandomBytesState :: RNG -> Int -> (ByteString, RNG)
genRandomBytesState RNG
rng Int
n
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
chunkSize = RNG -> (ByteString, RNG)
genNextChunk RNG
rng
    | Bool
otherwise      = let ([ByteString]
bs, RNG
rng') = Int -> [ByteString] -> RNG -> ([ByteString], RNG)
acc Int
0 [] RNG
rng
                        in ([ByteString] -> ByteString
B.concat [ByteString]
bs, RNG
rng')
  where acc :: Int -> [ByteString] -> RNG -> ([ByteString], RNG)
acc Int
l [ByteString]
bs RNG
g
            | Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
chunkSize Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = ([ByteString]
bs, RNG
g)
            | Bool
otherwise          = let (ByteString
b, RNG
g') = RNG -> (ByteString, RNG)
genNextChunk RNG
g
                                    in Int -> [ByteString] -> RNG -> ([ByteString], RNG)
acc (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (ByteString
bByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:[ByteString]
bs) RNG
g'

genRanBytesNoCheck :: AESRNG -> Int -> (ByteString, AESRNG)
genRanBytesNoCheck :: AESRNG -> Int -> (ByteString, AESRNG)
genRanBytesNoCheck AESRNG
rng Int
n
    | ByteString -> Int
B.length (AESRNG -> ByteString
aesrngCache AESRNG
rng) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = let (ByteString
b1,ByteString
b2) = Int -> ByteString -> (ByteString, ByteString)
B.splitAt Int
n (AESRNG -> ByteString
aesrngCache AESRNG
rng)
                                         in (ByteString
b1, AESRNG
rng { aesrngCache :: ByteString
aesrngCache = ByteString
b2 })
    | Bool
otherwise                       =
        let (ByteString
b, RNG
rng') = RNG -> Int -> (ByteString, RNG)
genRandomBytesState (AESRNG -> RNG
aesrngState AESRNG
rng) Int
n
            (ByteString
b1, ByteString
b2)  = Int -> ByteString -> (ByteString, ByteString)
B.splitAt Int
n ByteString
b
         in (ByteString
b1, AESRNG
rng { aesrngState :: RNG
aesrngState = RNG
rng', aesrngCache :: ByteString
aesrngCache = ByteString
b2 })

-- | generate a random set of bytes
genRanBytes :: AESRNG -> Int -> (ByteString, AESRNG)
genRanBytes :: AESRNG -> Int -> (ByteString, AESRNG)
genRanBytes AESRNG
rng Int
n = (AESRNG -> AESRNG) -> (ByteString, AESRNG) -> (ByteString, AESRNG)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AESRNG -> AESRNG
reseedThreshold ((ByteString, AESRNG) -> (ByteString, AESRNG))
-> (ByteString, AESRNG) -> (ByteString, AESRNG)
forall a b. (a -> b) -> a -> b
$ AESRNG -> Int -> (ByteString, AESRNG)
genRanBytesNoCheck AESRNG
rng Int
n

reseedThreshold :: AESRNG -> AESRNG
reseedThreshold :: AESRNG -> AESRNG
reseedThreshold AESRNG
rng
    | RNG -> Int
getNbChunksGenerated (AESRNG -> RNG
aesrngState AESRNG
rng) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
lvl =
         let newRngState :: RNG
newRngState = ByteString -> RNG
makeRNG (ByteString -> RNG) -> ByteString -> RNG
forall a b. (a -> b) -> a -> b
$ SecureMem -> ByteString
forall a. Byteable a => a -> ByteString
toBytes (SecureMem -> ByteString) -> SecureMem -> ByteString
forall a b. (a -> b) -> a -> b
$ Int -> EntropyPool -> SecureMem
grabEntropy Int
64 (AESRNG -> EntropyPool
aesrngEntropy AESRNG
rng)
          in AESRNG
rng { aesrngState :: RNG
aesrngState = RNG
newRngState }
    | Bool
otherwise  = AESRNG
rng
  where lvl :: Int
lvl = AESRNG -> Int
aesrngThreshold AESRNG
rng

instance CPRG AESRNG where
    cprgCreate :: EntropyPool -> AESRNG
cprgCreate                      = EntropyPool -> AESRNG
make
    cprgSetReseedThreshold :: Int -> AESRNG -> AESRNG
cprgSetReseedThreshold Int
lvl AESRNG
rng  = AESRNG -> AESRNG
reseedThreshold (AESRNG
rng { aesrngThreshold :: Int
aesrngThreshold = if Int
nbChunks Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 then Int
nbChunks else Int
1 })
      where nbChunks :: Int
nbChunks = Int
lvl Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
chunkSize
    cprgGenerate :: Int -> AESRNG -> (ByteString, AESRNG)
cprgGenerate Int
len AESRNG
rng            = AESRNG -> Int -> (ByteString, AESRNG)
genRanBytes AESRNG
rng Int
len
    cprgGenerateWithEntropy :: Int -> AESRNG -> (ByteString, AESRNG)
cprgGenerateWithEntropy Int
len AESRNG
rng =
        let ent :: ByteString
ent        = SecureMem -> ByteString
forall a. Byteable a => a -> ByteString
toBytes (SecureMem -> ByteString) -> SecureMem -> ByteString
forall a b. (a -> b) -> a -> b
$ Int -> EntropyPool -> SecureMem
grabEntropy Int
len (AESRNG -> EntropyPool
aesrngEntropy AESRNG
rng)
            (ByteString
bs, AESRNG
rng') = AESRNG -> Int -> (ByteString, AESRNG)
genRanBytes AESRNG
rng Int
len
         in ([Word8] -> ByteString
B.pack ([Word8] -> ByteString) -> [Word8] -> ByteString
forall a b. (a -> b) -> a -> b
$ (Word8 -> Word8 -> Word8) -> ByteString -> ByteString -> [Word8]
forall a. (Word8 -> Word8 -> a) -> ByteString -> ByteString -> [a]
B.zipWith Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
xor ByteString
ent ByteString
bs, AESRNG
rng')
    cprgFork :: AESRNG -> (AESRNG, AESRNG)
cprgFork AESRNG
rng = let (ByteString
b,AESRNG
rng') = AESRNG -> Int -> (ByteString, AESRNG)
genRanBytes AESRNG
rng Int
64
                    in (AESRNG
rng', EntropyPool -> ByteString -> AESRNG
makeFrom (AESRNG -> EntropyPool
aesrngEntropy AESRNG
rng) ByteString
b)

{-
instance RandomGen AESRNG where
    next rng =
        let (bs, rng') = genRanBytes rng 16 in
        let (Word128 a _) = get128 bs in
        let n = fromIntegral (a .&. 0x7fffffff) in
        (n, rng')
    split rng =
        let rng' = make (aesrngEntropy rng)
         in (rng, rng')
    genRange _ = (0, 0x7fffffff)
-}