PKITypes.hs

-- PKITypes.hs: OpenPGP (RFC9580) data types for public/secret keys
-- Copyright © 2012-2026  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ViewPatterns #-}

module Codec.Encryption.OpenPGP.Types.Internal.PKITypes where

import qualified Data.Aeson as A
import qualified Data.ByteString as B
import Data.ByteString.Lazy (ByteString)
import qualified Data.ByteString.Lazy as BL
import Data.Data (Data (..))
import qualified Data.Data as DData
import Data.Hashable (Hashable (..))
import Data.Ord (comparing)
import Data.Typeable (Typeable)
import Data.Word (Word16)
import GHC.Generics (Generic)
import Prettyprinter (Pretty (..), (<+>))

import Codec.Encryption.OpenPGP.Types.Internal.Base
import Codec.Encryption.OpenPGP.Types.Internal.CryptonNewtypes

data EdSigningCurve
    = EdSigningCurve25519
    | EdSigningCurve448
    deriving (Data, Eq, Generic, Ord, Show, Typeable)

instance Hashable EdSigningCurve

instance Pretty EdSigningCurve where
    pretty EdSigningCurve25519 = pretty "Ed25519"
    pretty EdSigningCurve448 = pretty "Ed448"

instance A.FromJSON EdSigningCurve

instance A.ToJSON EdSigningCurve

newtype EPoint
    = EPoint
    { unEPoint :: Integer
    }
    deriving (Data, Eq, Generic, Ord, Pretty, Show, Typeable)

instance Hashable EPoint

instance A.FromJSON EPoint

instance A.ToJSON EPoint

data EdPoint
    = PrefixedNativeEPoint EPoint
    | NativeEPoint EPoint
    deriving (Data, Eq, Generic, Ord, Show, Typeable)

instance Hashable EdPoint

instance Pretty EdPoint where
    pretty (PrefixedNativeEPoint ep) = pretty "prefixed-native" <+> pretty ep
    pretty (NativeEPoint ep) = pretty "native" <+> pretty ep

instance A.FromJSON EdPoint

instance A.ToJSON EdPoint

data PKey
    = RSAPubKey RSA_PublicKey
    | DSAPubKey DSA_PublicKey
    | ElGamalPubKey Integer Integer Integer
    | ECDHPubKey PKey HashAlgorithm SymmetricAlgorithm
    | ECDSAPubKey ECDSA_PublicKey
    | EdDSAPubKey EdSigningCurve EdPoint
    | MLKEMPubKey B.ByteString
    | MLDSAPubKey B.ByteString
    | SLHDSAPubKey B.ByteString
    | UnknownPKey ByteString
    deriving (Data, Eq, Generic, Ord, Show, Typeable)

instance Hashable PKey

instance Pretty PKey where
    pretty (RSAPubKey p) = pretty "RSA" <+> pretty p
    pretty (DSAPubKey p) = pretty "DSA" <+> pretty p
    pretty (ElGamalPubKey p g y) =
        pretty "Elgamal" <+> pretty p <+> pretty g <+> pretty y
    pretty (ECDHPubKey p ha sa) =
        pretty "ECDH" <+> pretty p <+> pretty ha <+> pretty sa
    pretty (ECDSAPubKey p) = pretty "ECDSA" <+> pretty p
    pretty (EdDSAPubKey c ep) = pretty c <+> pretty ep
    pretty (MLKEMPubKey bs) = pretty "ML-KEM" <+> pretty (bsToHexUpper (BL.fromStrict bs))
    pretty (MLDSAPubKey bs) = pretty "ML-DSA" <+> pretty (bsToHexUpper (BL.fromStrict bs))
    pretty (SLHDSAPubKey bs) = pretty "SLH-DSA" <+> pretty (bsToHexUpper (BL.fromStrict bs))
    pretty (UnknownPKey bs) = pretty "<unknown>" <+> pretty (bsToHexUpper bs)

