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 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.Bits ((.&.), (.|.), shiftL, shiftR)
import Data.Conduit ((.|), runConduitRes)
import qualified Data.Conduit.Binary as CB
import qualified Data.Conduit.List as CL
import Data.Conduit.OpenPGP.Keyring (conduitToTKsDropping)
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 HOpenPGP.Tools.Common.HKP (FetchValidationMethod(..))
import HOpenPGP.Tools.Common.TKUtils (processTK)
import Network.HTTP.Client
  ( Manager
  , Response(..)
  , httpLbs
  , newManager
  , parseUrlThrow
  , setQueryString
  )
import Network.HTTP.Client.TLS (tlsManagerSettings)
import Network.HTTP.Types.Status (ok200)

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 .| conduitToTKsDropping .| 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 .| conduitToTKsDropping .| 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