-- 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 DataKinds #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

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

import Control.Error.Util (note)
import qualified Crypto.PubKey.ECC.ECDSA as ECDSA
import Data.Bifunctor (first)
import qualified Data.ByteString as B
import GHC.TypeNats (KnownNat)

import Codec.Encryption.OpenPGP.BlockCipher
    ( keySize
    )
import Codec.Encryption.OpenPGP.Fingerprint (fingerprint)
import Codec.Encryption.OpenPGP.Internal
    ( FixedWidthBytes
    , bsToFixedWidth
    , byteWidth
    , curveFromCurve
    , curveToCurveoidBS
    )
import Codec.Encryption.OpenPGP.Policy (ecdhKdfHashDigest)
import Codec.Encryption.OpenPGP.Types

normalizeMontgomeryPublic
    :: forall n
     . KnownNat n
    => String
    -> B.ByteString
    -> Either String (FixedWidthBytes n)
normalizeMontgomeryPublic :: forall (n :: Nat).
KnownNat n =>
String -> ByteString -> Either String (FixedWidthBytes n)
normalizeMontgomeryPublic String
label ByteString
bs
    | ByteString -> Int
B.length ByteString
bs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
targetLen =
        String
-> Maybe (FixedWidthBytes n) -> Either String (FixedWidthBytes n)
forall a b. a -> Maybe b -> Either a b
note (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)) (forall (n :: Nat).
KnownNat n =>
ByteString -> Maybe (FixedWidthBytes n)
bsToFixedWidth @n ByteString
bs)
    | ByteString -> Int
B.length ByteString
bs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
targetLen =
        String
-> Maybe (FixedWidthBytes n) -> Either String (FixedWidthBytes n)
forall a b. a -> Maybe b -> Either a b
note String
"leftPadTo: input exceeds target" (forall (n :: Nat).
KnownNat n =>
ByteString -> Maybe (FixedWidthBytes n)
bsToFixedWidth @n 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 =
        String
-> Maybe (FixedWidthBytes n) -> Either String (FixedWidthBytes n)
forall a b. a -> Maybe b -> Either a b
note
            (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))
            (forall (n :: Nat).
KnownNat n =>
ByteString -> Maybe (FixedWidthBytes n)
bsToFixedWidth @n (HasCallStack => ByteString -> ByteString
ByteString -> ByteString
B.tail ByteString
bs))
    | Bool
otherwise = String -> Either String (FixedWidthBytes n)
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))
  where
    targetLen :: Int
targetLen = forall (n :: Nat). KnownNat n => Int
byteWidth @n

buildECDHKDFParam
    :: SomePKPayload
    -> PubKeyAlgorithm
    -> PKey
    -> HashAlgorithm
    -> SymmetricAlgorithm
    -> Either CipherError B.ByteString
buildECDHKDFParam :: SomePKPayload
-> PubKeyAlgorithm
-> PKey
-> HashAlgorithm
-> SymmetricAlgorithm
-> Either CipherError 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
<> Fingerprint -> ByteString
unFingerprint (SomePKPayload -> Fingerprint
fingerprint SomePKPayload
recipientPKP)
    )
        (ByteString -> ByteString)
-> Either CipherError ByteString -> Either CipherError ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either CipherError ByteString
encodedCurveOid
  where
    encodedCurveOid :: Either CipherError 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 CipherError ByteString -> Either CipherError ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either CipherError ByteString
curveOid
    curveOid :: Either CipherError ByteString
curveOid =
        case PKey
recipientECDHPub of
            ECDSAPubKey (ECDSA_PublicKey (ECDSA.PublicKey Curve
curve PublicPoint
_)) ->
                (CurveConversionError -> CipherError)
-> Either CurveConversionError ByteString
-> Either CipherError ByteString
forall a b c. (a -> b) -> Either a c -> Either b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first
                    CurveConversionError -> CipherError
CipherCurveConversionFailed
                    (ECCCurve -> Either CurveConversionError ByteString
curveToCurveoidBS (Curve -> ECCCurve
curveFromCurve Curve
curve))
            EdDSAPubKey EdSigningCurve
EdSigningCurve25519 EdPoint
_ ->
                (CurveConversionError -> CipherError)
-> Either CurveConversionError ByteString
-> Either CipherError ByteString
forall a b c. (a -> b) -> Either a c -> Either b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first CurveConversionError -> CipherError
CipherCurveConversionFailed (ECCCurve -> Either CurveConversionError ByteString
curveToCurveoidBS ECCCurve
Curve25519)
            EdDSAPubKey EdSigningCurve
EdSigningCurve448 EdPoint
_ ->
                (CurveConversionError -> CipherError)
-> Either CurveConversionError ByteString
-> Either CipherError ByteString
forall a b c. (a -> b) -> Either a c -> Either b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first CurveConversionError -> CipherError
CipherCurveConversionFailed (ECCCurve -> Either CurveConversionError ByteString
curveToCurveoidBS ECCCurve
Curve448)
            PKey
_ -> CipherError -> Either CipherError ByteString
forall a b. a -> Either a b
Left CipherError
CipherInvalidECDHRecipient

deriveECDHKek
    :: HashAlgorithm
    -> SymmetricAlgorithm
    -> B.ByteString
    -> B.ByteString
    -> Either CipherError B.ByteString
deriveECDHKek :: HashAlgorithm
-> SymmetricAlgorithm
-> ByteString
-> ByteString
-> Either CipherError ByteString
deriveECDHKek HashAlgorithm
kdfHA SymmetricAlgorithm
kdfSA ByteString
sharedSecret ByteString
kdfParam = do
    digest <-
        HashAlgorithm -> ByteString -> Either CipherError 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 <- keySize kdfSA
    if B.length digest < kekLen
        then
            Left
                ( CipherKeyWrapInvalidInput
                    "ECDH KDF digest is shorter than required KEK length"
                )
        else Right (B.take kekLen digest)