instance A.ToJSON PKey where
    toJSON (RSAPubKey p) = A.toJSON p
    toJSON (DSAPubKey p) = A.toJSON p
    toJSON (ElGamalPubKey p g y) = A.toJSON (p, g, y)
    toJSON (ECDHPubKey p ha sa) = A.toJSON (p, ha, sa)
    toJSON (ECDSAPubKey p) = A.toJSON p
    toJSON (EdDSAPubKey c ep) = A.toJSON (c, ep)
    toJSON (MLKEMPubKey bs) = A.toJSON (B.unpack bs)
    toJSON (MLDSAPubKey bs) = A.toJSON (B.unpack bs)
    toJSON (SLHDSAPubKey bs) = A.toJSON (B.unpack bs)
    toJSON (UnknownPKey bs) = A.toJSON (BL.unpack bs)

data SKey
    = RSAPrivateKey RSA_PrivateKey
    | DSAPrivateKey DSA_PrivateKey
    | ElGamalPrivateKey Integer
    | ECDHPrivateKey ECDSA_PrivateKey
    | ECDSAPrivateKey ECDSA_PrivateKey
    | EdDSAPrivateKey EdSigningCurve B.ByteString
    | X25519PrivateKey B.ByteString
    | X448PrivateKey B.ByteString
    | MLKEMPrivateKey B.ByteString
    | MLDSAPrivateKey B.ByteString
    | SLHDSAPrivateKey B.ByteString
    | UnknownSKey ByteString
    deriving (Data, Eq, Generic, Ord, Show, Typeable)

instance Hashable SKey

instance Pretty SKey where
    pretty (RSAPrivateKey p) = pretty "RSA" <+> pretty p
    pretty (DSAPrivateKey p) = pretty "DSA" <+> pretty p
    pretty (ElGamalPrivateKey p) = pretty "Elgamal" <+> pretty p
    pretty (ECDHPrivateKey p) = pretty "ECDH" <+> pretty p
    pretty (ECDSAPrivateKey p) = pretty "ECDSA" <+> pretty p
    pretty (EdDSAPrivateKey c bs) =
        pretty c <+> pretty (bsToHexUpper (BL.fromStrict bs))
    pretty (X25519PrivateKey bs) = pretty "X25519" <+> pretty (bsToHexUpper (BL.fromStrict bs))
    pretty (X448PrivateKey bs) = pretty "X448" <+> pretty (bsToHexUpper (BL.fromStrict bs))
    pretty (MLKEMPrivateKey bs) =
        pretty "ML-KEM-priv" <+> pretty (bsToHexUpper (BL.fromStrict bs))
    pretty (MLDSAPrivateKey bs) =
        pretty "ML-DSA-priv" <+> pretty (bsToHexUpper (BL.fromStrict bs))
    pretty (SLHDSAPrivateKey bs) =
        pretty "SLH-DSA-priv"
            <+> pretty (bsToHexUpper (BL.fromStrict bs))
    pretty (UnknownSKey bs) = pretty "<unknown>" <+> pretty (bsToHexUpper bs)

instance A.ToJSON SKey where
    toJSON (RSAPrivateKey k) = A.toJSON k
    toJSON (DSAPrivateKey k) = A.toJSON k
    toJSON (ElGamalPrivateKey k) = A.toJSON k
    toJSON (ECDHPrivateKey k) = A.toJSON k
    toJSON (ECDSAPrivateKey k) = A.toJSON k
    toJSON (EdDSAPrivateKey c bs) = A.toJSON (c, B.unpack bs)
    toJSON (X25519PrivateKey bs) = A.toJSON (B.unpack bs)
    toJSON (X448PrivateKey bs) = A.toJSON (B.unpack bs)
    toJSON (MLKEMPrivateKey bs) = A.toJSON (B.unpack bs)
    toJSON (MLDSAPrivateKey bs) = A.toJSON (B.unpack bs)
    toJSON (SLHDSAPrivateKey bs) = A.toJSON (B.unpack bs)
    toJSON (UnknownSKey bs) = A.toJSON (BL.unpack bs)

