-- Subpackets.hs: Type-safe subpacket builders with phantom types
-- Copyright © 2012-2026  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}

module Codec.Encryption.OpenPGP.Subpackets
  ( -- * Safe subpacket construction
    SafeSubpacket
  , mkSafeSubpacket
  , safePayload
  , safeSubpacket

    -- * Critical subpacket handling
  , CriticalSubpacket
  , mkCritical
  , canBeCritical

    -- * Text normalization mode
  , TextNormalizationMode(..)

    -- * Builder API for signature composition
  , SigBuilder
  , LegalSubpacket(..)
  , legalSubpacket
  , singleLegalSub
  , consLegalSub
  , listToLegalSubs
  , sbSigType
  , sbPubKeyAlgo
  , sbHashAlgo
  , sbHashedSubs
  , sbUnhashedSubs
  , sbSalt
  , sbTextNormMode
  , buildSigV4
  , buildSigV6
  , sigBuilderInit
  , sigBuilderInitTyped
  , sigBuilderInitRuntime
  , sigBuilderInitV6
  , sigBuilderInitV6Typed
  , sigBuilderInitV6Runtime
  , addHashedSubs
  , addUnhashedSubs
  , KnownPubKeyAlgorithm(..)

    -- * Private key wrapper (algorithm-specific)
  , PrivateKeyFor(..)

    -- * Subpacket list utilities
  , HashedSubpackets
  , UnhashedSubpackets
  , consHashedSub
  , singleHashedSub
  , singleUnhashedSub
  , listToHashedSubs
  , listToUnhashedSubs
  ) where

import Data.Kind (Type)
import Data.List.NonEmpty (NonEmpty)
import Data.Proxy (Proxy(..))
import Data.Word (Word16)

import qualified Crypto.PubKey.DSA as DSA
import qualified Crypto.PubKey.ECC.ECDSA as ECDSA
import qualified Crypto.PubKey.Ed25519 as Ed25519
import qualified Crypto.PubKey.Ed448 as Ed448
import qualified Crypto.PubKey.RSA.Types as RSATypes

import Codec.Encryption.OpenPGP.Types
  ( Hashed
  , Unhashed
  , V4Sig
  , V6Sig
  , Fingerprint
  , EightOctetKeyId
  , ThirtyTwoBitTimeStamp
  , MPI
  , SigSubPacket(..)
  , SigSubPacketPayload(..)
  , SignaturePayload(..)
  , SignatureSalt
  , SigType
  , HashAlgorithm
  , SubpacketList(..)
  )
import Codec.Encryption.OpenPGP.Types.Internal.Base
  ( PubKeyAlgorithm(..)
  , IssuerFingerprintVersion(..)
  )
import Codec.Encryption.OpenPGP.Policy
  ( HashAlgoAllowedFor
  , OpenPGPRFC
  , OpenPGPRFCW(..)
  , SomeOpenPGPRFCW(..)
  , HashAlgorithmW(..)
  , SomeHashAlgorithmW(..)
  , promoteOpenPGPRFC
  , promoteHashAlgorithm
  , demoteHashAlgorithmW
  , policyForRFC
  , policyGenerationDeprecations
  , deprecatedHashAlgorithms
  )

-- | Type alias for readability: hashed subpacket list
type HashedSubpackets v = SubpacketList Hashed v

-- | Type alias for readability: unhashed subpacket list
type UnhashedSubpackets v = SubpacketList Unhashed v

