    {-# LANGUAGE TypeApplications #-}

    module Crypto.Random.HmacDRG (HmacDRG, initial, update) where

    import Crypto.Hash
    import Crypto.MAC.HMAC (HMAC (..), hmac)
    import Crypto.Random.Types
    import Data.ByteArray (ByteArrayAccess, Bytes, ScrubbedBytes)
    import qualified Data.ByteArray as M
    import Data.Maybe

    -- | HMAC-based Deterministic Random Generator
    --
    -- Adapted from NIST Special Publication 800-90A Revision 1, Section 10.1.2
    data HmacDRG hash = HmacDRG (Digest hash) (Digest hash)

    -- | The initial DRG state. It should be seeded via 'update' before use.
    initial :: HashAlgorithm hash => hash -> HmacDRG hash
    initial :: forall hash. HashAlgorithm hash => hash -> HmacDRG hash
    initial hash
    algorithm = Digest hash -> Digest hash -> HmacDRG hash
    forall hash. Digest hash -> Digest hash -> HmacDRG hash
    HmacDRG (Word8 -> Digest hash
    constant Word8
    0x00) (Word8 -> Digest hash
    constant Word8
    0x01)
      where
        constant :: Word8 -> Digest hash
    constant =
            Maybe (Digest hash) -> Digest hash
    forall a. HasCallStack => Maybe a -> a
    fromJust (Maybe (Digest hash) -> Digest hash)
    -> (Word8 -> Maybe (Digest hash)) -> Word8 -> Digest hash
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Bytes -> Maybe (Digest hash)
    forall a ba.
    (HashAlgorithm a, ByteArrayAccess ba) =>
    ba -> Maybe (Digest a)
    digestFromByteString (Bytes -> Maybe (Digest hash))
    -> (Word8 -> Bytes) -> Word8 -> Maybe (Digest hash)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . forall ba. ByteArray ba => Int -> Word8 -> ba
    M.replicate @Bytes (hash -> Int
    forall a. HashAlgorithm a => a -> Int
    hashDigestSize hash
    algorithm)

    -- | Update the DRG state with optional provided data.
    update
        :: ByteArrayAccess input
        => HashAlgorithm hash
        => input -> HmacDRG hash -> HmacDRG hash
    update :: forall input hash.
    (ByteArrayAccess input, HashAlgorithm hash) =>
    input -> HmacDRG hash -> HmacDRG hash
    update input
    input HmacDRG hash
    state0 = if input -> Bool
    forall a. ByteArrayAccess a => a -> Bool
    M.null input
    input then HmacDRG hash
    state1 else HmacDRG hash
    state2
      where
        state1 :: HmacDRG hash
    state1 = Word8 -> HmacDRG hash -> HmacDRG hash
    forall {hash} {hash}.
    HashAlgorithm hash =>
    Word8 -> HmacDRG hash -> HmacDRG hash
    step Word8
    0x00 HmacDRG hash
    state0
        state2 :: HmacDRG hash
    state2 = Word8 -> HmacDRG hash -> HmacDRG hash
    forall {hash} {hash}.
    HashAlgorithm hash =>
    Word8 -> HmacDRG hash -> HmacDRG hash
    step Word8
    0x01 HmacDRG hash
    state1
        step :: Word8 -> HmacDRG hash -> HmacDRG hash
    step Word8
    byte (HmacDRG Digest hash
    key Digest hash
    value) = Digest hash -> Digest hash -> HmacDRG hash
    forall hash. Digest hash -> Digest hash -> HmacDRG hash
    HmacDRG Digest hash
    keyNew Digest hash
    valueNew
          where
            keyNew :: Digest hash
    keyNew =
                HMAC hash -> Digest hash
    forall a. HMAC a -> Digest a
    hmacGetDigest (HMAC hash -> Digest hash) -> HMAC hash -> Digest hash
    forall a b. (a -> b) -> a -> b
    $
                    Digest hash -> ScrubbedBytes -> HMAC hash
    forall key message a.
    (ByteArrayAccess key, ByteArrayAccess message, HashAlgorithm a) =>
    key -> message -> HMAC a
    hmac Digest hash
    key (ScrubbedBytes -> HMAC hash) -> ScrubbedBytes -> HMAC hash
    forall a b. (a -> b) -> a -> b
    $
                        Digest hash -> ScrubbedBytes
    forall bin bout.
    (ByteArrayAccess bin, ByteArray bout) =>
    bin -> bout
    M.convert Digest hash
    value ScrubbedBytes -> ScrubbedBytes -> ScrubbedBytes
    forall a. Semigroup a => a -> a -> a
    <> forall a. ByteArray a => Word8 -> a
    M.singleton @ScrubbedBytes Word8
    byte ScrubbedBytes -> ScrubbedBytes -> ScrubbedBytes
    forall a. Semigroup a => a -> a -> a
    <> input -> ScrubbedBytes
    forall bin bout.
    (ByteArrayAccess bin, ByteArray bout) =>
    bin -> bout
    M.convert input
    input
            valueNew :: Digest hash
    valueNew = HMAC hash -> Digest hash
    forall a. HMAC a -> Digest a
    hmacGetDigest (HMAC hash -> Digest hash) -> HMAC hash -> Digest hash
    forall a b. (a -> b) -> a -> b
    $ Digest hash -> Digest hash -> HMAC hash
    forall key message a.
    (ByteArrayAccess key, ByteArrayAccess message, HashAlgorithm a) =>
    key -> message -> HMAC a
    hmac Digest hash
    keyNew Digest hash
    value

    instance HashAlgorithm hash => DRG (HmacDRG hash) where
        randomBytesGenerate :: forall byteArray.
    ByteArray byteArray =>
    Int -> HmacDRG hash -> (byteArray, HmacDRG hash)
    randomBytesGenerate Int
    count (HmacDRG Digest hash
    key Digest hash
    value) = (byteArray
    output, HmacDRG hash
    state)
          where
            output :: byteArray
    output = Int -> byteArray -> byteArray
    forall bs. ByteArray bs => Int -> bs -> bs
    M.take Int
    count byteArray
    result
            state :: HmacDRG hash
    state = forall input hash.
    (ByteArrayAccess input, HashAlgorithm hash) =>
    input -> HmacDRG hash -> HmacDRG hash
    update @Bytes Bytes
    forall a. ByteArray a => a
    M.empty (HmacDRG hash -> HmacDRG hash) -> HmacDRG hash -> HmacDRG hash
    forall a b. (a -> b) -> a -> b
    $ Digest hash -> Digest hash -> HmacDRG hash
    forall hash. Digest hash -> Digest hash -> HmacDRG hash
    HmacDRG Digest hash
    key Digest hash
    new
            (byteArray
    result, Digest hash
    new) = byteArray -> Digest hash -> (byteArray, Digest hash)
    forall {t} {a}.
    (ByteArray t, HashAlgorithm a) =>
    t -> Digest a -> (t, Digest a)
    go byteArray
    forall a. ByteArray a => a
    M.empty Digest hash
    value
            go :: t -> Digest a -> (t, Digest a)
    go t
    buffer Digest a
    current
                | t -> Int
    forall ba. ByteArrayAccess ba => ba -> Int
    M.length t
    buffer Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    >= Int
    count = (t
    buffer, Digest a
    current)
                | Bool
    otherwise = t -> Digest a -> (t, Digest a)
    go (t
    buffer t -> t -> t
    forall a. Semigroup a => a -> a -> a
    <> Digest a -> t
    forall bin bout.
    (ByteArrayAccess bin, ByteArray bout) =>
    bin -> bout
    M.convert Digest a
    next) Digest a
    next
              where
                next :: Digest a
    next = HMAC a -> Digest a
    forall a. HMAC a -> Digest a
    hmacGetDigest (HMAC a -> Digest a) -> HMAC a -> Digest a
    forall a b. (a -> b) -> a -> b
    $ Digest hash -> Digest a -> HMAC a
    forall key message a.
    (ByteArrayAccess key, ByteArrayAccess message, HashAlgorithm a) =>
    key -> message -> HMAC a
    hmac Digest hash
    key Digest a
    current