-- CFB.hs: OpenPGP (RFC9580) CFB mode
-- Copyright © 2013-2026  Clint Adams
-- Copyright © 2013  Daniel Kahn Gillmor
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}

module Codec.Encryption.OpenPGP.CFB
  ( decrypt
  , decryptPreservingNonce
  , decryptNoNonce
  , decryptOpenPGPCfb
  , decryptOpenPGPCfbWithNonce
  , OpenPGPCFBMode(..)
  , OpenPGPCFBModeW(..)
  , encryptNoNonce
  , encryptOpenPGPCfbRaw
  , mdcTrailerForSEIPDv1
  , seipdv1NonceFromIV
  , validateSEIPD1MDC
  , calculateMDC
  ) where

import Codec.Encryption.OpenPGP.BlockCipher (CipherError(..), withSymmetricCipher)
import Codec.Encryption.OpenPGP.Internal.HOBlockCipher
import Codec.Encryption.OpenPGP.Types
import qualified Crypto.Hash as CH
import qualified Crypto.Hash.Algorithms as CHA
import qualified Data.ByteArray as BA
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import Control.Monad (when)

data OpenPGPCFBMode
  = OpenPGPCFBResync
  | OpenPGPCFBNoResync
  deriving (Eq, Show)

data OpenPGPCFBModeW (mode :: OpenPGPCFBMode) where
  OpenPGPCFBResyncW :: OpenPGPCFBModeW 'OpenPGPCFBResync
  OpenPGPCFBNoResyncW :: OpenPGPCFBModeW 'OpenPGPCFBNoResync

decryptOpenPGPCfb ::
     SymmetricAlgorithm
  -> B.ByteString
  -> B.ByteString
  -> Either CipherError B.ByteString
decryptOpenPGPCfb sa ciphertext keydata =
  snd <$> decryptOpenPGPCfbWithNonce sa ciphertext keydata

decryptOpenPGPCfbWithNonce ::
     SymmetricAlgorithm
  -> B.ByteString
  -> B.ByteString
  -> Either CipherError (B.ByteString, B.ByteString)
decryptOpenPGPCfbWithNonce Plaintext ciphertext _ = return (mempty, ciphertext)
decryptOpenPGPCfbWithNonce sa ciphertext keydata =
  withSymmetricCipher sa keydata $ \bc -> do
    nonce <- decrypt1 ciphertext bc
    cleartext <- decrypt2 ciphertext bc
    if nonceCheck bc nonce
      then return (nonce, cleartext)
      else Left "Session key quickcheck failed"
  where
    decrypt1 ::
         HOBlockCipher cipher
      => B.ByteString
      -> cipher
      -> Either String B.ByteString
    decrypt1 ct cipher =
      paddedCfbDecrypt
        cipher
        (B.replicate (blockSize cipher) 0)
        (B.take (blockSize cipher + 2) ct)
    decrypt2 ::
         HOBlockCipher cipher
      => B.ByteString
      -> cipher
      -> Either String B.ByteString
    decrypt2 ct cipher =
      let i = B.take (blockSize cipher) (B.drop 2 ct)
       in paddedCfbDecrypt cipher i (B.drop (blockSize cipher + 2) ct)

-- should deprecate this?
decrypt ::
     SymmetricAlgorithm
  -> B.ByteString
  -> B.ByteString
  -> Either CipherError B.ByteString
decrypt x y z = snd <$> (decryptPreservingNonce x y z)

decryptPreservingNonce ::
     SymmetricAlgorithm
  -> B.ByteString
  -> B.ByteString
  -> Either CipherError (B.ByteString, B.ByteString)
decryptPreservingNonce Plaintext ciphertext _ = return (mempty, ciphertext)
decryptPreservingNonce sa ciphertext keydata =
  withSymmetricCipher sa keydata $ \bc -> do
    let bs = blockSize bc
    decrypted <- paddedCfbDecrypt bc (B.replicate bs 0) ciphertext
    let (nonce, cleartext) = B.splitAt (bs + 2) decrypted
    if nonceCheck bc nonce
      then return (nonce, cleartext)
      else Left "Session key quickcheck failed"

decryptNoNonce ::
     SymmetricAlgorithm
  -> IV
  -> B.ByteString
  -> B.ByteString
  -> Either CipherError B.ByteString
