WKD.hs

-- WKD.hs: hOpenPGP key tool
-- Copyright © 2026  Clint Adams
--
-- vim: softtabstop=4:shiftwidth=4:expandtab
--
-- This program is free software: you can redistribute it and/or modify
-- it under the terms of the GNU Affero General Public License as
-- published by the Free Software Foundation, either version 3 of the
-- License, or (at your option) any later version.
--
-- This program is distributed in the hope that it will be useful,
-- but WITHOUT ANY WARRANTY; without even the implied warranty of
-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-- GNU Affero General Public License for more details.
--
-- You should have received a copy of the GNU Affero General Public License
-- along with this program.  If not, see <http://www.gnu.org/licenses/>.
{-# LANGUAGE OverloadedStrings #-}

module HOpenPGP.Tools.Common.WKD
    ( fetchKeys
    , parseMailbox
    ) where

import qualified Codec.Encryption.OpenPGP.ASCIIArmor as AA
import Codec.Encryption.OpenPGP.ASCIIArmor.Types
    ( Armor (Armor)
    , ArmorType (ArmorPublicKeyBlock)
    )
import Codec.Encryption.OpenPGP.Types (TKUnknown (..))
import Control.Arrow ((&&&))
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Trans.Except (ExceptT (..), throwE)
import qualified Crypto.Hash as CH
import qualified Crypto.Hash.Algorithms as CHA
import Data.Binary (get)
import Data.Bits (shiftL, shiftR, (.&.), (.|.))
import qualified Data.ByteArray as BA
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as BC8
import qualified Data.ByteString.Lazy as BL
import Data.Conduit (runConduitRes, (.|))
import qualified Data.Conduit.Binary as CB
import qualified Data.Conduit.List as CL
import Data.Conduit.OpenPGP.Keyring (conduitToTKsDroppingEither)
import Data.Conduit.Serialization.Binary (conduitGet)
import Data.Either (rights)
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import Data.Time.Clock.POSIX (getPOSIXTime)
import Data.Word (Word8)
import Network.HTTP.Client
    ( Manager
    , Response (..)
    , httpLbs
    , newManager
    , parseUrlThrow
    , setQueryString
    )
import Network.HTTP.Client.TLS (tlsManagerSettings)
import Network.HTTP.Types.Status (ok200)

import HOpenPGP.Tools.Common.HKP (FetchValidationMethod (..))
import HOpenPGP.Tools.Common.TKUtils (processTK)

fetchKeys
    :: FetchValidationMethod
    -> Text
    -> ExceptT String IO [TKUnknown]
fetchKeys fvm mailbox = do
    parsedMailbox <- ExceptT . return $ parseMailbox mailbox
    manager <- liftIO $ newManager tlsManagerSettings
    response <- fetchWKD manager parsedMailbox
    body <-
        if responseStatus response == ok200
            then return (responseBody response)
            else throwE ("HTTP status: " ++ show (responseStatus response))
    validateAndFilterKeys fvm parsedMailbox body

parseMailbox :: Text -> Either String (Text, Text)
parseMailbox rawMailbox =
    let mailbox = T.strip rawMailbox
        parts = T.splitOn "@" mailbox
     in case parts of
            [localPart, domain]
                | T.null localPart -> Left "mailbox local part cannot be empty"
                | T.null domain -> Left "mailbox domain cannot be empty"
                | T.any (== ' ') mailbox -> Left "mailbox cannot contain spaces"
                | otherwise -> Right (localPart, T.toLower domain)
            _ -> Left "mailbox must contain exactly one @"

fetchWKD
    :: Manager
    -> (Text, Text)
    -> ExceptT String IO (Response BL.ByteString)
fetchWKD manager (localPart, domain) = do
    let localPartLower = T.toLower localPart
        hu = BC8.unpack . zbase32 . sha1 . TE.encodeUtf8 $ localPartLower
        advancedUrl =
            "https://openpgpkey."
                <> T.unpack domain
                <> "/.well-known/openpgpkey/"
                <> T.unpack domain
                <> "/hu/"
                <> hu
        directUrl =
            "https://"
                <> T.unpack domain
                <> "/.well-known/openpgpkey/hu/"
                <> hu
        mailboxParam = TE.encodeUtf8 localPartLower
        withMailbox req = setQueryString [("l", Just mailboxParam)] req
    advancedRequest <- liftIO $ parseUrlThrow advancedUrl
    advancedResponse <-
        liftIO $ httpLbs (withMailbox advancedRequest) manager
    if responseStatus advancedResponse == ok200
        then return advancedResponse
        else do
            directRequest <- liftIO $ parseUrlThrow directUrl
            liftIO $ httpLbs (withMailbox directRequest) manager

validateAndFilterKeys
    :: FetchValidationMethod
    -> (Text, Text)
    -> BL.ByteString
    -> ExceptT String IO [TKUnknown]
validateAndFilterKeys fvm mailbox body = do
    keys <- decodeWkdResponse body
    cpt <- liftIO getPOSIXTime
    let processedKeys =
            rights $
                map (uncurry (liftA2 (,)) . (pure &&& processTK (Just cpt))) keys
        mailboxFiltered = filter (mailboxMatchesKey mailbox . snd) processedKeys
    return $
        map fst $
            case fvm of
                AnySelfSigned -> processedKeys
                MatchPrimaryKeyFingerprint -> mailboxFiltered
                MatchPrimaryOrAnySubkeyFingerprint -> mailboxFiltered

decodeWkdResponse
    :: BL.ByteString -> ExceptT String IO [TKUnknown]
decodeWkdResponse body =
    if isArmored body
        then decodeArmored body
        else decodeBinary body

decodeBinary :: BL.ByteString -> ExceptT String IO [TKUnknown]
decodeBinary bytes =
    liftIO . runConduitRes $
        CB.sourceLbs bytes
            .| conduitGet get
            .| conduitToTKsDroppingEither
            .| CL.mapFoldable id
            .| CL.mapMaybe id
            .| CL.consume

decodeArmored :: BL.ByteString -> ExceptT String IO [TKUnknown]
decodeArmored larmors = do
    bytestrings <-
        ExceptT . return $
            fmap (mconcat . map armorToBS) (AA.decodeLazy larmors)
    liftIO . runConduitRes $
        CB.sourceLbs bytestrings
            .| conduitGet get
            .| conduitToTKsDroppingEither
            .| CL.mapFoldable id
            .| CL.mapMaybe id
            .| CL.consume
  where
    armorToBS (Armor ArmorPublicKeyBlock _ bs) = bs
    armorToBS _ = mempty

isArmored :: BL.ByteString -> Bool
isArmored =
    BC8.isPrefixOf "-----BEGIN PGP PUBLIC KEY BLOCK-----"
        . BL.toStrict
        . BL.take 40

mailboxMatchesKey :: (Text, Text) -> TKUnknown -> Bool
mailboxMatchesKey (localPart, domain) tk =
    let mailbox = T.toLower (localPart <> "@" <> domain)
        bracketedMailbox = "<" <> mailbox <> ">"
     in any
            ( \uid ->
                let lowered = T.toLower uid
                 in lowered == mailbox || bracketedMailbox `T.isInfixOf` lowered
            )
            (map fst (_tkuUIDs tk))

sha1 :: B.ByteString -> B.ByteString
sha1 bs = BA.convert (CH.hashWith CHA.SHA1 bs :: CH.Digest CHA.SHA1)

zbase32 :: B.ByteString -> B.ByteString
zbase32 = BC8.pack . encodeZBase32 . B.unpack

encodeZBase32 :: [Word8] -> String
encodeZBase32 = go 0 0
  where
    alphabet = "ybndrfg8ejkmcpqxot1uwisza345h769"
    pick i = alphabet !! i
    go _ 0 [] = []
    go acc bits [] =
        [pick (fromIntegral (((acc `shiftL` (5 - bits)) .&. 31) :: Int))]
    go acc bits (x : xs)
        | bits >= 5 =
            pick (fromIntegral (((acc `shiftR` (bits - 5)) .&. 31) :: Int))
                : go acc (bits - 5) (x : xs)
        | otherwise =
            go ((acc `shiftL` 8) .|. fromIntegral x) (bits + 8) xs