data PKPayload (v :: KeyVersion) where
    PKPayloadV3
        :: ThirtyTwoBitTimeStamp
        -> V3Expiration
        -> PubKeyAlgorithm
        -> PKey
        -> PKPayload 'DeprecatedV3
    PKPayloadV4
        :: ThirtyTwoBitTimeStamp
        -> PubKeyAlgorithm
        -> PKey
        -> PKPayload 'V4
    PKPayloadV6
        :: ThirtyTwoBitTimeStamp
        -> PubKeyAlgorithm
        -> PKey
        -> PKPayload 'V6

deriving instance Eq (PKPayload v)
deriving instance Ord (PKPayload v)
deriving instance Show (PKPayload v)

instance Hashable (PKPayload v) where
    hashWithSalt s = hashWithSalt s . pkPayloadFields

instance Pretty (PKPayload v) where
    pretty pkp =
        let (kv, ts, v3e, pka, p) = pkPayloadFields pkp
         in pretty kv
                <+> pretty ts
                <+> pretty v3e
                <+> pretty pka
                <+> pretty p

instance A.ToJSON (PKPayload v) where
    toJSON = A.toJSON . pkPayloadFields

data SomePKPayload where
    SomePKPayload :: PKPayload v -> SomePKPayload

deriving instance Show SomePKPayload
deriving instance Typeable SomePKPayload

pkPayloadDataType :: DData.DataType
pkPayloadDataType =
    DData.mkDataType
        "Codec.Encryption.OpenPGP.Types.Internal.PKITypes.SomePKPayload"
        [pkPayloadConstr]

pkPayloadConstr :: DData.Constr
pkPayloadConstr = DData.mkConstr pkPayloadDataType "PKPayload" [] DData.Prefix

instance Data SomePKPayload where
    gfoldl f z (PKPayload kv ts v3e pka p) =
        z PKPayload `f` kv `f` ts `f` v3e `f` pka `f` p
    gunfold k z c
        | c == pkPayloadConstr = k (k (k (k (k (z PKPayload)))))
        | otherwise =
            error "gunfold: invalid constructor for SomePKPayload"
    toConstr _ = pkPayloadConstr
    dataTypeOf _ = pkPayloadDataType

instance Eq SomePKPayload where
    a == b = somePKPayloadFields a == somePKPayloadFields b

instance Ord SomePKPayload where
    compare = comparing somePKPayloadFields

instance Hashable SomePKPayload where
    hashWithSalt s = hashWithSalt s . somePKPayloadFields

instance Pretty SomePKPayload where
    pretty (PKPayload kv ts v3e pka p) =
        pretty kv
            <+> pretty ts
            <+> pretty v3e
            <+> pretty pka
            <+> pretty p

instance A.ToJSON SomePKPayload where
    toJSON = A.toJSON . somePKPayloadFields

pkPayloadFields
    :: PKPayload v
    -> ( KeyVersion
       , ThirtyTwoBitTimeStamp
       , V3Expiration
       , PubKeyAlgorithm
       , PKey
       )
pkPayloadFields (PKPayloadV3 ts v3e pka p) = (DeprecatedV3, ts, v3e, pka, p)
pkPayloadFields (PKPayloadV4 ts pka p) = (V4, ts, 0, pka, p)
pkPayloadFields (PKPayloadV6 ts pka p) = (V6, ts, 0, pka, p)

somePKPayloadFields
    :: SomePKPayload
    -> ( KeyVersion
       , ThirtyTwoBitTimeStamp
       , V3Expiration
       , PubKeyAlgorithm
       , PKey
       )
somePKPayloadFields (SomePKPayload pkp) = pkPayloadFields pkp

pattern PKPayload
    :: KeyVersion
    -> ThirtyTwoBitTimeStamp
    -> V3Expiration
    -> PubKeyAlgorithm
    -> PKey
    -> SomePKPayload
