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 FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeSynonymInstances #-}

module Data.Conduit.OpenPGP.Keyring.Instances
  (
  ) where

import Codec.Encryption.OpenPGP.Fingerprint (eightOctetKeyID, fingerprint)
import Codec.Encryption.OpenPGP.Internal (issuer)
import Codec.Encryption.OpenPGP.SignatureQualities (sigCT)
import Codec.Encryption.OpenPGP.Types
import Control.Arrow (second)
import Control.Lens ((^.), (^..), _1, folded)
import Data.Data.Lens (biplate)
import Data.Either (rights)
import Data.Function (on)
import qualified Data.HashMap.Lazy as HashMap
import Data.IxSet.Typed (Indexable(..), ixFun, ixList)
import Data.List (nub, sort)
import qualified Data.List.NonEmpty as NE
import qualified Data.Map as Map
import Data.Semigroup (Semigroup, (<>))
import Data.Text (Text)

instance Indexable KeyringIxs TKUnknown where
  indices = ixList (ixFun getEOKIs) (ixFun getFingerprints) (ixFun getUIDs)

getEOKIs :: TKUnknown -> [EightOctetKeyId]
getEOKIs tk = rights (map eightOctetKeyID (tk ^.. biplate :: [SomePKPayload]))

getFingerprints :: TKUnknown -> [Fingerprint]
getFingerprints tk = map fingerprint (tk ^.. biplate :: [SomePKPayload])

getUIDs :: TKUnknown -> [Text]
getUIDs tk = (tk ^. tkuUIDs) ^.. folded . _1

instance Semigroup TKUnknown where
  (<>) a b =
    TKUnknown
      (_tkuKey a)
      (nub . sort $ _tkuRevs a ++ _tkuRevs b)
      ((kvmerge `on` _tkuUIDs) a b)
      ((kvmerge `on` _tkuUAts) a b)
      ((ukvmerge `on` _tkuSubs) a b)
    where
      kvmerge x y =
        Map.toList (Map.unionWith nsa (Map.fromList x) (Map.fromList y))
      ukvmerge x y =
        HashMap.toList
          (HashMap.unionWith nsa (HashMap.fromList x) (HashMap.fromList y))
      nsa x y = nub . sort $ x ++ y

instance Semigroup TKWithWireRep where
  (<>) a b =
    let mergedTK = _tkValue a <> _tkValue b
        mergedPackets =
          selectPacketRefsByValue
            (flattenTKPackets mergedTK)
            (dedupePacketRefsById (_tkPackets a ++ _tkPackets b))
     in TKWithWireRep
          (mergeWireRepRefs (_tkWireRepRefs a) (_tkWireRepRefs b))
          (mergedWireRepRange mergedPackets)
          mergedPackets
          mergedTK

flattenTKPackets :: TKUnknown -> [Pkt]
flattenTKPackets tk =
  [someKeyPktToPkt (mkPrimaryKeyPkt pkp mska)] ++
  map SignaturePkt (_tkuRevs tk) ++
  concatMap flattenUID (_tkuUIDs tk) ++
  concatMap flattenUAT (_tkuUAts tk) ++
  concatMap flattenSub (_tkuSubs tk)
  where
    (pkp, mska) = _tkuKey tk
    flattenUID (uid, sigs) = UserIdPkt uid : map SignaturePkt sigs
    flattenUAT (uat, sigs) = UserAttributePkt uat : map SignaturePkt sigs
    flattenSub (pkt, sigs) = pkt : map SignaturePkt sigs

mergeWireRepRefs :: WireRepRefs -> WireRepRefs -> WireRepRefs
mergeWireRepRefs left right =
  case dedupe (NE.toList left ++ NE.toList right) of
    [] -> left
    (x:xs) -> x NE.:| xs
  where
    dedupe [] = []
    dedupe (x:xs) = x : dedupe (filter (/= x) xs)

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 _ [] = Nothing
extractFirstByValue expected (pkt:rest)
  | _pktValue pkt == expected = Just (pkt, rest)
  | otherwise = second (pkt :) <$> extractFirstByValue expected rest

-- | Extract all SomePKPayloads from a TK (primary + subkeys) without biplate
tkPKPayloads :: TK k -> [SomePKPayload]
tkPKPayloads tk =
  keyPktPKPayload (_tkPrimaryKey tk) :
  map (keyPktPKPayload . 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


instance Semigroup (TK k) where
  a <> b =
    TK
      (_tkPrimaryKey a)
      (nub . sort $ _tkRevs a ++ _tkRevs b)
      ((kvmerge `on` _tkUIDs) a b)
      ((kvmerge `on` _tkUAts) a b)
      ((ukvmerge `on` _tkSubs) a b)
    where
      kvmerge x y =
        Map.toList (Map.unionWith nsa (Map.fromList x) (Map.fromList y))
      ukvmerge x y =
        HashMap.toList
          (HashMap.unionWith nsa (HashMap.fromList x) (HashMap.fromList y))
      nsa x y = nub . sort $ x ++ y