-- SerializeForSigs.hs: OpenPGP (RFC9580) special serialization for signature purposes
-- Copyright © 2012-2026  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}

module Codec.Encryption.OpenPGP.SerializeForSigs
    ( putPKPforFingerprinting
    , putPartialSigforSigning
    , putSigTrailer
    , putUforSigning
    , putUIDforSigning
    , putUAtforSigning
    , putKeyforSigning
    , putSigforSigning
    , payloadForSig
    , payloadForSigWith
    ) where

import Control.Lens ((^.))
import Crypto.Number.Serialize (i2osp)
import Data.Binary (put)
import Data.Binary.Put
    ( Put
    , putByteString
    , putLazyByteString
    , putWord16be
    , putWord32be
    , putWord8
    , runPut
    )
import qualified Data.ByteString as B
import Data.ByteString.Lazy (ByteString)
import qualified Data.ByteString.Lazy as BL
import Data.Text.Encoding (encodeUtf8)
import Data.Word (Word8)

import Codec.Encryption.OpenPGP.Internal
    ( PktStreamContext (..)
    , pubkeyToMPIs
    )
import Codec.Encryption.OpenPGP.Internal.Whitespace
    ( canonicalizeLineEndings
    , stripTrailingWhitespacePerLine
    )
import Codec.Encryption.OpenPGP.Serialize ()
import Codec.Encryption.OpenPGP.Subpackets
    ( TextNormalizationMode (..)
    )
import Codec.Encryption.OpenPGP.Types

data SignatureSerializationCase where
    SignatureSerializationCaseV4
        :: SignaturePayloadV 'SigPayloadV4 -> SignatureSerializationCase
    SignatureSerializationCaseV6
        :: SignaturePayloadV 'SigPayloadV6 -> SignatureSerializationCase

fromPktSignatureSerializationCase
    :: Pkt -> Maybe SignatureSerializationCase
fromPktSignatureSerializationCase pkt =
    case fromPktEitherSomeSignatureV pkt of
        Right (SomeSignatureV (SignatureV4Packet payload)) ->
            Just (SignatureSerializationCaseV4 payload)
        Right (SomeSignatureV (SignatureV6Packet payload)) ->
            Just (SignatureSerializationCaseV6 payload)
        _ -> Nothing

putPartialSigforSigningCase :: SignatureSerializationCase -> Put
putPartialSigforSigningCase
    ( SignatureSerializationCaseV4
            (SigPayloadV4Data st pka ha hashed _ _ _)
        ) = do
        putWord8 4
        put st
        put pka
        put ha
        let hb = runPut $ mapM_ put hashed
        putWord16be . fromIntegral . BL.length $ hb
        putLazyByteString hb
putPartialSigforSigningCase
    ( SignatureSerializationCaseV6
            (SigPayloadV6Data st pka ha _salt hashed _ _ _)
        ) = do
        putWord8 6
        put st
        put pka
        put ha
        let hb = runPut $ mapM_ put hashed
        putWord32be . fromIntegral . BL.length $ hb
        putLazyByteString hb

putSigTrailerCase :: SignatureSerializationCase -> Put
putSigTrailerCase (SignatureSerializationCaseV4 (SigPayloadV4Data _ _ _ hs _ _ _)) = do
    putWord8 0x04
    putWord8 0xff
    putWord32be . fromIntegral . (+ 6) . BL.length $
        runPut $
            mapM_ put hs
-- this +6 seems like a bug in RFC4880
putSigTrailerCase signatureCase@(SignatureSerializationCaseV6 _) = do
    putWord8 0x06
    putWord8 0xff
    putWord32be . fromIntegral . (+ 6) . BL.length $
        runPut (putPartialSigforSigningCase signatureCase)

putSigforSigningCase :: SignatureSerializationCase -> Put
putSigforSigningCase
    ( SignatureSerializationCaseV4
            (SigPayloadV4Data st pka ha hashed _ left16 mpis)
        ) = do
        putWord8 0x88
        let bs = runPut $ put (SigV4 st pka ha hashed [] left16 mpis)
        putWord32be . fromIntegral . BL.length $ bs
        putLazyByteString bs
putSigforSigningCase
    ( SignatureSerializationCaseV6
            (SigPayloadV6Data st pka ha salt hashed _ left16 mpis)
        ) = do
        putWord8 0xC2
        let bs = runPut $ put (SigV6 st pka ha salt hashed [] left16 mpis)
        putWord32be . fromIntegral . BL.length $ bs
        putLazyByteString bs

putPKPforFingerprinting :: Pkt -> Put
putPKPforFingerprinting (PublicKeyPkt (PKPayload DeprecatedV3 _ _ _ pk)) =
    mapM_ putMPIforFingerprinting (pubkeyToMPIs pk)
putPKPforFingerprinting (PublicKeyPkt pkp@(PKPayload V4 _ _ _ _)) = do
    putWord8 0x99
    let bs = runPut $ put pkp
    putWord16be . fromIntegral $ BL.length bs
    putLazyByteString bs
putPKPforFingerprinting (PublicKeyPkt pkp@(PKPayload V6 _ _ _ _)) = do
    putWord8 0x9B
    let bs = runPut $ put pkp
    putWord32be . fromIntegral $ BL.length bs
    putLazyByteString bs
putPKPforFingerprinting _ =
    error "This should never happen (putPKPforFingerprinting)"

putMPIforFingerprinting :: MPI -> Put
putMPIforFingerprinting (MPI i) =
    let bs = i2osp i
     in putByteString bs

putPartialSigforSigning :: Pkt -> Put
putPartialSigforSigning pkt =
    case fromPktSignatureSerializationCase pkt of
        Just signatureCase -> putPartialSigforSigningCase signatureCase
        Nothing ->
            error
                ( "putPartialSigforSigning: unsupported signature packet version: "
                    ++ show (pktTag pkt)
                )