decryptNoNonce Plaintext _ ciphertext _ = return ciphertext
decryptNoNonce sa iv ciphertext keydata =
  withSymmetricCipher sa keydata (decrypt' ciphertext)
  where
    decrypt' ::
         HOBlockCipher cipher
      => B.ByteString
      -> cipher
      -> Either String B.ByteString
    decrypt' ct cipher = paddedCfbDecrypt cipher (unIV iv) ct

nonceCheck :: HOBlockCipher cipher => cipher -> B.ByteString -> Bool
nonceCheck bc =
  (==) <$> B.take 2 . B.drop (blockSize bc - 2) <*> B.drop (blockSize bc)

encryptNoNonce ::
     SymmetricAlgorithm
  -> S2K
  -> IV
  -> B.ByteString
  -> B.ByteString
  -> Either CipherError B.ByteString
encryptNoNonce Plaintext _ _ payload _ = return payload
encryptNoNonce sa s2k iv payload keydata =
  withSymmetricCipher sa keydata (encrypt' payload)
  where
    encrypt' ::
         HOBlockCipher cipher
      => B.ByteString
      -> cipher
      -> Either String B.ByteString
    encrypt' ct cipher = paddedCfbEncrypt cipher (unIV iv) ct

encryptOpenPGPCfbRaw ::
    OpenPGPCFBModeW mode
  -> SymmetricAlgorithm
  -> IV
  -> B.ByteString  -- ^ plaintext (with MDC trailer already appended)
  -> B.ByteString  -- ^ raw session key bytes
  -> Either CipherError B.ByteString
encryptOpenPGPCfbRaw _ Plaintext _ cleartext _ = Right cleartext
encryptOpenPGPCfbRaw mode sa iv cleartext keydata =
  withSymmetricCipher sa keydata $ \cipher -> do
    let initialVector = unIV iv
        bs            = blockSize cipher
    if B.length initialVector /= bs
      then Left
             ("IV length mismatch for " ++
              show sa ++
              ": expected " ++ show bs ++ ", got " ++ show (B.length initialVector))
      else do
        let prefix = initialVector <> B.drop (bs - 2) initialVector
        case mode of
          OpenPGPCFBResyncW -> do
            nonceAndCheck <- paddedCfbEncrypt cipher (B.replicate bs 0) prefix
            encryptedPayload <-
              paddedCfbEncrypt cipher (B.take bs (B.drop 2 nonceAndCheck)) cleartext
            return (nonceAndCheck <> encryptedPayload)
          OpenPGPCFBNoResyncW ->
            paddedCfbEncrypt cipher (B.replicate bs 0) (prefix <> cleartext)

-- | Compute the MDC trailer appended to SEIPDv1 plaintext before encryption.
-- The trailer is: @0xd3 0x14 SHA1(nonce || plaintext || 0xd3 0x14)@.
mdcTrailerForSEIPDv1 :: IV -> B.ByteString -> B.ByteString
mdcTrailerForSEIPDv1 iv plaintext = mdcHeader <> digest
  where
    mdcHeader = B.pack [0xd3, 0x14]
    nonce     = seipdv1NonceFromIV iv
    digest    = BA.convert (CH.hash (nonce <> plaintext <> mdcHeader) :: CH.Digest CHA.SHA1)

-- | The SEIPDv1 nonce: the IV bytes followed by its last two bytes (resync prefix).
seipdv1NonceFromIV :: IV -> B.ByteString
seipdv1NonceFromIV (IV ivBytes) = ivBytes <> B.drop (B.length ivBytes - 2) ivBytes

calculateMDC :: B.ByteString -> B.ByteString -> Maybe BL.ByteString
calculateMDC nonce garbage
  | B.length garbage < 23 = Nothing
  | otherwise =
    let digest = CH.hash (nonce <> B.take (B.length garbage - 22) garbage <> B.pack [211, 20]) :: CH.Digest CHA.SHA1
     in Just (BL.fromStrict (BA.convert digest :: B.ByteString))

-- | Verify the MDC trailer of a decrypted SEIPDv1 payload.
-- Takes the CFB nonce (blockSize+2 prefix bytes retained from decryption)
-- and the full decrypted bytes (payload + MDC packet), and returns the
-- payload without the MDC trailer on success.
validateSEIPD1MDC :: B.ByteString -> B.ByteString -> Either String B.ByteString
validateSEIPD1MDC nonce decrypted = do
  when (B.length decrypted < 22) $
    Left "SEIPD1 cleartext too short to contain MDC trailer"
  let (payload, trailer) = B.splitAt (B.length decrypted - 22) decrypted
  when (B.take 2 trailer /= B.pack [211, 20]) $
    Left "SEIPD1 cleartext missing MDC packet trailer (tag 19)"
  expectedMdc <-
    case calculateMDC nonce decrypted of
     Nothing -> Left "SEIPD1 cleartext too short for MDC calculation"
     Just x -> Right x
  let actualMdc = BL.fromStrict (B.drop 2 trailer)
  when (expectedMdc /= actualMdc) $
    Left "MDC indicates tampering"
  Right payload
