-- SEIPDv2.hs: OpenPGP (RFC9580) SEIPDv2 and SKESK v6 crypto helpers
-- Copyright © 2012-2026  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PackageImports #-}
{-# LANGUAGE TypeApplications #-}

module Codec.Encryption.OpenPGP.SEIPDv2
    ( SEIPDv2Failure (..)
    , aeadModeAndNonceSizeForSEIPDv2
    , seipdv2SymmetricKeySize
    , deriveSKESK6KEK
    , encryptSKESK6SessionKey
    , decryptSKESK6SessionKey
    , renderSEIPDv2Failure
    ) where

import Control.Error.Util (note)
import qualified Crypto.Error as CE
import qualified Crypto.Hash.Algorithms as CHA
import Crypto.KDF.HKDF (expand, extract)
import Data.Bifunctor (first)
import qualified Data.ByteArray as BA
import qualified Data.ByteString as B
import qualified "crypton" Crypto.Cipher.Types as CCT

import Codec.Encryption.OpenPGP.BlockCipher
    ( CipherError (..)
    , renderCipherError
    )
import Codec.Encryption.OpenPGP.Internal.CryptoAES
    ( withAESCipher
    )
import Codec.Encryption.OpenPGP.Internal.RFC7253OCB
    ( decryptWithOCBRFC7253With
    , encryptWithOCBRFC7253
    )
import Codec.Encryption.OpenPGP.S2K (S2KError, renderS2KError)
import Codec.Encryption.OpenPGP.Types

data SEIPDv2Failure
    = SEIPDv2UnsupportedAEADAlgorithm AEADAlgorithm
    | SEIPDv2UnsupportedSymmetricAlgorithm SymmetricAlgorithm
    | SEIPDv2InvalidSaltLength
    | SEIPDv2InvalidIVLength
    | SEIPDv2InvalidChunkSize
    | SEIPDv2CiphertextTooShort
    | SEIPDv2MalformedChunkLengths
    | SEIPDv2MissingFinalTag
    | SEIPDv2ChunkAuthFailed AEADAlgorithm Int
    | SEIPDv2FinalTagFailed AEADAlgorithm
    | SEIPDv2AuthFailed
    | SEIPDv2CipherInitFailed CE.CryptoError
    | SEIPDv2CipherFailed CipherError
    | SEIPDv2SessionKeyError S2KError
    deriving (Eq, Show)

renderSEIPDv2Failure :: SEIPDv2Failure -> String
renderSEIPDv2Failure (SEIPDv2UnsupportedAEADAlgorithm EAX) =
    "EAX is currently unsupported by the crypton AEAD backend"
renderSEIPDv2Failure (SEIPDv2UnsupportedAEADAlgorithm alg) =
    "unsupported AEAD algorithm: " ++ show alg
renderSEIPDv2Failure (SEIPDv2UnsupportedSymmetricAlgorithm _) =
    "SEIPD v2 encrypt currently supports AES-128/192/256 only"
renderSEIPDv2Failure SEIPDv2InvalidSaltLength =
    "SEIPD v2 salt must be exactly 32 octets"
renderSEIPDv2Failure SEIPDv2InvalidIVLength =
    "SKESK v6 IV length does not match AEAD algorithm"
renderSEIPDv2Failure SEIPDv2InvalidChunkSize =
    "SEIPD v2 chunk size octet must be between 0 and 16"
renderSEIPDv2Failure SEIPDv2CiphertextTooShort =
    "SEIPD v2 ciphertext must include at least one chunk tag and a final tag"
renderSEIPDv2Failure SEIPDv2MalformedChunkLengths =
    "SEIPD v2 malformed chunk lengths"
renderSEIPDv2Failure SEIPDv2MissingFinalTag =
    "SEIPD v2 missing final authentication tag"
renderSEIPDv2Failure (SEIPDv2ChunkAuthFailed algo chunk) =
    "AEAD chunk authentication failed for "
        ++ show algo
        ++ " at chunk "
        ++ show chunk
renderSEIPDv2Failure (SEIPDv2FinalTagFailed algo) =
    "AEAD final tag verification failed for " ++ show algo
renderSEIPDv2Failure SEIPDv2AuthFailed =
    "SKESK v6 authentication failed"
renderSEIPDv2Failure (SEIPDv2CipherInitFailed err) =
    "AEAD initialization failed: " ++ show err
renderSEIPDv2Failure (SEIPDv2CipherFailed err) =
    "AEAD/cipher operation failed: " ++ renderCipherError err
renderSEIPDv2Failure (SEIPDv2SessionKeyError err) =
    renderS2KError err

aeadModeAndNonceSizeForSEIPDv2
    :: AEADAlgorithm -> Either SEIPDv2Failure (CCT.AEADMode, Int)
aeadModeAndNonceSizeForSEIPDv2 EAX =
    Left $ SEIPDv2UnsupportedAEADAlgorithm EAX
aeadModeAndNonceSizeForSEIPDv2 OCB = Right (CCT.AEAD_OCB, 15)
aeadModeAndNonceSizeForSEIPDv2 GCM = Right (CCT.AEAD_GCM, 12)
aeadModeAndNonceSizeForSEIPDv2 (OtherAEADAlgo _) =
    Left . SEIPDv2UnsupportedAEADAlgorithm $ OtherAEADAlgo 0

seipdv2SymmetricKeySize
    :: SymmetricAlgorithm -> Either SEIPDv2Failure Int
seipdv2SymmetricKeySize symalgo =
    case symalgo of
        AES128 -> Right 16
        AES192 -> Right 24
        AES256 -> Right 32
        _ -> Left $ SEIPDv2UnsupportedSymmetricAlgorithm symalgo

skeskV6Info
    :: SymmetricAlgorithm -> AEADAlgorithm -> B.ByteString
skeskV6Info symalgo aead = B.pack [0xc3, 6, fromFVal symalgo, fromFVal aead]

deriveSKESK6KEK
    :: SymmetricAlgorithm
    -> AEADAlgorithm
    -> B.ByteString
    -> Either SEIPDv2Failure B.ByteString
deriveSKESK6KEK symalgo aead ikm = do
    keyLen <- seipdv2SymmetricKeySize symalgo
    let prk = extract @CHA.SHA256 B.empty ikm
    pure (expand @CHA.SHA256 prk (skeskV6Info symalgo aead) keyLen)

encryptSKESK6SessionKey
    :: SymmetricAlgorithm
    -> AEADAlgorithm
    -> B.ByteString
    -> B.ByteString
    -> B.ByteString
    -> Either SEIPDv2Failure (B.ByteString, B.ByteString)
encryptSKESK6SessionKey symalgo aead kek iv sessionKey = do
    (mode, nonceSize) <- aeadModeAndNonceSizeForSEIPDv2 aead
    if B.length iv /= nonceSize
        then Left SEIPDv2InvalidIVLength
        else
            withAESCipher
                SEIPDv2CipherInitFailed
                (SEIPDv2UnsupportedSymmetricAlgorithm symalgo)
                symalgo
                kek
                ( \cipher ->
                    if mode == CCT.AEAD_OCB
                        then do
                            (tag, ciphertext) <-
                                encryptWithOCBRFC7253
                                    cipher
                                    iv
                                    (skeskV6Info symalgo aead)
                                    sessionKey
                            pure (ciphertext, authTagToBS tag)
                        else do
                            aeadCtx <-
                                first SEIPDv2CipherInitFailed . CE.eitherCryptoError $
                                    CCT.aeadInit mode cipher iv
                            let (tag, ciphertext) =
                                    CCT.aeadSimpleEncrypt
                                        aeadCtx
                                        (skeskV6Info symalgo aead)
                                        sessionKey
                                        16
                            pure (ciphertext, authTagToBS tag)
                )

decryptSKESK6SessionKey
    :: SymmetricAlgorithm
    -> AEADAlgorithm
    -> B.ByteString
    -> B.ByteString
    -> B.ByteString
    -> B.ByteString
    -> Either SEIPDv2Failure B.ByteString
decryptSKESK6SessionKey symalgo aead kek iv ciphertext tag = do
    (mode, nonceSize) <- aeadModeAndNonceSizeForSEIPDv2 aead
    if B.length iv /= nonceSize
        then Left SEIPDv2InvalidIVLength
        else
            withAESCipher
                SEIPDv2CipherInitFailed
                (SEIPDv2UnsupportedSymmetricAlgorithm symalgo)
                symalgo
                kek
                ( \cipher ->
                    if mode == CCT.AEAD_OCB
                        then
                            decryptWithOCBRFC7253With
                                (\_ _ _ _ _ _ -> SEIPDv2AuthFailed)
                                cipher
                                iv
                                (skeskV6Info symalgo aead)
                                ciphertext
                                (mkAuthTag tag)
                        else do
                            aeadCtx <-
                                first SEIPDv2CipherInitFailed . CE.eitherCryptoError $
                                    CCT.aeadInit mode cipher iv
                            note
                                SEIPDv2AuthFailed
                                ( CCT.aeadSimpleDecrypt
                                    aeadCtx
                                    (skeskV6Info symalgo aead)
                                    ciphertext
                                    (mkAuthTag tag)
                                )
                )

authTagToBS :: CCT.AuthTag -> B.ByteString
authTagToBS = BA.convert . CCT.unAuthTag

mkAuthTag :: B.ByteString -> CCT.AuthTag
mkAuthTag = CCT.AuthTag . BA.convert
