{-# LANGUAGE OverloadedStrings #-}

module RoundTrips (tests) where

import Distribution.TestSuite
import Base (simpleTest)
import Data.ByteString (ByteString)
import qualified Data.ByteString.Lazy as LBS
import qualified Data.ByteString.Lazy.Char8 as C8
import qualified Data.ByteString.Char8 as SC8

import Crypto.Error
import Crypto.Chunked
import Crypto.ChunkedAlgorithms

tests :: IO [Test]
tests = return [
  simpleTest "Check MAC size" checkMacSize,
  simpleTest "Good round trip" goodRoundTrip,
  simpleTest "Auth error" authError
  ]

testKey :: ByteString
testKey = SC8.pack "01234567890123456789012345678901"
testNonce :: ByteString
testNonce = SC8.pack "012345678901"
testString :: LBS.ByteString
testString = C8.pack $ unlines [
  "ChaCha20 and Poly1305 for IETF Protocols",
  "Abstract",
  "This document defines the ChaCha20 stream cipher as well as the use",
  "of the Poly1305 authenticator, both as stand-alone algorithms and as",
  "a combined mode, or Authenticated Encryption with Associated Data",
  "(AEAD) algorithm.",
  "This document does not introduce any new crypto, but is meant to",
  "serve as a stable reference and an implementation guide.  It is a",
  "product of the Crypto Forum Research Group (CFRG).",
  "Status of This Memo",
   "This document is not an Internet Standards Track specification; it is",
   "published for informational purposes.",
   "This document is a product of the Internet Research Task Force",
   "(IRTF).  The IRTF publishes the results of Internet-related research",
   "and development activities.  These results might not be suitable for",
   "deployment.  This RFC represents the consensus of the Crypto Forum",
   "Research Group of the Internet Research Task Force (IRTF).  Documents",
   "approved for publication by the IRSG are not a candidate for any",
   "level of Internet Standard; see Section 2 of RFC 5741.",
   "Information about the current status of this document, any errata,",
   "and how to provide feedback on it may be obtained at",
   "http://www.rfc-editor.org/info/rfc7539."
  ]

checkMacSize :: IO Progress
checkMacSize = do
  let algo = chunkedChaChaPoly1305 32 testKey
      (enc, fails1) = encrypt algo testNonce 0 testString
  case sequence fails1 of
    CryptoFailed e -> return . Finished . Fail $ "Crypto failed: " ++ show e
    CryptoPassed _ -> let
      nchunks = LBS.length testString `div` plainSize algo
      modchunks = LBS.length testString `mod` plainSize algo
      expsize = nchunks * encryptedSize algo + modchunks + (encryptedSize algo - plainSize algo)
      in if expsize == LBS.length enc
         then return . Finished $ Pass
         else return . Finished . Fail $ "Sizes mismatch. Original size is " ++ show (LBS.length testString) ++
              " encrypted text size is " ++ show (LBS.length enc) ++ " extimated size is " ++ show expsize

goodRoundTrip :: IO Progress
goodRoundTrip = do
  let algo = chunkedChaChaPoly1305 32 testKey
      (enc, fails1) = encrypt algo testNonce 0 testString
      (dec, fails2) = decrypt algo testNonce 0 enc
  case sequence_ fails1 >> sequence fails2 of
    CryptoPassed _ -> if dec == testString
                      then return . Finished $ Pass
                      else return . Finished . Fail $ "Text differs: " ++ C8.unpack dec
    CryptoFailed e -> return . Finished . Fail $ "Crypto failed: " ++ show e

authError :: IO Progress
authError = do
  let algo = chunkedChaChaPoly1305 32 testKey
      (enc, fails1) = encrypt algo testNonce 0 testString
  case LBS.uncons enc of
    Nothing -> return . Finished . Fail $ "Empty encrypted text!"
    Just (t, tt) -> do
      let enc' = LBS.cons (t+1) tt
          (dec, fails2) = decrypt algo testNonce 0 enc'
      case sequence_ fails1 >> sequence fails2 of
        CryptoFailed CryptoError_MacKeyInvalid -> return . Finished $ Pass
        CryptoFailed e -> return . Finished . Fail $ "Crypto failed: " ++ show e
        CryptoPassed _ -> return . Finished . Fail $ "Got decrypted text: " ++ C8.unpack dec