-- 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 :: 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
Ed25519 EdPoint
_ ->
          ECCCurve -> Either String ByteString
curveToCurveoidBS ECCCurve
Curve25519
        EdDSAPubKey EdSigningCurve
Ed448 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)