-- | A subpacket that has passed basic safety checks
-- (does not guarantee it's critical-safe; use CriticalSubpacket for that)
newtype SafeSubpacket = SafeSubpacket SigSubPacket
  deriving (Eq, Ord, Show)

-- | Create a safe subpacket from a payload and criticality flag
-- Basic validation that the payload is well-formed
mkSafeSubpacket :: Bool -> SigSubPacketPayload -> Either String SafeSubpacket
mkSafeSubpacket crit payload = do
  -- Validate criticality constraints: some subpackets should never be critical
  if crit && not (canBeCritical payload)
    then Left $ "Subpacket type cannot be marked critical: " ++ show payload
    else Right (SafeSubpacket (SigSubPacket crit payload))

-- | Extract the payload from a safe subpacket
safePayload :: SafeSubpacket -> SigSubPacketPayload
safePayload (SafeSubpacket (SigSubPacket _ payload)) = payload

-- | Extract the underlying SigSubPacket from a safe subpacket
safeSubpacket :: SafeSubpacket -> SigSubPacket
safeSubpacket (SafeSubpacket ssp) = ssp

-- | A subpacket that is guaranteed to be both well-formed AND can be safely marked critical
-- Used to prevent accidentally marking non-critical-safe types as critical
newtype CriticalSubpacket = CriticalSubpacket SigSubPacket
  deriving (Eq, Ord, Show)

-- | Create a critical subpacket only if the payload type allows it
-- RFC9580 section 5.2.3.5: only certain subpacket types can be marked critical
mkCritical :: SigSubPacketPayload -> Either String CriticalSubpacket
mkCritical payload
  | canBeCritical payload = Right $ CriticalSubpacket (SigSubPacket True payload)
  | otherwise = Left $ "Cannot mark as critical: " ++ payloadType payload
  where
    payloadType (SigCreationTime _) = "SigCreationTime"
    payloadType (SigExpirationTime _) = "SigExpirationTime"
    payloadType (ExportableCertification _) = "ExportableCertification"
    payloadType (TrustSignature {}) = "TrustSignature"
    payloadType (RegularExpression _) = "RegularExpression"
    payloadType (Revocable _) = "Revocable"
    payloadType (KeyExpirationTime _) = "KeyExpirationTime"
    payloadType (PreferredSymmetricAlgorithms _) = "PreferredSymmetricAlgorithms"
    payloadType (RevocationKey {}) = "RevocationKey"
    payloadType (Issuer _) = "Issuer"
    payloadType (NotationData {}) = "NotationData"
    payloadType (PreferredHashAlgorithms _) = "PreferredHashAlgorithms"
    payloadType (PreferredCompressionAlgorithms _) = "PreferredCompressionAlgorithms"
    payloadType (KeyServerPreferences _) = "KeyServerPreferences"
    payloadType (PreferredKeyServer _) = "PreferredKeyServer"
    payloadType (PrimaryUserId _) = "PrimaryUserId"
    payloadType (PolicyURL _) = "PolicyURL"
    payloadType (KeyFlags _) = "KeyFlags"
    payloadType (SignersUserId _) = "SignersUserId"
    payloadType (ReasonForRevocation {}) = "ReasonForRevocation"
    payloadType (Features _) = "Features"
    payloadType (SignatureTarget {}) = "SignatureTarget"
    payloadType (EmbeddedSignature _) = "EmbeddedSignature"
    payloadType (IssuerFingerprint {}) = "IssuerFingerprint"
    payloadType (UserDefinedSigSub _ _) = "UserDefinedSigSub"
    payloadType (OtherSigSub _ _) = "OtherSigSub"
    payloadType (IntendedRecipient _ _) = "IntendedRecipient"
    payloadType (PreferredAEADCiphersuites _) = "PreferredAEADCiphersuites"

-- | RFC9580 §5.2.3.5 lists which subpacket types can be marked critical
canBeCritical :: SigSubPacketPayload -> Bool
canBeCritical KeyFlags {} = True
canBeCritical IssuerFingerprint {} = True
canBeCritical EmbeddedSignature {} = True
-- Future-proofing: unknown critical types are allowed to be critical
canBeCritical UserDefinedSigSub {} = True
canBeCritical OtherSigSub {} = True
canBeCritical _ = False

-- | Wrapper GADT for algorithm-specific private keys
-- Encodes the algorithm at the type level to ensure type-safe key/builder matching
data PrivateKeyFor (algo :: PubKeyAlgorithm) where
  RSAPrivateKey :: RSATypes.PrivateKey -> PrivateKeyFor 'RSA
  Ed25519PrivateKey :: Ed25519.SecretKey -> PrivateKeyFor 'Ed25519
  Ed448PrivateKey :: Ed448.SecretKey -> PrivateKeyFor 'Ed448
  DSAPrivateKey :: DSA.PrivateKey -> PrivateKeyFor 'DSA
  ECDSAPrivateKey :: ECDSA.PrivateKey -> PrivateKeyFor 'ECDSA

class KnownPubKeyAlgorithm (algo :: PubKeyAlgorithm) where
  demotePubKeyAlgorithmT :: Proxy algo -> PubKeyAlgorithm

instance KnownPubKeyAlgorithm 'RSA where
  demotePubKeyAlgorithmT _ = RSA

instance KnownPubKeyAlgorithm 'DSA where
  demotePubKeyAlgorithmT _ = DSA

instance KnownPubKeyAlgorithm 'ECDSA where
  demotePubKeyAlgorithmT _ = ECDSA

instance KnownPubKeyAlgorithm 'EdDSA where
  demotePubKeyAlgorithmT _ = EdDSA

instance KnownPubKeyAlgorithm 'Ed25519 where
  demotePubKeyAlgorithmT _ = Ed25519

instance KnownPubKeyAlgorithm 'Ed448 where
  demotePubKeyAlgorithmT _ = Ed448

-- | Legal subpackets encoded by placement and signature-version constraints.
-- This expands compile-time legality beyond bare hashed/unhashed staging.
data LegalSubpacket (h :: Type) (v :: Type) where
  -- RFC9580: signature creation time is a hashed subpacket.
  LegalSigCreationTime :: ThirtyTwoBitTimeStamp -> LegalSubpacket Hashed v
  -- v4 signatures use an issuer fingerprint subpacket with version marker 4 in hashed area.
  LegalIssuerFingerprintV4 :: Fingerprint -> LegalSubpacket Hashed V4Sig
  -- v6 signatures use an issuer fingerprint subpacket with version marker 6 in hashed area.
  LegalIssuerFingerprintV6 :: Fingerprint -> LegalSubpacket Hashed V6Sig
  -- Legacy issuer key ID subpacket is accepted only in v4 and only unhashed.
  LegalIssuerV4 :: EightOctetKeyId -> LegalSubpacket Unhashed V4Sig

legalSubpacket :: LegalSubpacket h v -> SigSubPacket
legalSubpacket (LegalSigCreationTime ts) = SigSubPacket False (SigCreationTime ts)
legalSubpacket (LegalIssuerFingerprintV4 fp) = SigSubPacket False (IssuerFingerprint IssuerFingerprintV4 fp)
legalSubpacket (LegalIssuerFingerprintV6 fp) = SigSubPacket False (IssuerFingerprint IssuerFingerprintV6 fp)
legalSubpacket (LegalIssuerV4 keyId) = SigSubPacket False (Issuer keyId)

singleLegalSub :: LegalSubpacket h v -> SubpacketList h v
singleLegalSub = SubpacketList . (: []) . legalSubpacket

consLegalSub :: LegalSubpacket h v -> SubpacketList h v -> SubpacketList h v
consLegalSub legal (SubpacketList sps) = SubpacketList (legalSubpacket legal : sps)

listToLegalSubs :: [LegalSubpacket h v] -> SubpacketList h v
listToLegalSubs = SubpacketList . fmap legalSubpacket

-- | Construct a single hashed subpacket into a hashed subpacket list
singleHashedSub :: SigSubPacket -> HashedSubpackets v
singleHashedSub sp = SubpacketList [sp]

-- | Construct a single unhashed subpacket into an unhashed subpacket list
singleUnhashedSub :: SigSubPacket -> UnhashedSubpackets v
singleUnhashedSub sp = SubpacketList [sp]

-- | Prepend a subpacket to a hashed subpacket list
consHashedSub :: SigSubPacket -> HashedSubpackets v -> HashedSubpackets v
consHashedSub sp (SubpacketList sps) = SubpacketList (sp : sps)

-- | Controls how text payload is normalized for CanonicalTextSig.
--
-- RFC 9580 §5.2.1.2 requires only CRLF normalization for type 0x01
-- signatures.  Trailing-whitespace stripping is only mandated by
-- §7 (Cleartext Signature Framework).  Use 'CleartextCompat' for
-- GnuPG-compatible behaviour when producing cleartext-armored
-- signatures; use 'RFC9580Strict' for inline text signatures.
data TextNormalizationMode
  = RFC9580Strict
    -- ^ CRLF normalization only; no trailing-whitespace stripping.
    --   Correct for inline type 0x01 document signatures.
  | CleartextCompat
    -- ^ CRLF normalization *plus* per-line trailing-whitespace stripping.
    --   Required by the Cleartext Signature Framework (RFC 9580 §7).
  deriving (Eq, Show)

-- | Staged signature builder that enforces completion order
-- 
-- Usage pattern:
--   builder <- sigBuilderInit SigTypeBinary RSA SHA256
--   builder' <- addHashedSubs hashedList builder
--   finalPayload <- buildSigV4 sigMPIs (addUnhashedSubs unhashedList builder')
data SigBuilder (hashedness :: Type) (v :: Type) (algo :: PubKeyAlgorithm) = SigBuilder
  { sbSigType :: SigType
  , sbHashAlgo :: HashAlgorithm
  , sbHashedSubs :: [SigSubPacket]
  , sbUnhashedSubs :: [SigSubPacket]
  , sbSalt :: BuilderSalt v
  , sbTextNormMode :: TextNormalizationMode
    -- ^ Normalization mode for CanonicalTextSig payloads.
    --   Defaults to 'CleartextCompat' for backward compatibility.
  }

type family BuilderSalt (v :: Type) where
  BuilderSalt V4Sig = ()
  BuilderSalt V6Sig = SignatureSalt

sbPubKeyAlgo :: forall hashedness v algo. KnownPubKeyAlgorithm algo => SigBuilder hashedness v algo -> PubKeyAlgorithm
sbPubKeyAlgo _ = demotePubKeyAlgorithmT (Proxy @algo)

-- | Initialize a builder for a v4 signature
-- Starts with no subpackets and must call addHashedSubs and addUnhashedSubs
sigBuilderInit ::
     forall algo
   . KnownPubKeyAlgorithm algo
  => SigType
  -> HashAlgorithm
  -> SigBuilder Hashed V4Sig algo
sigBuilderInit st ha = SigBuilder
  { sbSigType = st
  , sbHashAlgo = ha
  , sbHashedSubs = []
  , sbUnhashedSubs = []
  , sbSalt = ()
  , sbTextNormMode = CleartextCompat
  }

-- | Initialize a builder for a v4 signature with compile-time RFC/hash policy enforcement.
sigBuilderInitTyped ::
     forall algo rfc h
   . (KnownPubKeyAlgorithm algo, HashAlgoAllowedFor rfc h)
  => OpenPGPRFCW rfc
  -> SigType
  -> HashAlgorithmW h
  -> SigBuilder Hashed V4Sig algo
sigBuilderInitTyped _ st hashW =
  sigBuilderInit @algo st (demoteHashAlgorithmW hashW)

-- | Initialize a v4 builder from runtime RFC/hash inputs with policy validation
-- and witness promotion at the API boundary.
sigBuilderInitRuntime ::
     forall algo
   . KnownPubKeyAlgorithm algo
  => OpenPGPRFC
  -> SigType
  -> HashAlgorithm
  -> Either String (SigBuilder Hashed V4Sig algo)
sigBuilderInitRuntime rfc st ha =
  withGenerationHashWitness
    rfc
    ha
    (\rfcW hashW -> sigBuilderInitTyped @algo rfcW st hashW)

-- | Initialize a builder for a v6 signature
sigBuilderInitV6 ::
     forall algo
   . KnownPubKeyAlgorithm algo
  => SigType
  -> HashAlgorithm
  -> SignatureSalt
  -> SigBuilder Hashed V6Sig algo
sigBuilderInitV6 st ha salt = SigBuilder
  { sbSigType = st
  , sbHashAlgo = ha
  , sbHashedSubs = []
  , sbUnhashedSubs = []
  , sbSalt = salt
  , sbTextNormMode = CleartextCompat
  }

-- | Initialize a builder for a v6 signature with compile-time RFC/hash policy enforcement.
sigBuilderInitV6Typed ::
     forall algo rfc h
   . (KnownPubKeyAlgorithm algo, HashAlgoAllowedFor rfc h)
  => OpenPGPRFCW rfc
  -> SigType
  -> HashAlgorithmW h
  -> SignatureSalt
  -> SigBuilder Hashed V6Sig algo
sigBuilderInitV6Typed _ st hashW salt =
  sigBuilderInitV6 @algo st (demoteHashAlgorithmW hashW) salt

-- | Initialize a v6 builder from runtime RFC/hash inputs with policy validation
-- and witness promotion at the API boundary.
sigBuilderInitV6Runtime ::
     forall algo
   . KnownPubKeyAlgorithm algo
  => OpenPGPRFC
  -> SigType
  -> HashAlgorithm
  -> SignatureSalt
  -> Either String (SigBuilder Hashed V6Sig algo)
sigBuilderInitV6Runtime rfc st ha salt =
  withGenerationHashWitness
    rfc
    ha
    (\rfcW hashW -> sigBuilderInitV6Typed @algo rfcW st hashW salt)

withGenerationHashWitness ::
     OpenPGPRFC
  -> HashAlgorithm
  -> (forall rfc h. HashAlgoAllowedFor rfc h => OpenPGPRFCW rfc -> HashAlgorithmW h -> a)
  -> Either String a
withGenerationHashWitness rfc ha mk
  | ha `elem` deprecatedHashAlgorithms (policyGenerationDeprecations (policyForRFC rfc)) =
      Left (hashPolicyDisallowedMessage rfc ha)
  | otherwise =
      case promoteHashAlgorithm ha of
        Nothing ->
          Left (hashNotTypedBuilderMessage ha)
        Just (SomeHashAlgorithmW hashW) ->
          case promoteOpenPGPRFC rfc of
            SomeOpenPGPRFCW RFC2440W ->
              Right $
                case hashW of
                  DeprecatedMD5W -> mk RFC2440W DeprecatedMD5W
                  SHA1W -> mk RFC2440W SHA1W
                  RIPEMD160W -> mk RFC2440W RIPEMD160W
                  SHA256W -> mk RFC2440W SHA256W
                  SHA384W -> mk RFC2440W SHA384W
                  SHA512W -> mk RFC2440W SHA512W
                  SHA224W -> mk RFC2440W SHA224W
                  SHA3_256W -> mk RFC2440W SHA3_256W
                  SHA3_512W -> mk RFC2440W SHA3_512W
            SomeOpenPGPRFCW RFC4880W ->
              case hashW of
                SHA1W -> Right (mk RFC4880W SHA1W)
                RIPEMD160W -> Right (mk RFC4880W RIPEMD160W)
                SHA256W -> Right (mk RFC4880W SHA256W)
                SHA384W -> Right (mk RFC4880W SHA384W)
                SHA512W -> Right (mk RFC4880W SHA512W)
                SHA224W -> Right (mk RFC4880W SHA224W)
                SHA3_256W -> Right (mk RFC4880W SHA3_256W)
                SHA3_512W -> Right (mk RFC4880W SHA3_512W)
                DeprecatedMD5W ->
                  Left (hashPolicyDisallowedMessage rfc ha)
            SomeOpenPGPRFCW RFC9580W ->
              case hashW of
                SHA256W -> Right (mk RFC9580W SHA256W)
                SHA384W -> Right (mk RFC9580W SHA384W)
                SHA512W -> Right (mk RFC9580W SHA512W)
                SHA224W -> Right (mk RFC9580W SHA224W)
                SHA3_256W -> Right (mk RFC9580W SHA3_256W)
                SHA3_512W -> Right (mk RFC9580W SHA3_512W)
                DeprecatedMD5W ->
                  Left (hashPolicyDisallowedMessage rfc ha)
                SHA1W ->
                  Left (hashPolicyDisallowedMessage rfc ha)
                RIPEMD160W ->
                  Left (hashPolicyDisallowedMessage rfc ha)

hashPolicyDisallowedMessage :: OpenPGPRFC -> HashAlgorithm -> String
hashPolicyDisallowedMessage rfc ha =
  "signature hash algorithm disallowed by RFC policy (" ++ show rfc ++ "): " ++ show ha

hashNotTypedBuilderMessage :: HashAlgorithm -> String
hashNotTypedBuilderMessage ha =
  "signature hash algorithm is not supported by typed builder API: " ++ show ha

-- | Add hashed subpackets to a builder
-- Must be called before addUnhashedSubs
addHashedSubs :: HashedSubpackets v -> SigBuilder Hashed v algo -> SigBuilder Unhashed v algo
addHashedSubs (SubpacketList sps) builder = builder { sbHashedSubs = sps }

-- | Add unhashed subpackets to a builder
-- Must be called after addHashedSubs
addUnhashedSubs :: UnhashedSubpackets v -> SigBuilder Unhashed v algo -> SigBuilder Unhashed v algo
addUnhashedSubs (SubpacketList sps) builder = builder { sbUnhashedSubs = sps }

-- | Build a v4 signature from a completed builder and MPI values
buildSigV4 ::
     KnownPubKeyAlgorithm algo
  => SigBuilder Unhashed V4Sig algo
  -> Word16
  -> NonEmpty MPI
  -> SignaturePayload
buildSigV4 builder hashLeft mpis = SigV4
  (sbSigType builder)
  (sbPubKeyAlgo builder)
  (sbHashAlgo builder)
  (sbHashedSubs builder)
  (sbUnhashedSubs builder)
  hashLeft
  mpis

-- | Build a v6 signature from a completed builder and MPI values
buildSigV6 ::
     KnownPubKeyAlgorithm algo
  => SigBuilder Unhashed V6Sig algo
  -> Word16
  -> NonEmpty MPI
  -> SignaturePayload
buildSigV6 builder hashLeft mpis =
  SigV6
    (sbSigType builder)
    (sbPubKeyAlgo builder)
    (sbHashAlgo builder)
    (sbSalt builder)
    (sbHashedSubs builder)
    (sbUnhashedSubs builder)
    hashLeft
    mpis

-- | Convert a list to a hashed subpacket list
-- Used to bridge between list-based and phantom-typed APIs
listToHashedSubs :: [SigSubPacket] -> HashedSubpackets v
listToHashedSubs sps = SubpacketList sps

-- | Convert a list to an unhashed subpacket list
-- Used to bridge between list-based and phantom-typed APIs
listToUnhashedSubs :: [SigSubPacket] -> UnhashedSubpackets v
listToUnhashedSubs sps = SubpacketList sps
