-- Keyring.hs: OpenPGP (RFC9580) transferable keys parsing
-- Copyright © 2012-2026  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}

module Data.Conduit.OpenPGP.Keyring
    ( TypedTKConduitError (..)
    , conduitToSomeTKsEither
    , conduitToSomeTKsDroppingEither
    , AuthSecretSubkeyUID (..)
    , AuthSecretSubkeyAtTime (..)
    , AuthSecretSubkeyRejectionReason (..)
    , AuthSecretSubkeyRejectedAtTime (..)
    , AuthSecretSubkeysAtReport (..)
    , authSecretSubkeysAt
    , authSecretSubkeysAtReport
    , conduitToAuthSecretSubkeysAt
    , conduitToAuthSecretSubkeysAtReport
    , conduitToTKsEither
    , conduitToTKsDroppingEither
    , conduitToTKsWithWireRepEither
    , conduitToTKsDroppingWithWireRepEither
    , conduitDropErrorsAndNothings
    , KeyringChunkParseError (..)
    , sinkPublicKeyringMap
    , sinkSecretKeyringMap
    , publicTKToKeyring
    , secretTKToKeyring
    , partitionSomeTKs
    ) where

import Control.Error.Util (hush)
import Control.Lens ((^.))
import Control.Monad (join)
import Data.Bifunctor (first)
import Data.Conduit
import qualified Data.Conduit.List as CL
import Data.IxSet.Typed (empty, insert)
import Data.List (find)
import Data.Maybe (mapMaybe, maybeToList)
import qualified Data.Set as Set
import Data.Text (Text)
import Data.Time.Clock (UTCTime)

import Codec.Encryption.OpenPGP.Expirations
    ( isCertificationSig
    , isPKTimeValidWithSelfSignatures
    , isTKTimeValid
    , newestByCreationTime
    , signatureCreationTime
    , signatureEffectiveAt
    )
import Codec.Encryption.OpenPGP.KeyringParser
    ( KeyringChunkParseError (..)
    , anyTK
    , anyTKWithWireRep
    , finalizeParsingEither
    , parseAChunkEither
    )
import Codec.Encryption.OpenPGP.Ontology
    ( isSubkeyBindingSig
    , isTrustPkt
    )
import Codec.Encryption.OpenPGP.Policy
    ( defaultVerificationPolicy
    )
import Codec.Encryption.OpenPGP.SignatureQualities
    ( signatureHashedSubpacketsKnown
    )
import Codec.Encryption.OpenPGP.Signatures
    ( verifyAgainstKeys
    , verifySigWith
    , verifyTKWith
    )
import Codec.Encryption.OpenPGP.Types
import Data.Conduit.OpenPGP.Keyring.Instances ()

data Phase
    = MainKey
    | Revs
    | Uids
    | UAts
    | Subs
    | SkippingBroken
    deriving (Eq, Ord, Show)

data TypedTKConduitError
    = TypedTKParseError KeyringChunkParseError
    | TypedTKConversionError TKConversionError
    deriving (Eq, Show)

-- | Canonical strict typed conduit with explicit parse+conversion error channel.
conduitToSomeTKsEither
    :: (Monad m)
    => ConduitT Pkt (Either TypedTKConduitError (Maybe SomeTK)) m ()
conduitToSomeTKsEither =
    conduitToTKsEither
        .| CL.map toTypedSomeTKEither

{- | Tolerant typed conduit (broken transferable-key chunks may be omitted),
while still surfacing parse+conversion failures.
-}
conduitToSomeTKsDroppingEither
    :: (Monad m)
    => ConduitT Pkt (Either TypedTKConduitError (Maybe SomeTK)) m ()
conduitToSomeTKsDroppingEither =
    conduitToTKsDroppingEither
        .| CL.map toTypedSomeTKEither

toTypedSomeTKEither
    :: Either KeyringChunkParseError (Maybe TKUnknown)
    -> Either TypedTKConduitError (Maybe SomeTK)
