-- 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 :: Int -> String -> ByteString -> Either String ByteString
normalizeMontgomeryPublic Int
targetLen String
label ByteString
bs
    | ByteString -> Int
B.length ByteString
bs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
targetLen = ByteString -> Either String ByteString
forall a b. b -> Either a b
Right ByteString
bs
    | ByteString -> Int
B.length ByteString
bs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
targetLen = ByteString -> Either String ByteString
forall a b. b -> Either a b
Right (Int -> ByteString -> ByteString
leftPadTo Int
targetLen ByteString
bs)
    | ByteString -> Int
B.length ByteString
bs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
targetLen Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Bool -> Bool -> Bool
&& HasCallStack => ByteString -> Word8
ByteString -> Word8
B.head ByteString
bs Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x40 =
        ByteString -> Either String ByteString
forall a b. b -> Either a b
Right (HasCallStack => ByteString -> ByteString
ByteString -> ByteString
B.tail ByteString
bs)
    | Bool
otherwise = String -> Either String ByteString
forall a b. a -> Either a b
Left (String
label String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (ByteString -> Int
B.length ByteString
bs))

buildECDHKDFParam
    :: SomePKPayload
    -> PubKeyAlgorithm
    -> PKey
    -> HashAlgorithm
    -> SymmetricAlgorithm
    -> Either String B.ByteString
buildECDHKDFParam :: SomePKPayload
-> PubKeyAlgorithm
-> PKey
-> HashAlgorithm
-> SymmetricAlgorithm
-> Either String ByteString
buildECDHKDFParam SomePKPayload
recipientPKP PubKeyAlgorithm
pka PKey
recipientECDHPub HashAlgorithm
kdfHA SymmetricAlgorithm
kdfSA =
    ( ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<>
        [Word8] -> ByteString
B.pack [PubKeyAlgorithm -> Word8
forall a. FutureVal a => a -> Word8
fromFVal PubKeyAlgorithm
pka, Word8
0x03, Word8
0x01, HashAlgorithm -> Word8
forall a. FutureVal a => a -> Word8
fromFVal HashAlgorithm
kdfHA, SymmetricAlgorithm -> Word8
forall a. FutureVal a => a -> Word8
fromFVal SymmetricAlgorithm
kdfSA]
            ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
"Anonymous Sender    "
            ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> LazyByteString -> ByteString
BL.toStrict (Fingerprint -> LazyByteString
unFingerprint (SomePKPayload -> Fingerprint
fingerprint SomePKPayload
recipientPKP))
    )
        (ByteString -> ByteString)
-> Either String ByteString -> Either String ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either String ByteString
encodedCurveOid
  where
    encodedCurveOid :: Either String ByteString
encodedCurveOid =
        ((\ByteString
oid -> Word8 -> ByteString
B.singleton (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
B.length ByteString
oid)) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
oid) (ByteString -> ByteString)
-> Either String ByteString -> Either String ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>)
            Either String ByteString
curveOid
    curveOid :: Either String ByteString
curveOid =
        case PKey
recipientECDHPub of
            ECDSAPubKey (ECDSA_PublicKey (ECDSA.PublicKey Curve
curve PublicPoint
_)) ->
                ECCCurve -> Either String ByteString
curveToCurveoidBS (Curve -> ECCCurve
curveFromCurve Curve
curve)
            EdDSAPubKey EdSigningCurve
EdSigningCurve25519 EdPoint
_ ->
                ECCCurve -> Either String ByteString
curveToCurveoidBS ECCCurve
Curve25519
            EdDSAPubKey EdSigningCurve
EdSigningCurve448 EdPoint
_ ->
                ECCCurve -> Either String ByteString
curveToCurveoidBS ECCCurve
Curve448
            PKey
_ -> String -> Either String ByteString
forall a b. a -> Either a b
Left String
"ECDH KDF param requires ECDH recipient key"

deriveECDHKek
    :: HashAlgorithm
    -> SymmetricAlgorithm
    -> B.ByteString
    -> B.ByteString
    -> Either String B.ByteString
deriveECDHKek :: HashAlgorithm
-> SymmetricAlgorithm
-> ByteString
-> ByteString
-> Either String ByteString
deriveECDHKek HashAlgorithm
kdfHA SymmetricAlgorithm
kdfSA ByteString
sharedSecret ByteString
kdfParam = do
    digest <-
        HashAlgorithm -> ByteString -> Either String ByteString
ecdhKdfHashDigest
            HashAlgorithm
kdfHA
            ([Word8] -> ByteString
B.pack [Word8
0, Word8
0, Word8
0, Word8
1] ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
sharedSecret ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
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)