putSigTrailer :: Pkt -> Put
putSigTrailer pkt =
    case fromPktSignatureSerializationCase pkt of
        Just signatureCase -> putSigTrailerCase signatureCase
        Nothing -> error "This should never happen (putSigTrailer)"

putUforSigning :: Pkt -> Put
putUforSigning u@(UserIdPkt _) = putUIDforSigning u
putUforSigning u@(UserAttributePkt _) = putUAtforSigning u
putUforSigning _ = error "This should never happen (putUforSigning)"

putUIDforSigning :: Pkt -> Put
putUIDforSigning (UserIdPkt u) = do
    putWord8 0xB4
    let bs = encodeUtf8 u
    putWord32be . fromIntegral . B.length $ bs
    putByteString bs
putUIDforSigning _ = error "This should never happen (putUIDforSigning)"

putUAtforSigning :: Pkt -> Put
putUAtforSigning (UserAttributePkt us) = do
    putWord8 0xD1
    let bs = runPut (mapM_ put us)
    putWord32be . fromIntegral . BL.length $ bs
    putLazyByteString bs
putUAtforSigning _ = error "This should never happen (putUAtforSigning)"

putSigforSigning :: Pkt -> Put
putSigforSigning pkt =
    case fromPktSignatureSerializationCase pkt of
        Just signatureCase -> putSigforSigningCase signatureCase
        Nothing ->
            error
                ( "putSigforSigning: unsupported signature packet version: "
                    ++ show (pktTag pkt)
                )

putKeyforSigning :: Pkt -> Put
putKeyforSigning (PublicKeyPkt pkp) = putKeyForSigning' pkp
putKeyforSigning (PublicSubkeyPkt pkp) = putKeyForSigning' pkp
putKeyforSigning (SecretKeyPkt pkp _) = putKeyForSigning' pkp
putKeyforSigning (SecretSubkeyPkt pkp _) = putKeyForSigning' pkp
putKeyforSigning x =
    error
        ( "This should never happen (putKeyforSigning) "
            ++ show (pktTag x)
            ++ "/"
            ++ show x
        )

putKeyForSigning' :: SomePKPayload -> Put
putKeyForSigning' pkp@(PKPayload V6 _ _ _ _) = do
    putWord8 0x9B
    let bs = runPut $ put pkp
    putWord32be . fromIntegral . BL.length $ bs
    putLazyByteString bs
putKeyForSigning' pkp = do
    putWord8 0x99
    let bs = runPut $ put pkp
    putWord16be . fromIntegral . BL.length $ bs
    putLazyByteString bs

payloadForSig :: SigType -> PktStreamContext -> ByteString
payloadForSig BinarySig state =
    case (fromPktEither (lastLD state) :: Either String LiteralData) of
        Right ld -> ld ^. literalDataPayload
        Left err -> error ("payloadForSig expected literal data packet: " ++ err)
payloadForSig CanonicalTextSig state =
    stripTrailingWhitespacePerLine
        (canonicalizeLineEndings (payloadForSig BinarySig state))
payloadForSig StandaloneSig _ = BL.empty
payloadForSig GenericCert state =
    kandUPayload (lastPrimaryKey state) (lastUIDorUAt state)
payloadForSig PersonaCert state = payloadForSig GenericCert state
payloadForSig CasualCert state = payloadForSig GenericCert state
payloadForSig PositiveCert state = payloadForSig GenericCert state
payloadForSig SubkeyBindingSig state =
    kandKPayload (lastPrimaryKey state) (lastSubkey state)
payloadForSig PrimaryKeyBindingSig state =
    kandKPayload (lastPrimaryKey state) (lastSubkey state)
payloadForSig SignatureDirectlyOnAKey state =
    runPut (putKeyforSigning (lastPrimaryKey state))
payloadForSig KeyRevocationSig state =
    payloadForSig SignatureDirectlyOnAKey state
payloadForSig SubkeyRevocationSig state =
    kandKPayload (lastPrimaryKey state) (lastSubkey state)
payloadForSig CertRevocationSig state =
    -- RFC 9580 §5.2.1: 0x30 revokes a UID certification when a UID/UAt is in
    -- scope, but when there is no UID/UAt in scope it revokes a direct-key sig
    -- (0x1F) and the payload is just the primary key material.
    case lastUIDorUAt state of
        UserIdPkt _ ->
            kandUPayload (lastPrimaryKey state) (lastUIDorUAt state)
        UserAttributePkt _ -> kandUPayload (lastPrimaryKey state) (lastUIDorUAt state)
        _ ->
            runPut (putKeyforSigning (lastPrimaryKey state))
payloadForSig st _ = error ("payloadForSig: unhandled signature type " ++ show st)

{- | Like 'payloadForSig' but accepts an explicit 'TextNormalizationMode'
that controls whether trailing whitespace is stripped for CanonicalTextSig.

Use 'RFC9580Strict' for inline type 0x01 document signatures.
Use 'CleartextCompat' (or 'payloadForSig') for cleartext-armored messages.
-}
payloadForSigWith
    :: TextNormalizationMode
    -> SigType
    -> PktStreamContext
    -> ByteString
payloadForSigWith RFC9580Strict CanonicalTextSig state =
    canonicalizeLineEndings (payloadForSig BinarySig state)
payloadForSigWith _ st state = payloadForSig st state

kandUPayload :: Pkt -> Pkt -> ByteString
kandUPayload k u = runPut (sequence_ [putKeyforSigning k, putUforSigning u])

kandKPayload :: Pkt -> Pkt -> ByteString
kandKPayload k1 k2 =
    runPut (sequence_ [putKeyforSigning k1, putKeyforSigning k2])