toTypedSomeTKEither =
    either
        (Left . TypedTKParseError)
        ( \maybeUnknown ->
            case maybeUnknown of
                Nothing -> Right Nothing
                Just unknown ->
                    first
                        TypedTKConversionError
                        (Just <$> fromUnknownToTKEither unknown)
        )

data AuthSecretSubkeyUID
    = AuthSecretSubkeyUID
    { authSecretSubkeyUIDValue :: Text
    , authSecretSubkeyUIDIsPrimary :: Bool
    }
    deriving (Eq, Show)

data AuthSecretSubkeyAtTime
    = AuthSecretSubkeyAtTime
    { authSecretSubkeyPrimaryKey :: KeyPkt 'SecretPkt
    , authSecretSubkeyValue :: KeyPkt 'SecretPkt
    , authSecretSubkeyUIDs :: [AuthSecretSubkeyUID]
    , authSecretSubkeyPrimaryUID :: Maybe Text
    }
    deriving (Eq, Show)

data AuthSecretSubkeyRejectionReason
    = AuthSecretSubkeyTKVerificationFailed
    | AuthSecretSubkeyPrimaryKeyInvalidAtTime
    | AuthSecretSubkeyNotSecretSubkeyPacket
    | AuthSecretSubkeyNotSubkeyPacket
    | AuthSecretSubkeySubkeyInvalidAtTime
    | AuthSecretSubkeyMissingAuthCapability
    deriving (Eq, Show)

