{-# LANGUAGE TupleSections #-}
module Crypto.Random.DRBG.CTR
    ( State
    , getCounter
    , reseedInterval
    , update
    , instantiate
    , reseed
    , generate
    ) where

import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import Crypto.Classes
import Data.Serialize
import Crypto.Types
import Crypto.Random.DRBG.Types
import Data.Word (Word64)

data State a = St { State a -> Word64
counter     :: {-# UNPACK #-} !Word64
                  , State a -> IV a
value       :: !(IV a)
                  , State a -> a
key         :: a
                  }

instance Serialize a => Serialize (State a) where
    get :: Get (State a)
get = do Word64
c <- Get Word64
getWord64be
             ByteString
v <- Get ByteString
forall t. Serialize t => Get t
get
             a
k <- Get a
forall t. Serialize t => Get t
get
             State a -> Get (State a)
forall (m :: * -> *) a. Monad m => a -> m a
return (State a -> Get (State a)) -> State a -> Get (State a)
forall a b. (a -> b) -> a -> b
$ Word64 -> IV a -> a -> State a
forall a. Word64 -> IV a -> a -> State a
St Word64
c (ByteString -> IV a
forall k. ByteString -> IV k
IV ByteString
v) a
k
    put :: Putter (State a)
put (St c :: Word64
c (IV v :: ByteString
v) k :: a
k) = Putter Word64
putWord64be Word64
c PutM () -> PutM () -> PutM ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Putter ByteString
forall t. Serialize t => Putter t
put ByteString
v PutM () -> PutM () -> PutM ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Putter a
forall t. Serialize t => Putter t
put a
k

-- |Get a count of how many times this generator has been used since
-- instantiation or reseed.
getCounter :: State a -> Word64
getCounter :: State a -> Word64
getCounter = State a -> Word64
forall a. State a -> Word64
counter

-- |Update the RNG
update :: BlockCipher a => ByteString -> State a -> Maybe (State a)
update :: ByteString -> State a -> Maybe (State a)
update provided_data :: ByteString
provided_data st :: State a
st
    | ByteString -> Int
B.length ByteString
provided_data Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
seedLen = Maybe (State a)
forall a. Maybe a
Nothing
    | Bool
otherwise =
        let (temp :: ByteString
temp,_) = a -> IV a -> ByteString -> (ByteString, IV a)
forall k.
BlockCipher k =>
k -> IV k -> ByteString -> (ByteString, IV k)
ctr (State a -> a
forall a. State a -> a
key State a
st) (State a -> IV a
forall a. State a -> IV a
value State a
st) (Int -> Word8 -> ByteString
B.replicate Int
seedLen 0)
            (keyBytes :: ByteString
keyBytes,valBytes :: ByteString
valBytes) = Int -> ByteString -> (ByteString, ByteString)
B.splitAt Int
keyLen (ByteString -> ByteString -> ByteString
zwp' ByteString
temp ByteString
provided_data)
            newValue :: IV k
newValue = ByteString -> IV k
forall k. ByteString -> IV k
IV ByteString
valBytes
            newKey :: Maybe a
newKey   = ByteString -> Maybe a
forall k. BlockCipher k => ByteString -> Maybe k
buildKey ByteString
keyBytes
        in Word64 -> IV a -> a -> State a
forall a. Word64 -> IV a -> a -> State a
St (State a -> Word64
forall a. State a -> Word64
counter State a
st) IV a
forall k. IV k
newValue (a -> State a) -> Maybe a -> Maybe (State a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` Maybe a
newKey
  where
    keyLen :: Int
keyLen  = Tagged a Int
forall k. BlockCipher k => Tagged k Int
keyLengthBytes Tagged a Int -> a -> Int
forall a b. Tagged a b -> a -> b
`for` State a -> a
forall a. State a -> a
key State a
st
    blkLen :: Int
blkLen  = Tagged a Int
forall k. BlockCipher k => Tagged k Int
blockSizeBytes Tagged a Int -> a -> Int
forall a b. Tagged a b -> a -> b
`for` State a -> a
forall a. State a -> a
key State a
st
    seedLen :: Int
seedLen = Int
keyLen Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
blkLen
{-# INLINEABLE update #-}

-- | Instantiate a new CTR based counter.  This assumes the block cipher is
-- safe for generating 2^48 seperate bitstrings (e.g. For SP800-90 we
-- assume AES and not 3DES)
instantiate :: BlockCipher a => Entropy -> PersonalizationString -> Maybe (State a)
instantiate :: ByteString -> ByteString -> Maybe (State a)
instantiate ent :: ByteString
ent perStr :: ByteString
perStr = Maybe (State a)
st
  where
  seedLen :: Int
seedLen   = Int
blockLen Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
keyLen
  blockLen :: Int
blockLen  = Tagged a Int
forall k. BlockCipher k => Tagged k Int
blockSizeBytes Tagged a Int -> a -> Int
forall a b. Tagged a b -> a -> b
`for` Maybe (State a) -> a
forall a. Maybe (State a) -> a
keyOfState Maybe (State a)
st
  keyLen :: Int
keyLen    = Tagged a Int
forall k. BlockCipher k => Tagged k Int
keyLengthBytes Tagged a Int -> a -> Int
forall a b. Tagged a b -> a -> b
`for` Maybe (State a) -> a
forall a. Maybe (State a) -> a
keyOfState Maybe (State a)
st
  temp :: ByteString
temp      = Int -> ByteString -> ByteString
B.take Int
seedLen (ByteString -> ByteString -> ByteString
B.append ByteString
perStr (Int -> Word8 -> ByteString
B.replicate Int
seedLen 0))
  seedMat :: ByteString
seedMat   = ByteString -> ByteString -> ByteString
zwp' ByteString
ent ByteString
temp
  key0 :: Maybe a
key0      = ByteString -> Maybe a
forall k. BlockCipher k => ByteString -> Maybe k
buildKey (Int -> Word8 -> ByteString
B.replicate Int
keyLen 0)
  v0 :: IV a
v0        = ByteString -> IV a
forall k. ByteString -> IV k
IV (Int -> Word8 -> ByteString
B.replicate Int
blockLen 0)
  st :: Maybe (State a)
st        = do a
k <- Maybe a
key0
                 ByteString -> State a -> Maybe (State a)
forall a. BlockCipher a => ByteString -> State a -> Maybe (State a)
update ByteString
seedMat (Word64 -> IV a -> a -> State a
forall a. Word64 -> IV a -> a -> State a
St 1 IV a
v0 a
k)
{-# INLINABLE instantiate #-}

keyOfState :: Maybe (State a) -> a
keyOfState :: Maybe (State a) -> a
keyOfState = a -> Maybe (State a) -> a
forall a b. a -> b -> a
const a
forall a. HasCallStack => a
undefined

-- |@reseed oldRNG entropy additionalInfo@
--
-- Reseed a DRBG with some entropy ('ent' must be at least seedlength, which is the
-- block length plus the key length)
reseed :: BlockCipher a => State a -> Entropy -> AdditionalInput -> Maybe (State a)
reseed :: State a -> ByteString -> ByteString -> Maybe (State a)
reseed st0 :: State a
st0 ent :: ByteString
ent ai :: ByteString
ai = Maybe (State a)
st1
  where
  seedLen :: Int
seedLen = (Tagged a Int
forall k. BlockCipher k => Tagged k Int
blockSizeBytes Tagged a Int -> a -> Int
forall a b. Tagged a b -> a -> b
`for` State a -> a
forall a. State a -> a
key State a
st0) Int -> Int -> Int
forall a. Num a => a -> a -> a
+
            (Tagged a Int
forall k. BlockCipher k => Tagged k Int
keyLengthBytes Tagged a Int -> a -> Int
forall a b. Tagged a b -> a -> b
`for` State a -> a
forall a. State a -> a
key State a
st0)
  newAI :: ByteString
newAI   = Int -> ByteString -> ByteString
B.take Int
seedLen (ByteString -> ByteString -> ByteString
B.append ByteString
ai (Int -> Word8 -> ByteString
B.replicate Int
seedLen 0))
  seedMat :: ByteString
seedMat = ByteString -> ByteString -> ByteString
zwp' ByteString
ent ByteString
newAI
  st1 :: Maybe (State a)
st1     = ByteString -> State a -> Maybe (State a)
forall a. BlockCipher a => ByteString -> State a -> Maybe (State a)
update ByteString
seedMat (State a
st0 { counter :: Word64
counter = 1} )
{-# INLINABLE reseed #-}

-- |Generate new bytes of data, stepping the generator.
generate :: BlockCipher a => State a -> ByteLength -> AdditionalInput -> Maybe (RandomBits, State a)
generate :: State a -> Int -> ByteString -> Maybe (ByteString, State a)
generate st0 :: State a
st0 len :: Int
len ai0 :: ByteString
ai0
  | State a -> Word64
forall a. State a -> Word64
counter State a
st0 Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
> Word64
reseedInterval = Maybe (ByteString, State a)
forall a. Maybe a
Nothing
  | Bool -> Bool
not (ByteString -> Bool
B.null ByteString
ai0) =
      let aiNew :: ByteString
aiNew = Int -> ByteString -> ByteString
B.take Int
seedLen (ByteString -> ByteString -> ByteString
B.append ByteString
ai0 (Int -> Word8 -> ByteString
B.replicate Int
seedLen 0))
      in do State a
st' <- ByteString -> State a -> Maybe (State a)
forall a. BlockCipher a => ByteString -> State a -> Maybe (State a)
update ByteString
aiNew State a
st0
            State a -> ByteString -> Maybe (ByteString, State a)
forall a.
BlockCipher a =>
State a -> ByteString -> Maybe (ByteString, State a)
go State a
st' ByteString
aiNew
  | Bool
otherwise = State a -> ByteString -> Maybe (ByteString, State a)
forall a.
BlockCipher a =>
State a -> ByteString -> Maybe (ByteString, State a)
go State a
st0 (Int -> Word8 -> ByteString
B.replicate Int
seedLen 0)
  where
  outLen :: Int
outLen  = Tagged a Int
forall k. BlockCipher k => Tagged k Int
blockSizeBytes Tagged a Int -> a -> Int
forall a b. Tagged a b -> a -> b
`for` State a -> a
forall a. State a -> a
key State a
st0
  keyLen :: Int
keyLen  = Tagged a Int
forall k. BlockCipher k => Tagged k Int
keyLengthBytes Tagged a Int -> a -> Int
forall a b. Tagged a b -> a -> b
`for` State a -> a
forall a. State a -> a
key State a
st0
  seedLen :: Int
seedLen = Int
outLen Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
keyLen
  -- go :: BlockCipher a => State a
  --                     -> AdditionalInput
  --                     -> Maybe (RandomBits, State a)
  go :: State a -> ByteString -> Maybe (ByteString, State a)
go st :: State a
st ai :: ByteString
ai =
      let (temp :: ByteString
temp,v2 :: IV a
v2) = a -> IV a -> ByteString -> (ByteString, IV a)
forall k.
BlockCipher k =>
k -> IV k -> ByteString -> (ByteString, IV k)
ctr (State a -> a
forall a. State a -> a
key State a
st) (State a -> IV a
forall a. State a -> IV a
value State a
st) (Int -> Word8 -> ByteString
B.replicate Int
len 0)
          st1 :: Maybe (State a)
st1       = ByteString -> State a -> Maybe (State a)
forall a. BlockCipher a => ByteString -> State a -> Maybe (State a)
update ByteString
ai (State a
st { value :: IV a
value = IV a
v2
                                    , counter :: Word64
counter = State a -> Word64
forall a. State a -> Word64
counter State a
st Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ 1 })
      in (State a -> (ByteString, State a))
-> Maybe (State a) -> Maybe (ByteString, State a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ByteString
temp,) Maybe (State a)
st1
{-# INLINABLE generate #-}

-- |The reseed interval
reseedInterval :: Word64
reseedInterval :: Word64
reseedInterval = 2Word64 -> Integer -> Word64
forall a b. (Num a, Integral b) => a -> b -> a
^48