{-# LANGUAGE Safe #-}
{-|
  Module        : KAT
  Description   : Known Answer Tests
  Copyright     : © Jeremy Bornstein 2019
  License       : Apache 2.0
  Maintainer    : jeremy@bornstein.org
  Stability     : experimental
  Portability   : portable

  KAT stands for Known Answer Test(s). The validity of an
  implementation may be judged in part by its ability to produce the
  correct output.

  This module generates test vectors (according to the NIST PQC spec)
  which should be identical (in filenames and contents) to the ones
  generated by the original NewHope reference code.

  The automated tests in this codebase contain tests which take those
  reference implementation vectors as input to verify that we generate
  the same data with the below code.

-}

module KAT where

import qualified Data.ByteString            as BS
import           Data.ByteString.Builder
import qualified Data.ByteString.Char8      as BSC
import qualified Data.ByteString.Lazy.Char8 as BSLC
import           Data.Char
import           Data.Semigroup             ((<>))

import           AuxUtil
import qualified Crypto.NewHope                  as NewHope
import qualified Crypto.NewHope.Internal.CCA_KEM as CCA_KEM
import qualified Crypto.NewHope.Internal.CPA_KEM as CPA_KEM
import qualified Crypto.NewHope.Internal.RNG     as RNG

type Filename = String
type VectorGenerator = NewHope.N -> Int -> (Filename, Builder)


recordsToGenerate :: Int
recordsToGenerate = 100





-- | output ByteString as uppercase hex string representation
processBS :: BS.ByteString -> String
processBS bs = toUpper <$> (BSLC.unpack . toLazyByteString . byteStringHex) bs





-- | Note that the seeds used are the same for both algorithms and both variants thereof.
getSeeds :: [RNG.RandomSeed]
getSeeds = go initialCTX
  where
    initialSeed = RNG.makeRandomSeed $ BS.pack [0 .. 47]
    initialCTX = RNG.randomBytesInit initialSeed Nothing 256
    go ctx = newSeed : go nextCTX
      where
        (newSeed, nextCTX) = (RNG.makeRandomSeed seed', ctx')
          where
            (seed', ctx') = RNG.randomBytes ctx 48


ccaKemTestVectors :: VectorGenerator
ccaKemTestVectors n count = (filename, start <> contents)
  where
    filename = "PQCkemKAT_" ++ show (CCA_KEM.secretKeyBytes n) ++ ".rsp"
    start = lazyByteString . BSLC.pack $ "# " ++ show (WrapN n) ++ "-CCAKEM\n\n"
    contents = foldr1 mappend $ outputRecord <$> Prelude.zip [0 ..] (take count getSeeds)

    outputRecord :: (Int, RNG.RandomSeed) -> Builder
    outputRecord (i, seed) = header <> record
      where
        header = lazyByteString . BSLC.pack $ "count = " ++ show i ++ "\n"
        record = testVectorsFromSeed seed

    testVectorsFromSeed :: RNG.RandomSeed -> Builder
    testVectorsFromSeed seed = seq check output'
      where
        ctx = RNG.randomBytesInit seed Nothing 256
        (pk, sk, ctx') = CCA_KEM.keypair ctx n
        (ct, ss, _) = CCA_KEM.encrypt ctx' pk
        check = let (success, ss') = CCA_KEM.decrypt ct sk
                    ssOK = success && ss == ss'
                in if not ssOK
                    then error "Shared secret could not be generated."
                    else ()

        seedData = RNG.getRandomSeedData seed
        output = (lazyByteString . BSLC.pack) <$> [ "seed = " ++ processBS seedData ++ "\n"
                                                  , "pk = " ++ processBS (CCA_KEM.getPKData pk) ++ "\n"
                                                  , "sk = " ++ processBS (CCA_KEM.getSKData sk) ++ "\n"
                                                  , "ct = " ++ processBS (CCA_KEM.getCTData ct) ++ "\n"
                                                  , "ss = " ++ processBS (CCA_KEM.getSSData ss) ++ "\n"
                                                  , "\n"
                                                  ]
        output' = foldr go (byteString $ BSC.pack "") output
          where
            go a as = a <> as


cpaKemTestVectors :: VectorGenerator
cpaKemTestVectors n count = (filename, start <> contents)
  where
    filename = "PQCkemKAT_" ++ show (CPA_KEM.secretKeyBytes n) ++ ".rsp"
    start = lazyByteString . BSLC.pack $ "# " ++ show (WrapN n) ++ "-CPAKEM\n\n"
    contents = foldr1 mappend $ outputRecord <$> Prelude.zip [0 ..] (take count getSeeds)

    outputRecord :: (Int, RNG.RandomSeed) -> Builder
    outputRecord (i, seed) = header <> record
      where
        header = lazyByteString . BSLC.pack $ "count = " ++ show i ++ "\n"
        record = testVectorsFromSeed seed

    testVectorsFromSeed :: RNG.RandomSeed -> Builder
    testVectorsFromSeed seed = seq check output'
      where
        ctx = RNG.randomBytesInit seed Nothing 256
        (pk, sk, ctx') = CPA_KEM.keypair ctx n
        (ct, ss, _) = CPA_KEM.encrypt ctx' pk
        check = let mss = CPA_KEM.decrypt ct sk
                    ssOK = ss == mss
                in if not ssOK
                   then error "Shared secret could not be generated."
                   else ()

        seedData = RNG.getRandomSeedData seed
        output = (lazyByteString . BSLC.pack) <$> [ "seed = " ++ processBS seedData ++ "\n"
                                                  , "pk = " ++ processBS (CPA_KEM.getPKData pk) ++ "\n"
                                                  , "sk = " ++ processBS (CPA_KEM.getSKData sk) ++ "\n"
                                                  , "ct = " ++ processBS (CPA_KEM.getCTData ct) ++ "\n"
                                                  , "ss = " ++ processBS (CPA_KEM.getSSData ss) ++ "\n"
                                                  , "\n"
                                                  ]
        output' = foldr go (byteString $ BSC.pack "") output
          where
            go a as = a <> as