Instances.hs

-- Instances.hs: OpenPGP (RFC9580) additional types for transferable keys
-- 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 FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeSynonymInstances #-}

module Data.Conduit.OpenPGP.Keyring.Instances
    ( flattenTKPackets
    , flattenTKPacketsMixedTK
    , getEOKIsMixed
    , getFingerprintsMixed
    , getUIDsMixed
    ) where

import Control.Arrow (second)
import Control.Lens (folded, (^.), (^..), _1)
import Data.Containers.ListUtils (nubOrd)
import Data.Data.Lens (biplate)
import Data.Either (rights)
import Data.IxSet.Typed (Indexable (..), ixFun, ixList)
import qualified Data.List.NonEmpty as NE
import Data.Text (Text)

import Codec.Encryption.OpenPGP.Fingerprint
    ( eightOctetKeyID
    , fingerprint
    )
import Codec.Encryption.OpenPGP.Types

instance Indexable KeyringIxs (TK 'MixedTK) where
    indices =
        ixList
            (ixFun getEOKIsMixed)
            (ixFun getFingerprintsMixed)
            (ixFun getUIDsMixed)

getEOKIsMixed :: TK 'MixedTK -> [EightOctetKeyId]
getEOKIsMixed tk =
    rights (map eightOctetKeyID (tk ^.. biplate :: [SomePKPayload]))

getFingerprintsMixed :: TK 'MixedTK -> [Fingerprint]
getFingerprintsMixed tk = map fingerprint (tk ^.. biplate :: [SomePKPayload])

getUIDsMixed :: TK 'MixedTK -> [Text]
getUIDsMixed tk = (tk ^. tkUIDs) ^.. folded . _1

instance Semigroup TKWithWireRep where
    (<>) a b =
        let mergedTK = _tkValue a <> _tkValue b
            mergedPackets =
                selectPacketRefsByValue
                    (flattenTKPackets mergedTK)
                    (dedupePacketRefsById (_tkPackets a ++ _tkPackets b))
            mergedRefs =
                case nubOrd
                    (NE.toList (_tkWireRepRefs a) ++ NE.toList (_tkWireRepRefs b)) of
                    [] -> _tkWireRepRefs a
                    (x : xs) -> x NE.:| xs
         in TKWithWireRep
                mergedRefs
                (mergedWireRepRange mergedPackets)
                mergedPackets
                mergedTK

flattenTKPackets :: TK 'MixedTK -> [Pkt]
flattenTKPackets = flattenTKPacketsMixedTK

flattenTKPacketsMixedTK :: TK 'MixedTK -> [Pkt]
flattenTKPacketsMixedTK tk =
    [someKeyPktToPkt (_tkPrimaryKey tk)]
        ++ map SignaturePkt (_tkRevs tk)
        ++ map SignaturePkt (_tkDirectKeySigs tk)
        ++ concatMap flattenUID (_tkUIDs tk)
        ++ concatMap flattenUAt (_tkUAts tk)
        ++ concatMap flattenSubMixed (_tkSubs tk)
  where
    flattenUID (uid, sigs) = UserIdPkt uid : map SignaturePkt sigs
    flattenUAt (uat, sigs) = UserAttributePkt uat : map SignaturePkt sigs
    flattenSubMixed (kp, sigs) = someKeyPktToPkt kp : map SignaturePkt sigs

mergedWireRepRange :: [PktWithWireRep] -> Maybe ByteRange
mergedWireRepRange [] = Nothing
mergedWireRepRange (pkt : rest)
    | all ((== _pktWireRepRef pkt) . _pktWireRepRef) rest =
        spanByteRanges (map _pktRange (pkt : rest))
    | otherwise = Nothing

dedupePacketRefsById :: [PktWithWireRep] -> [PktWithWireRep]
dedupePacketRefsById = go []
  where
    go _ [] = []
    go seen (pkt : rest) =
        let packetRefId = packetRefIdOf pkt
         in if packetRefId `elem` seen
                then go seen rest
                else pkt : go (packetRefId : seen) rest

selectPacketRefsByValue
    :: [Pkt] -> [PktWithWireRep] -> [PktWithWireRep]
selectPacketRefsByValue expected available = go expected available []
  where
    go [] _ acc = reverse acc
    go (pkt : pktRest) refs acc =
        case extractFirstByValue pkt refs of
            Nothing ->
                error
                    ( "TKWithWireRep Semigroup merge missing packet reference for tag "
                        ++ show (pktTag pkt)
                    )
            Just (matched, remaining) -> go pktRest remaining (matched : acc)

extractFirstByValue
    :: Pkt
    -> [PktWithWireRep]
    -> Maybe (PktWithWireRep, [PktWithWireRep])
extractFirstByValue expected = go []
  where
    go _ [] = Nothing
    go seen (pkt : rest)
        | pkt ^. pktWireRep . pktValue == expected =
            Just (pkt, seen ++ rest)
        | otherwise = second (pkt :) <$> go seen rest

-- | Extract all SomePKPayloads from a TK (primary + subkeys) without biplate
class TKPKPayloads (k :: TKKind) where
    tkPKPayloads :: TK k -> [SomePKPayload]

instance TKPKPayloads 'PublicTK where
    tkPKPayloads tk =
        keyPktPKPayload (_tkPrimaryKey tk)
            : map (keyPktPKPayload . fst) (_tkSubs tk)

instance TKPKPayloads 'SecretTK where
    tkPKPayloads tk =
        keyPktPKPayload (_tkPrimaryKey tk)
            : map (keyPktPKPayload . fst) (_tkSubs tk)

instance TKPKPayloads 'MixedTK where
    tkPKPayloads tk =
        someKeyPktPKPayload (_tkPrimaryKey tk)
            : map (someKeyPktPKPayload . fst) (_tkSubs tk)

-- | Index public TKs by key ID, fingerprint, and UID
instance Indexable KeyringIxs (TK 'PublicTK) where
    indices =
        ixList
            (ixFun getEOKIsPublic)
            (ixFun getFingerprintsPublic)
            (ixFun getUIDsPublic)

getEOKIsPublic :: TK 'PublicTK -> [EightOctetKeyId]
getEOKIsPublic tk = rights (map eightOctetKeyID (tkPKPayloads tk))

getFingerprintsPublic :: TK 'PublicTK -> [Fingerprint]
getFingerprintsPublic tk = map fingerprint (tkPKPayloads tk)

getUIDsPublic :: TK 'PublicTK -> [Text]
getUIDsPublic tk = (tk ^. tkUIDs) ^.. folded . _1

-- | Index secret TKs by key ID, fingerprint, and UID
instance Indexable KeyringIxs (TK 'SecretTK) where
    indices =
        ixList
            (ixFun getEOKIsSecret)
            (ixFun getFingerprintsSecret)
            (ixFun getUIDsSecret)

getEOKIsSecret :: TK 'SecretTK -> [EightOctetKeyId]
getEOKIsSecret tk = rights (map eightOctetKeyID (tkPKPayloads tk))

getFingerprintsSecret :: TK 'SecretTK -> [Fingerprint]
getFingerprintsSecret tk = map fingerprint (tkPKPayloads tk)

getUIDsSecret :: TK 'SecretTK -> [Text]
getUIDsSecret tk = (tk ^. tkUIDs) ^.. folded . _1