module Network.TLS.Crypto.DH
    (
    -- * DH types
      DHParams
    , DHPublic
    , DHPrivate
    , DHKey

    -- * DH methods
    , dhPublic
    , dhPrivate
    , dhParams
    , dhParamsGetP
    , dhParamsGetG
    , dhParamsGetBits
    , dhGenerateKeyPair
    , dhGetShared
    , dhValid
    , dhUnwrap
    , dhUnwrapPublic
    ) where

import qualified Crypto.PubKey.DH as DH
import           Crypto.Number.Basic (numBits)
import qualified Data.ByteArray as B
import           Network.TLS.RNG

type DHPublic   = DH.PublicNumber
type DHPrivate  = DH.PrivateNumber
type DHParams   = DH.Params
type DHKey      = DH.SharedKey

dhPublic :: Integer -> DHPublic
dhPublic :: Integer -> DHPublic
dhPublic = Integer -> DHPublic
DH.PublicNumber

dhPrivate :: Integer -> DHPrivate
dhPrivate :: Integer -> DHPrivate
dhPrivate = Integer -> DHPrivate
DH.PrivateNumber

dhParams :: Integer -> Integer -> DHParams
dhParams :: Integer -> Integer -> DHParams
dhParams Integer
p Integer
g = Integer -> Integer -> Int -> DHParams
DH.Params Integer
p Integer
g (Integer -> Int
numBits Integer
p)

dhGenerateKeyPair :: MonadRandom r => DHParams -> r (DHPrivate, DHPublic)
dhGenerateKeyPair :: forall (r :: * -> *).
MonadRandom r =>
DHParams -> r (DHPrivate, DHPublic)
dhGenerateKeyPair DHParams
params = do
    DHPrivate
priv <- forall (m :: * -> *). MonadRandom m => DHParams -> m DHPrivate
DH.generatePrivate DHParams
params
    let pub :: DHPublic
pub        = DHParams -> DHPrivate -> DHPublic
DH.calculatePublic DHParams
params DHPrivate
priv
    forall (m :: * -> *) a. Monad m => a -> m a
return (DHPrivate
priv, DHPublic
pub)

dhGetShared :: DHParams -> DHPrivate -> DHPublic -> DHKey
dhGetShared :: DHParams -> DHPrivate -> DHPublic -> DHKey
dhGetShared DHParams
params DHPrivate
priv DHPublic
pub =
    DHKey -> DHKey
stripLeadingZeros (DHParams -> DHPrivate -> DHPublic -> DHKey
DH.getShared DHParams
params DHPrivate
priv DHPublic
pub)
  where
    -- strips leading zeros from the result of DH.getShared, as required
    -- for DH(E) premaster secret in SSL/TLS before version 1.3.
    stripLeadingZeros :: DHKey -> DHKey
stripLeadingZeros (DH.SharedKey ScrubbedBytes
sb) = ScrubbedBytes -> DHKey
DH.SharedKey (forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall bs. ByteArray bs => (Word8 -> Bool) -> bs -> (bs, bs)
B.span (forall a. Eq a => a -> a -> Bool
== Word8
0) ScrubbedBytes
sb)

-- Check that group element in not in the 2-element subgroup { 1, p - 1 }.
-- See RFC 7919 section 3 and NIST SP 56A rev 2 section 5.6.2.3.1.
-- This verification is enough when using a safe prime.
dhValid :: DHParams -> Integer -> Bool
dhValid :: DHParams -> Integer -> Bool
dhValid (DH.Params Integer
p Integer
_ Int
_) Integer
y = Integer
1 forall a. Ord a => a -> a -> Bool
< Integer
y Bool -> Bool -> Bool
&& Integer
y forall a. Ord a => a -> a -> Bool
< Integer
p forall a. Num a => a -> a -> a
- Integer
1

dhUnwrap :: DHParams -> DHPublic -> [Integer]
dhUnwrap :: DHParams -> DHPublic -> [Integer]
dhUnwrap (DH.Params Integer
p Integer
g Int
_) (DH.PublicNumber Integer
y) = [Integer
p,Integer
g,Integer
y]

dhParamsGetP :: DHParams -> Integer
dhParamsGetP :: DHParams -> Integer
dhParamsGetP (DH.Params Integer
p Integer
_ Int
_) = Integer
p

dhParamsGetG :: DHParams -> Integer
dhParamsGetG :: DHParams -> Integer
dhParamsGetG (DH.Params Integer
_ Integer
g Int
_) = Integer
g

dhParamsGetBits :: DHParams -> Int
dhParamsGetBits :: DHParams -> Int
dhParamsGetBits (DH.Params Integer
_ Integer
_ Int
b) = Int
b

dhUnwrapPublic :: DHPublic -> Integer
dhUnwrapPublic :: DHPublic -> Integer
dhUnwrapPublic (DH.PublicNumber Integer
y) = Integer
y