{-# LANGUAGE Haskell2010       #-}

{-# LANGUAGE OverloadedStrings #-}

-- SPDX-License-Identifier: GPL-2.0-or-later

module CAVS(TestVec(..), TestVecHdr(..), EncDec(..), PT(..), parseTestVecs, hex, unhex) where

import           Control.Monad
import           Data.ByteString            (ByteString)
import qualified Data.ByteString.Base16     as B16
import qualified Data.ByteString.Char8      as BSC
import qualified Data.ByteString.Lazy.Char8 as BLC
import           Data.Char                  (isHexDigit)
import           Data.Either
import           Data.Maybe
import           Text.Read                  (readMaybe)

-- WARNING: the code that follows will make you cry;
--          a safety pig is provided below for your benefit.
--
--                          _
--  _._ _..._ .-',     _.._(`))
-- '-. `     '  /-._.-'    ',/
--    )         \            '.
--   / _    _    |             \
--  |  a    a    /              |
--  \   .-.                     ;
--   '-('' ).-'       ,'       ;
--      '-;           |      .'
--         \           \    /
--         | 7  .__  _.-\   \
--         | |  |  ``/  /`  /
--        /,_|  |   /,_/   /
--           /,_/      '`-'
--

data EncDec = Enc | Dec deriving (Show,Eq)

data PT = PT !EncDec !ByteString | P'FAIL deriving (Show,Eq)

data TestVec = TestVec
  { tv'Count :: !Int
  , tv'Key   :: !ByteString
  , tv'IV    :: !ByteString
  , tv'AAD   :: !ByteString
  , tv'CT    :: !ByteString
  , tv'Tag   :: !ByteString
  , tv'PT    :: !PT
  } deriving Show

data TestVecHdr = TestVecHdr
  { tv'Keylen :: !Int
  , tv'IVlen  :: !Int
  , tv'PTlen  :: !Int
  , tv'AADlen :: !Int
  , tv'Taglen :: !Int
  } deriving Show

mkTestVec :: Int -> ByteString -> ByteString -> ByteString -> ByteString -> ByteString -> PT -> TestVec
mkTestVec count key0 iv0 aad0 ct0 tag0 pt0_ = TestVec count key iv aad ct tag pt_
  where
    pt0 = case pt0_ of { PT _ x -> x; P'FAIL -> mempty }
    merged0 = mconcat [key0, iv0, aad0, ct0, tag0, pt0]

    keyLen = BSC.length key0 `div` 2
    ivLen  = BSC.length iv0  `div` 2
    aadLen = BSC.length aad0 `div` 2
    ctLen  = BSC.length ct0  `div` 2
    tagLen = BSC.length tag0 `div` 2
    ptLen  = BSC.length pt0  `div` 2

    merged = unhex merged0
    slice ofs len
      | len > 0   = BSC.take len $ BSC.drop ofs merged
      | otherwise = BSC.empty

    key = slice 0 keyLen
    iv  = slice keyLen ivLen
    aad = slice (keyLen+ivLen) aadLen
    ct  = slice (keyLen+ivLen+aadLen) ctLen
    tag = slice (keyLen+ivLen+aadLen+ctLen) tagLen
    pt  = slice (keyLen+ivLen+aadLen+ctLen+tagLen) ptLen

    pt_ = case pt0_ of
            P'FAIL -> P'FAIL
            PT d _ -> PT d pt

parseTestVecs :: BLC.ByteString -> [(TestVecHdr,[TestVec])]
parseTestVecs = token2tvecs . tokenize

token2tvecs :: [Token] -> [(TestVecHdr,[TestVec])]
token2tvecs = groupEither . go0 . dedupEmpty . filter (not . isComment)
  where
    go0 [] = []
    go0 (Token'empty:Token'empty:rest) = go0 (Token'empty:rest)
    go0 (Token'empty:Token'parmLen Parm'Key key:Token'parmLen Parm'IV iv:Token'parmLen Parm'PT pt : Token'parmLen Parm'AAD aad:Token'parmLen Parm'Tag tag:Token'empty:rest) = Left (TestVecHdr key iv pt aad tag) : go1 rest
    go0 xs = error ("parseTestVecs/go0: " ++ show (take 10 xs))

    go1 [] = []
    go1 (Token'count cnt : Token'parmVal Parm'Key key:Token'parmVal Parm'IV iv:Token'parmVal Parm'CT ct : Token'parmVal Parm'AAD aad:Token'parmVal Parm'Tag tag:tok:Token'empty:rest)
      | Token'parmVal Parm'PT pt <- tok = (Right $! (mkTestVec cnt key iv aad ct tag (PT Dec pt))) : go1' rest
      | Token'fail               <- tok = (Right $! (mkTestVec cnt key iv aad ct tag P'FAIL))      : go1' rest
    go1 (Token'count cnt : Token'parmVal Parm'Key key:Token'parmVal Parm'IV iv:Token'parmVal Parm'PT pt : Token'parmVal Parm'AAD aad:Token'parmVal Parm'CT ct : Token'parmVal Parm'Tag tag:Token'empty:rest)
                                        = (Right $! (mkTestVec cnt key iv aad ct tag (PT Enc pt))) : go1' rest
    go1 xs = error ("parseTestVecs/go1: " ++ show (take 10 xs))

    go1' toks@(Token'parmLen _ _ :_) = go0 (Token'empty:toks)
    go1' toks                        = go1 toks

    dedupEmpty (Token'empty : rest@(Token'empty : _)) = dedupEmpty rest
    dedupEmpty [Token'empty]                          = [Token'empty]
    dedupEmpty (tok:rest)                             = tok : dedupEmpty rest
    dedupEmpty []                                     = [Token'empty]

groupEither :: [Either l r] -> [(l,[r])]
groupEither [] = []
groupEither (Left l : xs) = (l, [r | Right r <- rs]) : groupEither rest
  where
    (rs, rest) = span isRight xs
groupEither (Right _ : _) = error "groupEither: leading Rights"

data Parm = Parm'Key
          | Parm'IV
          | Parm'PT
          | Parm'AAD
          | Parm'CT
          | Parm'Tag
          deriving (Show,Eq)

bs2parm :: ByteString -> Maybe Parm
bs2parm "Key" = Just Parm'Key
bs2parm "IV"  = Just Parm'IV
bs2parm "PT"  = Just Parm'PT
bs2parm "AAD" = Just Parm'AAD
bs2parm "CT"  = Just Parm'CT
bs2parm "Tag" = Just Parm'Tag
bs2parm _     = Nothing


data Token = Token'empty -- empty line
           | Token'comment !ByteString
           | Token'parmLen Parm !Int
           | Token'parmVal Parm !ByteString
           | Token'fail
           | Token'count !Int
           | Token'UNPARSED !ByteString
           deriving (Show,Eq)

isComment :: Token -> Bool
isComment (Token'comment _) = True
isComment _                 = False

tokenize :: BLC.ByteString -> [Token]
tokenize = map (tokLine . BLC.toStrict) . BLC.lines
  where
    tokLine l0
      | Just c <- BSC.stripPrefix "#" l = Token'comment c
      | l == "" = Token'empty
      | l == "FAIL" = Token'fail
      | Just (k,v) <- decodeParmLen l = Token'parmLen k v
      | Just (k,v) <- decodeParmVal l = Token'parmVal k v
      | Just x <- decodeCount l = Token'count x
      | otherwise = Token'UNPARSED l0
      where
        l = fromMaybe l0 $ BSC.stripSuffix "\r" l0

    decodeCount l0 = do
      [ "Count", "=", v ] <- Just $ BSC.words l0
      readMaybe (BSC.unpack v)

    decodeParmLen :: ByteString -> Maybe (Parm,Int)
    decodeParmLen l0 = do
      l1 <- BSC.stripPrefix "[" l0
      l2 <- BSC.stripSuffix "]" l1

      [ k1, "=", v ] <- Just $ BSC.words l2

      k2 <- BSC.stripSuffix "len" k1
      k3 <- bs2parm k2

      v' <- readMaybe (BSC.unpack v)

      pure (k3,v')

    decodeParmVal :: ByteString -> Maybe (Parm,ByteString)
    decodeParmVal l0 = do
      (k1,v2) <- case BSC.words l0 of
        [ k1, "=", v1 ] -> Just (k1,v1)
        [ k1, "=" ]     -> Just (k1,"")
        _               -> Nothing

      k2 <- bs2parm k1

      guard (even $ BSC.length v2)
      guard (BSC.all isHexDigit v2)

      pure (k2,v2)

unhex :: ByteString -> ByteString
unhex = either (\err -> error ("mkTestVec/unhex: " ++ err)) id . B16.decode

hex :: ByteString -> ByteString
hex = B16.encode