{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}

module Simplex.Messaging.Crypto.SNTRUP761
  ( KEMHybridSecret (..),
    kcbDecrypt,
    kcbEncrypt,
    kemHybridSecret,
  ) where

import Crypto.Hash (Digest, SHA3_256, hash)
import Data.ByteArray (ScrubbedBytes)
import qualified Data.ByteArray as BA
import Data.ByteString (ByteString)
import Simplex.Messaging.Crypto
import qualified Simplex.Messaging.Crypto as C
import Simplex.Messaging.Crypto.SNTRUP761.Bindings

-- Hybrid shared secret for crypto_box is defined as SHA256(DHSecret || KEMSharedKey),
-- similar to https://datatracker.ietf.org/doc/draft-josefsson-ntruprime-hybrid/

newtype KEMHybridSecret = KEMHybridSecret ScrubbedBytes

-- | NaCl @crypto_box@ decrypt with a shared hybrid DH + KEM secret and 192-bit nonce.
kcbDecrypt :: KEMHybridSecret -> CbNonce -> ByteString -> Either CryptoError ByteString
kcbDecrypt :: KEMHybridSecret
-> CbNonce -> ByteString -> Either CryptoError ByteString
kcbDecrypt (KEMHybridSecret ScrubbedBytes
k) = ScrubbedBytes
-> CbNonce -> ByteString -> Either CryptoError ByteString
forall key.
ByteArrayAccess key =>
key -> CbNonce -> ByteString -> Either CryptoError ByteString
sbDecrypt_ ScrubbedBytes
k

-- | NaCl @crypto_box@ encrypt with a shared hybrid DH + KEM secret and 192-bit nonce.
kcbEncrypt :: KEMHybridSecret -> CbNonce -> ByteString -> Int -> Either CryptoError ByteString
kcbEncrypt :: KEMHybridSecret
-> CbNonce -> ByteString -> Int -> Either CryptoError ByteString
kcbEncrypt (KEMHybridSecret ScrubbedBytes
k) = ScrubbedBytes
-> CbNonce -> ByteString -> Int -> Either CryptoError ByteString
forall key.
ByteArrayAccess key =>
key
-> CbNonce -> ByteString -> Int -> Either CryptoError ByteString
sbEncrypt_ ScrubbedBytes
k

kemHybridSecret :: PublicKeyX25519 -> PrivateKeyX25519 -> KEMSharedKey -> KEMHybridSecret
kemHybridSecret :: PublicKeyX25519
-> PrivateKeyX25519 -> KEMSharedKey -> KEMHybridSecret
kemHybridSecret PublicKeyX25519
k PrivateKeyX25519
pk (KEMSharedKey ScrubbedBytes
kem) =
  let DhSecretX25519 DhSecret
dh = PublicKeyX25519 -> PrivateKeyX25519 -> DhSecret 'X25519
forall (a :: Algorithm).
DhAlgorithm a =>
PublicKey a -> PrivateKey a -> DhSecret a
C.dh' PublicKeyX25519
k PrivateKeyX25519
pk
   in ScrubbedBytes -> KEMHybridSecret
KEMHybridSecret (ScrubbedBytes -> KEMHybridSecret)
-> ScrubbedBytes -> KEMHybridSecret
forall a b. (a -> b) -> a -> b
$ Digest SHA3_256 -> ScrubbedBytes
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
bin -> bout
BA.convert (ScrubbedBytes -> Digest SHA3_256
forall ba a.
(ByteArrayAccess ba, HashAlgorithm a) =>
ba -> Digest a
hash (ScrubbedBytes -> Digest SHA3_256)
-> ScrubbedBytes -> Digest SHA3_256
forall a b. (a -> b) -> a -> b
$ DhSecret -> ScrubbedBytes
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
bin -> bout
BA.convert DhSecret
dh ScrubbedBytes -> ScrubbedBytes -> ScrubbedBytes
forall a. Semigroup a => a -> a -> a
<> ScrubbedBytes
kem :: Digest SHA3_256)