-- CryptoECDH.hs: OpenPGP (RFC9580) ECDH helper utilities
-- Copyright © 2012-2026  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).

{-# LANGUAGE OverloadedStrings #-}

module Codec.Encryption.OpenPGP.Internal.CryptoECDH
  ( normalizeMontgomeryPublic
  , buildECDHKDFParam
  , deriveECDHKek
  ) where

import Codec.Encryption.OpenPGP.BlockCipher (keySize, renderCipherError)
import Codec.Encryption.OpenPGP.Fingerprint (fingerprint)
import Codec.Encryption.OpenPGP.Internal (curveFromCurve, curveToCurveoidBS, leftPadTo)
import Codec.Encryption.OpenPGP.Policy (ecdhKdfHashDigest)
import Codec.Encryption.OpenPGP.Types
import qualified Crypto.PubKey.ECC.ECDSA as ECDSA
import Data.Bifunctor (first)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL

normalizeMontgomeryPublic ::
     Int
  -> String
  -> B.ByteString
  -> Either String B.ByteString
normalizeMontgomeryPublic targetLen label bs
  | B.length bs == targetLen = Right bs
  | B.length bs < targetLen = Right (leftPadTo targetLen bs)
  | B.length bs == targetLen + 1 && B.head bs == 0x40 = Right (B.tail bs)
  | otherwise = Left (label ++ show (B.length bs))

buildECDHKDFParam ::
     SomePKPayload
  -> PubKeyAlgorithm
  -> PKey
  -> HashAlgorithm
  -> SymmetricAlgorithm
  -> Either String B.ByteString
buildECDHKDFParam recipientPKP pka recipientECDHPub kdfHA kdfSA =
  (<>
   B.pack [fromFVal pka, 0x03, 0x01, fromFVal kdfHA, fromFVal kdfSA] <>
   "Anonymous Sender    " <>
   BL.toStrict (unFingerprint (fingerprint recipientPKP))) <$>
  encodedCurveOid
  where
    encodedCurveOid = ((\oid -> B.singleton (fromIntegral (B.length oid)) <> oid) <$>) curveOid
    curveOid =
      case recipientECDHPub of
        ECDSAPubKey (ECDSA_PublicKey (ECDSA.PublicKey curve _)) ->
          curveToCurveoidBS (curveFromCurve curve)
        EdDSAPubKey Ed25519 _ ->
          curveToCurveoidBS Curve25519
        EdDSAPubKey Ed448 _ ->
          curveToCurveoidBS Curve448
        _ -> Left "ECDH KDF param requires ECDH recipient key"

deriveECDHKek ::
     HashAlgorithm
  -> SymmetricAlgorithm
  -> B.ByteString
  -> B.ByteString
  -> Either String B.ByteString
deriveECDHKek kdfHA kdfSA sharedSecret kdfParam = do
  digest <- ecdhKdfHashDigest kdfHA (B.pack [0, 0, 0, 1] <> sharedSecret <> kdfParam)
  kekLen <- first renderCipherError (keySize kdfSA)
  if B.length digest < kekLen
    then Left "ECDH KDF digest is shorter than required KEK length"
    else Right (B.take kekLen digest)
