-- CryptoSEIPDv2.hs: OpenPGP (RFC9580) SEIPD-v2 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.Internal.CryptoSEIPDv2
  ( aeadModeAndNonceSizeForSEIPDv2
  , seipdv2SymmetricKeySize
  , deriveSKESK6KEK
  , encryptSKESK6SessionKey
  , decryptSKESK6SessionKey
  ) where

import Codec.Encryption.OpenPGP.Internal.CryptoAES (withAESCipher)
import Codec.Encryption.OpenPGP.Internal.RFC7253OCB
  ( decryptWithOCBRFC7253With
  , encryptWithOCBRFC7253
  )
import Codec.Encryption.OpenPGP.Types
import qualified "crypton" Crypto.Cipher.Types as CCT
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

aeadModeAndNonceSizeForSEIPDv2 ::
     String -> AEADAlgorithm -> Either String (CCT.AEADMode, Int)
aeadModeAndNonceSizeForSEIPDv2 otherAeadError EAX =
  Left "EAX is currently unsupported by the crypton AEAD backend"
aeadModeAndNonceSizeForSEIPDv2 otherAeadError OCB = Right (CCT.AEAD_OCB, 15)
aeadModeAndNonceSizeForSEIPDv2 otherAeadError GCM = Right (CCT.AEAD_GCM, 12)
aeadModeAndNonceSizeForSEIPDv2 otherAeadError (OtherAEADAlgo _) = Left otherAeadError

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

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

deriveSKESK6KEK ::
     SymmetricAlgorithm
  -> AEADAlgorithm
  -> B.ByteString
  -> Either String B.ByteString
deriveSKESK6KEK symalgo aead ikm = do
  keyLen <-
    seipdv2SymmetricKeySize
      "SKESK v6 currently supports AES-128/192/256 only"
      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 String (B.ByteString, B.ByteString)
encryptSKESK6SessionKey symalgo aead kek iv sessionKey = do
  (mode, nonceSize) <-
    aeadModeAndNonceSizeForSEIPDv2
      "unsupported AEAD algorithm for SKESK v6 encrypt"
      aead
  if B.length iv /= nonceSize
    then Left "SKESK v6 IV length does not match AEAD algorithm"
    else
      withAESCipher
        "SKESK v6 encrypt currently supports AES-128/192/256 only"
        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 show . 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 String B.ByteString
decryptSKESK6SessionKey symalgo aead kek iv ciphertext tag = do
  (mode, nonceSize) <-
    aeadModeAndNonceSizeForSEIPDv2
      "unsupported AEAD algorithm for SKESK v6 decrypt"
      aead
  if B.length iv /= nonceSize
    then Left "SKESK v6 IV length does not match AEAD algorithm"
    else
      withAESCipher
        "SKESK v6 decrypt currently supports AES-128/192/256 only"
        symalgo
        kek
        (\cipher ->
           if mode == CCT.AEAD_OCB
             then
               decryptWithOCBRFC7253With
                 (\_ _ _ _ _ _ -> "SKESK v6 authentication failed")
                 cipher
                 iv
                 (skeskV6Info symalgo aead)
                 ciphertext
                 (mkAuthTag tag)
             else do
               aeadCtx <- first show . CE.eitherCryptoError $ CCT.aeadInit mode cipher iv
               case CCT.aeadSimpleDecrypt aeadCtx (skeskV6Info symalgo aead) ciphertext (mkAuthTag tag) of
                 Nothing -> Left "SKESK v6 authentication failed"
                 Just plain -> Right plain)

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

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