-- 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
    ) where

import qualified Data.ByteString as B

import Codec.Encryption.OpenPGP.BlockCipher
    ( CipherError (..)
    , withSymmetricCipher
    )
import Codec.Encryption.OpenPGP.Internal.HOBlockCipher
import Codec.Encryption.OpenPGP.Types

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)