data AuthSecretSubkeyRejectedAtTime
    = AuthSecretSubkeyRejectedAtTime
    { authSecretSubkeyRejectedPrimaryKey :: KeyPkt 'SecretPkt
    , authSecretSubkeyRejectedValue :: Maybe (KeyPkt 'SecretPkt)
    , authSecretSubkeyRejectedUIDs :: [AuthSecretSubkeyUID]
    , authSecretSubkeyRejectedPrimaryUID :: Maybe Text
    , authSecretSubkeyRejectedReason :: AuthSecretSubkeyRejectionReason
    }
    deriving (Eq, Show)

data AuthSecretSubkeysAtReport
    = AuthSecretSubkeysAtReport
    { authSecretSubkeysAccepted :: [AuthSecretSubkeyAtTime]
    , authSecretSubkeysRejected :: [AuthSecretSubkeyRejectedAtTime]
    }
    deriving (Eq, Show)

conduitToAuthSecretSubkeysAtReport
    :: Monad m
    => UTCTime
    -> ConduitT (TK 'SecretTK) AuthSecretSubkeysAtReport m ()
conduitToAuthSecretSubkeysAtReport validationTime =
    CL.map (authSecretSubkeysAtReport validationTime)

conduitToAuthSecretSubkeysAt
    :: Monad m
    => UTCTime
    -> ConduitT (TK 'SecretTK) AuthSecretSubkeyAtTime m ()
conduitToAuthSecretSubkeysAt validationTime =
    CL.concatMap
        ( authSecretSubkeysAccepted
            . authSecretSubkeysAtReport validationTime
        )

authSecretSubkeysAt
    :: UTCTime -> TK 'SecretTK -> [AuthSecretSubkeyAtTime]
authSecretSubkeysAt validationTime =
    authSecretSubkeysAccepted
        . authSecretSubkeysAtReport validationTime

authSecretSubkeysAtReport
    :: UTCTime -> TK 'SecretTK -> AuthSecretSubkeysAtReport
authSecretSubkeysAtReport validationTime typedTk =
    case verifyTKWith
        ( verifySigWith
            defaultVerificationPolicy
            (verifyAgainstKeys [untyped])
        )
        (Just validationTime)
        typedTk of
        Left _ ->
            AuthSecretSubkeysAtReport
                []
                [ AuthSecretSubkeyRejectedAtTime
                    { authSecretSubkeyRejectedPrimaryKey = primaryKey
                    , authSecretSubkeyRejectedValue = Nothing
                    , authSecretSubkeyRejectedUIDs = []
                    , authSecretSubkeyRejectedPrimaryUID = Nothing
                    , authSecretSubkeyRejectedReason =
                        AuthSecretSubkeyTKVerificationFailed
                    }
                ]
        Right verifiedTk
            | not (isTKTimeValid validationTime verifiedTk) ->
                let uids = uidContextsAt validationTime verifiedTk
                    primaryUid =
                        authSecretSubkeyUIDValue
                            <$> find authSecretSubkeyUIDIsPrimary uids
                 in AuthSecretSubkeysAtReport
                        []
                        [ AuthSecretSubkeyRejectedAtTime
                            { authSecretSubkeyRejectedPrimaryKey = primaryKey
                            , authSecretSubkeyRejectedValue = Nothing
                            , authSecretSubkeyRejectedUIDs = uids
                            , authSecretSubkeyRejectedPrimaryUID = primaryUid
                            , authSecretSubkeyRejectedReason =
                                AuthSecretSubkeyPrimaryKeyInvalidAtTime
                            }
                        ]
            | otherwise ->
                let uids = uidContextsAt validationTime verifiedTk
                    primaryUid =
                        authSecretSubkeyUIDValue
                            <$> find authSecretSubkeyUIDIsPrimary uids
                 in foldr
                        ( \subCandidate acc ->
                            case classifySecretSubkeyAtTime
                                validationTime
                                primaryKey
                                uids
                                primaryUid
                                subCandidate of
                                Left rejected ->
                                    acc
                                        { authSecretSubkeysRejected =
                                            rejected : authSecretSubkeysRejected acc
                                        }
                                Right accepted ->
                                    acc
                                        { authSecretSubkeysAccepted =
                                            accepted : authSecretSubkeysAccepted acc
                                        }
                        )
                        (AuthSecretSubkeysAtReport [] [])
                        (_tkSubs verifiedTk)
  where
    untyped = publicViewTK typedTk
    primaryKey = _tkPrimaryKey typedTk

classifySecretSubkeyAtTime
    :: UTCTime
    -> KeyPkt 'SecretPkt
    -> [AuthSecretSubkeyUID]
    -> Maybe Text
    -> (KeyPkt 'SecretPkt, [SignaturePayload])
    -> Either AuthSecretSubkeyRejectedAtTime AuthSecretSubkeyAtTime
classifySecretSubkeyAtTime validationTime primaryKey uids primaryUid (subkey, sigs)
    | keyPktRole subkey /= KeyPktSubkey =
        Left
            AuthSecretSubkeyRejectedAtTime
                { authSecretSubkeyRejectedPrimaryKey = primaryKey
                , authSecretSubkeyRejectedValue = Just subkey
                , authSecretSubkeyRejectedUIDs = uids
                , authSecretSubkeyRejectedPrimaryUID = primaryUid
                , authSecretSubkeyRejectedReason = AuthSecretSubkeyNotSubkeyPacket
                }
    | not
        ( isPKTimeValidWithSelfSignatures
            validationTime
            (keyPktPKPayload subkey)
            sigs
        ) =
        Left
            AuthSecretSubkeyRejectedAtTime
                { authSecretSubkeyRejectedPrimaryKey = primaryKey
                , authSecretSubkeyRejectedValue = Just subkey
                , authSecretSubkeyRejectedUIDs = uids
                , authSecretSubkeyRejectedPrimaryUID = primaryUid
                , authSecretSubkeyRejectedReason =
                    AuthSecretSubkeySubkeyInvalidAtTime
                }
    | not (subkeyAuthCapableAt validationTime sigs) =
        Left
            AuthSecretSubkeyRejectedAtTime
                { authSecretSubkeyRejectedPrimaryKey = primaryKey
                , authSecretSubkeyRejectedValue = Just subkey
                , authSecretSubkeyRejectedUIDs = uids
                , authSecretSubkeyRejectedPrimaryUID = primaryUid
                , authSecretSubkeyRejectedReason =
                    AuthSecretSubkeyMissingAuthCapability
                }
    | otherwise =
        Right
            AuthSecretSubkeyAtTime
                { authSecretSubkeyPrimaryKey = primaryKey
                , authSecretSubkeyValue = subkey
                , authSecretSubkeyUIDs = uids
                , authSecretSubkeyPrimaryUID = primaryUid
                }

uidContextsAt :: UTCTime -> TK k -> [AuthSecretSubkeyUID]
uidContextsAt validationTime tk =
    map
        ( \(uid, _) ->
            AuthSecretSubkeyUID
                { authSecretSubkeyUIDValue = uid
                , authSecretSubkeyUIDIsPrimary = Just uid == primaryUid
                }
        )
        (tk ^. tkUIDs)
  where
    primaryUid = primaryUIDAt validationTime (tk ^. tkUIDs)

primaryUIDAt
    :: UTCTime -> [(Text, [SignaturePayload])] -> Maybe Text
primaryUIDAt validationTime uids =
    snd <$> newestByCreationTime candidates
  where
    candidates =
        [ (createdAt, uid)
        | (uid, sigs) <- uids
        , cert <-
            maybeToList (latestEffectiveCertificationAt validationTime sigs)
        , signatureMarksPrimaryUID cert
        , createdAt <- maybeToList (signatureCreationTime cert)
        ]

subkeyAuthCapableAt :: UTCTime -> [SignaturePayload] -> Bool
subkeyAuthCapableAt validationTime sigs =
    maybe
        False
        signatureHasAuthKeyFlag
        (latestEffectiveBindingSignatureAt validationTime sigs)

latestEffectiveSignatureAt
    :: (SignaturePayload -> Bool)
    -> UTCTime
    -> [SignaturePayload]
    -> Maybe SignaturePayload
latestEffectiveSignatureAt typePred validationTime sigs =
    snd
        <$> newestByCreationTime
            [ (createdAt, sig)
            | sig <- sigs
            , typePred sig
            , signatureEffectiveAt validationTime sig
            , createdAt <- maybeToList (signatureCreationTime sig)
            ]

latestEffectiveBindingSignatureAt
    :: UTCTime -> [SignaturePayload] -> Maybe SignaturePayload
latestEffectiveBindingSignatureAt =
    latestEffectiveSignatureAt isSubkeyBindingSig

latestEffectiveCertificationAt
    :: UTCTime -> [SignaturePayload] -> Maybe SignaturePayload
latestEffectiveCertificationAt =
    latestEffectiveSignatureAt isCertificationSig

signatureHasAuthKeyFlag :: SignaturePayload -> Bool
signatureHasAuthKeyFlag sig =
    Set.member AuthKey (signatureKeyFlags sig)

foldHasSubPacket
    :: (SigSubPacket -> Maybe b)
    -> (b -> Bool)
    -> SignaturePayload
    -> Bool
foldHasSubPacket extract predicate sig =
    any
        predicate
        ( mapMaybe
            extract
            (maybe [] id (signatureHashedSubpacketsKnown sig))
        )

signatureKeyFlags :: SignaturePayload -> Set.Set KeyFlag
signatureKeyFlags sig =
    foldr
        ( \sp acc ->
            case sp of
                SigSubPacket _ (KeyFlags flags) -> Set.union flags acc
                _ -> acc
        )
        Set.empty
        (maybe [] id (signatureHashedSubpacketsKnown sig))

signatureMarksPrimaryUID :: SignaturePayload -> Bool
signatureMarksPrimaryUID sig =
    foldHasSubPacket
        (\case SigSubPacket _ (PrimaryUserId p) -> Just p; _ -> Nothing)
        id
        sig

{-# DEPRECATED conduitToTKsEither "Use conduitToSomeTKsEither instead" #-}
conduitToTKsEither
    :: Monad m
    => ConduitT
        Pkt
        (Either KeyringChunkParseError (Maybe TKUnknown))
        m
        ()
conduitToTKsEither = conduitToTKsEither' True

{-# DEPRECATED
    conduitToTKsDroppingEither
    "Use conduitToSomeTKsDroppingEither instead"
    #-}
conduitToTKsDroppingEither
    :: Monad m
    => ConduitT
        Pkt
        (Either KeyringChunkParseError (Maybe TKUnknown))
        m
        ()
conduitToTKsDroppingEither = conduitToTKsEither' False

conduitToTKsWithWireRepEither
    :: Monad m
    => ConduitT
        PktWithWireRep
        (Either KeyringChunkParseError (Maybe TKWithWireRep))
        m
        ()
conduitToTKsWithWireRepEither = conduitToTKsWithWireRepEither' True

conduitToTKsDroppingWithWireRepEither
    :: Monad m
    => ConduitT
        PktWithWireRep
        (Either KeyringChunkParseError (Maybe TKWithWireRep))
        m
        ()
conduitToTKsDroppingWithWireRepEither = conduitToTKsWithWireRepEither' False

fakecmAccumEither
    :: Monad m
    => (accum -> Either e (accum, [b]))
    -> (a -> accum -> Either e (accum, [b]))
    -> accum
    -> ConduitT a (Either e b) m ()
fakecmAccumEither finalizer f initialAccum = loop initialAccum
  where
    loop accum =
        await
            >>= maybe
                ( case finalizer accum of
                    Left err -> yield (Left err)
                    Right (_, bs) -> mapM_ (yield . Right) bs
                )
                go
      where
        go a = do
            case f a accum of
                Left err -> do
                    yield (Left err)
                    loop initialAccum
                Right (accum', bs) -> do
                    mapM_ (yield . Right) bs
                    loop accum'

conduitDropErrorsAndNothings
    :: Monad m => ConduitT (Either e (Maybe a)) a m ()
conduitDropErrorsAndNothings =
    CL.mapMaybe (join . hush)

conduitToTKsEither'
    :: Monad m
    => Bool
    -> ConduitT
        Pkt
        (Either KeyringChunkParseError (Maybe TKUnknown))
        m
        ()
conduitToTKsEither' intolerant =
    CL.filter notTrustPacket
        .| CL.map (: [])
        .| fakecmAccumEither
            finalizeParsingEither
            (parseAChunkEither (anyTK intolerant))
            ([], Just (Nothing, anyTK intolerant))
  where
    notTrustPacket = not . isTrustPkt

conduitToTKsWithWireRepEither'
    :: Monad m
    => Bool
    -> ConduitT
        PktWithWireRep
        (Either KeyringChunkParseError (Maybe TKWithWireRep))
        m
        ()
conduitToTKsWithWireRepEither' intolerant =
    CL.filter notTrustPacket
        .| CL.map (: [])
        .| fakecmAccumEither
            finalizeParsingEither
            (parseAChunkEither (anyTKWithWireRep intolerant))
            ([], Just (Nothing, anyTKWithWireRep intolerant))
  where
    notTrustPacket = not . isTrustPkt . (^. pktWireRep . pktValue)

sinkPublicKeyringMap
    :: Monad m => ConduitT (TK 'PublicTK) Void m PublicKeyring
sinkPublicKeyringMap = CL.fold (flip insert) empty

sinkSecretKeyringMap
    :: Monad m => ConduitT (TK 'SecretTK) Void m SecretKeyring
sinkSecretKeyringMap = CL.fold (flip insert) empty

-- | Lift a single typed TK into its kinded keyring
publicTKToKeyring :: TK 'PublicTK -> PublicKeyring
publicTKToKeyring tk = insert tk empty

secretTKToKeyring :: TK 'SecretTK -> SecretKeyring
secretTKToKeyring tk = insert tk empty

-- | Partition a list of SomeTK into homogeneous public and secret keyrings
partitionSomeTKs :: [SomeTK] -> (PublicKeyring, SecretKeyring)
partitionSomeTKs = foldr step (empty, empty)
  where
    step (SomePublicTK tk) (pub, sec) = (insert tk pub, sec)
    step (SomeSecretTK tk) (pub, sec) = (pub, insert tk sec)
