-- 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.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

canonicalizeLineEndings :: ByteString -> ByteString
canonicalizeLineEndings = BL.pack . go . BL.unpack
  where
    go [] = []
    go (0x0d:0x0a:rest) = 0x0d : 0x0a : go rest
    go (0x0d:rest) = 0x0d : 0x0a : go rest
    go (0x0a:rest) = 0x0d : 0x0a : go rest
    go (w:rest) = w : go rest

stripTrailingWhitespacePerLine :: ByteString -> ByteString
stripTrailingWhitespacePerLine = BL.pack . go [] . BL.unpack
  where
    go lineRev [] = reverseTrimmed lineRev
    go lineRev (0x0d:0x0a:rest) =
      reverseTrimmed lineRev ++ [0x0d, 0x0a] ++ go [] rest
    go lineRev (w:rest) = go (w : lineRev) rest

    reverseTrimmed :: [Word8] -> [Word8]
    reverseTrimmed = reverse . dropWhile isTrailingWhitespace

    isTrailingWhitespace :: Word8 -> Bool
    isTrailingWhitespace w = w == 0x20 || w == 0x09

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