pattern PKPayload kv ts v3e pka p <-
    (somePKPayloadFields -> (kv, ts, v3e, pka, p))
    where
        PKPayload DeprecatedV3 ts v3e pka p = SomePKPayload (PKPayloadV3 ts v3e pka p)
        PKPayload V4 ts _ pka p = SomePKPayload (PKPayloadV4 ts pka p)
        PKPayload V6 ts _ pka p = SomePKPayload (PKPayloadV6 ts pka p)

{-# COMPLETE PKPayload #-}

_keyVersion :: SomePKPayload -> KeyVersion
_keyVersion (PKPayload kv _ _ _ _) = kv

_timestamp :: SomePKPayload -> ThirtyTwoBitTimeStamp
_timestamp (PKPayload _ ts _ _ _) = ts

_v3exp :: SomePKPayload -> V3Expiration
_v3exp (PKPayload _ _ v3e _ _) = v3e

_pkalgo :: SomePKPayload -> PubKeyAlgorithm
_pkalgo (PKPayload _ _ _ pka _) = pka

_pubkey :: SomePKPayload -> PKey
_pubkey (PKPayload _ _ _ _ p) = p

data SKAddendum
    = SUS16bit SymmetricAlgorithm S2K IV ByteString
    | SUSSHA1 SymmetricAlgorithm S2K IV ByteString
    | SUSAEAD SymmetricAlgorithm AEADAlgorithm S2K IV ByteString
    | SUSym SymmetricAlgorithm IV ByteString
    | SUUnencrypted SKey Word16
    deriving (Data, Eq, Generic, Show, Typeable)

instance Ord SKAddendum where
    compare (SUS16bit sa1 s2k1 iv1 bs1) (SUS16bit sa2 s2k2 iv2 bs2) =
        compare sa1 sa2
            <> compare s2k1 s2k2
            <> compare iv1 iv2
            <> compare bs1 bs2
    compare (SUSSHA1 sa1 s2k1 iv1 bs1) (SUSSHA1 sa2 s2k2 iv2 bs2) =
        compare sa1 sa2
            <> compare s2k1 s2k2
            <> compare iv1 iv2
            <> compare bs1 bs2
    compare (SUSAEAD sa1 aa1 s2k1 iv1 bs1) (SUSAEAD sa2 aa2 s2k2 iv2 bs2) =
        compare sa1 sa2
            <> compare aa1 aa2
            <> compare s2k1 s2k2
            <> compare iv1 iv2
            <> compare bs1 bs2
    compare (SUSym sa1 iv1 bs1) (SUSym sa2 iv2 bs2) =
        compare sa1 sa2 <> compare iv1 iv2 <> compare bs1 bs2
    compare (SUUnencrypted sk1 ck1) (SUUnencrypted sk2 ck2) =
        compare sk1 sk2 <> compare ck1 ck2
    compare SUS16bit {} SUSSHA1 {} = LT
    compare SUS16bit {} SUSAEAD {} = LT
    compare SUS16bit {} SUSym {} = LT
    compare SUS16bit {} SUUnencrypted {} = LT
    compare SUSSHA1 {} SUS16bit {} = GT
    compare SUSSHA1 {} SUSAEAD {} = LT
    compare SUSSHA1 {} SUSym {} = LT
    compare SUSSHA1 {} SUUnencrypted {} = LT
    compare SUSAEAD {} SUS16bit {} = GT
    compare SUSAEAD {} SUSSHA1 {} = GT
    compare SUSAEAD {} SUSym {} = LT
    compare SUSAEAD {} SUUnencrypted {} = LT
    compare SUSym {} SUS16bit {} = GT
    compare SUSym {} SUSSHA1 {} = GT
    compare SUSym {} SUSAEAD {} = GT
    compare SUSym {} SUUnencrypted {} = LT
    compare SUUnencrypted {} _ = GT

instance Hashable SKAddendum

instance Pretty SKAddendum where
    pretty (SUS16bit sa s2k iv bs) =
        pretty "SUS16bit"
            <+> pretty sa
            <+> pretty s2k
            <+> pretty iv
            <+> pretty (bsToHexUpper bs)
    pretty (SUSSHA1 sa s2k iv bs) =
        pretty "SUSSHA1"
            <+> pretty sa
            <+> pretty s2k
            <+> pretty iv
            <+> pretty (bsToHexUpper bs)
    pretty (SUSAEAD sa aa s2k iv bs) =
        pretty "SUSAEAD"
            <+> pretty sa
            <+> pretty aa
            <+> pretty s2k
            <+> pretty iv
            <+> pretty (bsToHexUpper bs)
    pretty (SUSym sa iv bs) =
        pretty "SUSym"
            <+> pretty sa
            <+> pretty iv
            <+> pretty (bsToHexUpper bs)
    pretty (SUUnencrypted s ck) =
        pretty "SUUnencrypted" <+> pretty s <+> pretty ck

instance A.ToJSON SKAddendum where
    toJSON (SUS16bit sa s2k iv bs) = A.toJSON (sa, s2k, iv, BL.unpack bs)
    toJSON (SUSSHA1 sa s2k iv bs) = A.toJSON (sa, s2k, iv, BL.unpack bs)
    toJSON (SUSAEAD sa aa s2k iv bs) = A.toJSON (sa, aa, s2k, iv, BL.unpack bs)
    toJSON (SUSym sa iv bs) = A.toJSON (sa, iv, BL.unpack bs)
    toJSON (SUUnencrypted s ck) = A.toJSON (s, ck)

class LegacyKeyVersion (v :: KeyVersion)

instance LegacyKeyVersion 'DeprecatedV3

instance LegacyKeyVersion 'V4

data SKAddendumV (v :: KeyVersion) where
    SKA16bit
        :: (LegacyKeyVersion v)
        => SymmetricAlgorithm
        -> S2K
        -> IV
        -> ByteString
        -> SKAddendumV v
    SKASHA1Legacy
        :: (LegacyKeyVersion v)
        => SymmetricAlgorithm
        -> S2K
        -> IV
        -> ByteString
        -> SKAddendumV v
    SKASHA1V6
        :: SymmetricAlgorithm
        -> S2K
        -> IV
        -> ByteString
        -> SKAddendumV 'V6
    SKAAEADV6
        :: SymmetricAlgorithm
        -> AEADAlgorithm
        -> S2K
        -> IV
        -> ByteString
        -> SKAddendumV 'V6
    SKAAEADLegacy
        :: (LegacyKeyVersion v)
        => SymmetricAlgorithm
        -> AEADAlgorithm
        -> S2K
        -> IV
        -> ByteString
        -> SKAddendumV v
    SKASymLegacy
        :: (LegacyKeyVersion v)
        => SymmetricAlgorithm
        -> IV
        -> ByteString
        -> SKAddendumV v
    SKASymV6
        :: SymmetricAlgorithm
        -> IV
        -> ByteString
        -> SKAddendumV 'V6
    SKAUnencryptedLegacy
        :: (LegacyKeyVersion v)
        => SKey
        -> Word16
        -> SKAddendumV v
    SKAUnencryptedV6
        :: SKey
        -> SKAddendumV 'V6

deriving instance Show (SKAddendumV v)

data SomeSKAddendumV where
    SomeSKAddendumV :: SKAddendumV v -> SomeSKAddendumV

deriving instance Show SomeSKAddendumV

toSKAddendum :: SKAddendumV v -> SKAddendum
toSKAddendum (SKA16bit sa s2k iv bs) = SUS16bit sa s2k iv bs
toSKAddendum (SKASHA1Legacy sa s2k iv bs) = SUSSHA1 sa s2k iv bs
toSKAddendum (SKASHA1V6 sa s2k iv bs) = SUSSHA1 sa s2k iv bs
toSKAddendum (SKAAEADV6 sa aa s2k iv bs) = SUSAEAD sa aa s2k iv bs
toSKAddendum (SKAAEADLegacy sa aa s2k iv bs) = SUSAEAD sa aa s2k iv bs
toSKAddendum (SKASymLegacy sa iv bs) = SUSym sa iv bs
toSKAddendum (SKASymV6 sa iv bs) = SUSym sa iv bs
toSKAddendum (SKAUnencryptedLegacy sk checksum) = SUUnencrypted sk checksum
toSKAddendum (SKAUnencryptedV6 sk) = SUUnencrypted sk 0

fromSKAddendumForKeyVersion
    :: KeyVersion
    -> SKAddendum
    -> Either String SomeSKAddendumV
fromSKAddendumForKeyVersion DeprecatedV3 (SUS16bit sa s2k iv bs) =
    Right
        ( SomeSKAddendumV
            (SKA16bit sa s2k iv bs :: SKAddendumV 'DeprecatedV3)
        )
fromSKAddendumForKeyVersion DeprecatedV3 (SUSSHA1 sa s2k iv bs) =
    Right
        ( SomeSKAddendumV
            (SKASHA1Legacy sa s2k iv bs :: SKAddendumV 'DeprecatedV3)
        )
fromSKAddendumForKeyVersion DeprecatedV3 (SUSym sa iv bs) =
    Right
        ( SomeSKAddendumV
            (SKASymLegacy sa iv bs :: SKAddendumV 'DeprecatedV3)
        )
fromSKAddendumForKeyVersion DeprecatedV3 (SUUnencrypted sk checksum) =
    Right
        ( SomeSKAddendumV
            (SKAUnencryptedLegacy sk checksum :: SKAddendumV 'DeprecatedV3)
        )
fromSKAddendumForKeyVersion DeprecatedV3 (SUSAEAD sa aa s2k iv bs) =
    Right
        ( SomeSKAddendumV
            (SKAAEADLegacy sa aa s2k iv bs :: SKAddendumV 'DeprecatedV3)
        )
fromSKAddendumForKeyVersion V4 (SUS16bit sa s2k iv bs) =
    Right
        (SomeSKAddendumV (SKA16bit sa s2k iv bs :: SKAddendumV 'V4))
fromSKAddendumForKeyVersion V4 (SUSSHA1 sa s2k iv bs) =
    Right
        (SomeSKAddendumV (SKASHA1Legacy sa s2k iv bs :: SKAddendumV 'V4))
fromSKAddendumForKeyVersion V4 (SUSym sa iv bs) =
    Right
        (SomeSKAddendumV (SKASymLegacy sa iv bs :: SKAddendumV 'V4))
fromSKAddendumForKeyVersion V4 (SUUnencrypted sk checksum) =
    Right
        ( SomeSKAddendumV
            (SKAUnencryptedLegacy sk checksum :: SKAddendumV 'V4)
        )
fromSKAddendumForKeyVersion V4 (SUSAEAD sa aa s2k iv bs) =
    Right
        ( SomeSKAddendumV
            (SKAAEADLegacy sa aa s2k iv bs :: SKAddendumV 'V4)
        )
fromSKAddendumForKeyVersion V6 (SUS16bit _ _ _ _) =
    Left
        "v6 secret keys must not use 16-bit checksum protected secret key addendums"
fromSKAddendumForKeyVersion V6 (SUSSHA1 sa s2k iv bs) =
    Right (SomeSKAddendumV (SKASHA1V6 sa s2k iv bs))
fromSKAddendumForKeyVersion V6 (SUSAEAD sa aa s2k iv bs) =
    Right (SomeSKAddendumV (SKAAEADV6 sa aa s2k iv bs))
fromSKAddendumForKeyVersion V6 (SUSym sa iv bs) =
    Right (SomeSKAddendumV (SKASymV6 sa iv bs))
fromSKAddendumForKeyVersion V6 (SUUnencrypted sk _) =
    Right (SomeSKAddendumV (SKAUnencryptedV6 sk))

fromSKAddendumForPKPayload
    :: SomePKPayload
    -> SKAddendum
    -> Either String SomeSKAddendumV
fromSKAddendumForPKPayload pkp =
    fromSKAddendumForKeyVersion (_keyVersion pkp)