{-# 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