-- 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 qualified Crypto.PubKey.ECC.ECDSA as ECDSA
import Data.Bifunctor (first)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL

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

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 EdSigningCurve25519 _ ->
                curveToCurveoidBS Curve25519
            EdDSAPubKey EdSigningCurve448 _ ->
                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)
