    {-# LANGUAGE FlexibleInstances #-}

    -- | Functions for RSA operations
    module Network.TLS.Handshake.Key (
        encryptRSA,
        signPrivate,
        decryptRSA,
        verifyPublic,
        generateDHE,
        generateECDHE,
        generateECDHEShared,
        generateFFDHE,
        generateFFDHEShared,
        versionCompatible,
        isDigitalSignaturePair,
        checkDigitalSignatureKey,
        getLocalPublicKey,
        satisfiesEcPredicate,
        logKey,
    ) where

    import Control.Monad.State.Strict

    import qualified Data.ByteString as B

    import Network.TLS.Context.Internal
    import Network.TLS.Crypto
    import Network.TLS.Handshake.State
    import Network.TLS.Imports
    import Network.TLS.State (withRNG)
    import Network.TLS.Struct
    import Network.TLS.Types
    import Network.TLS.X509

    {- if the RSA encryption fails we just return an empty bytestring, and let the protocol
     - fail by itself; however it would be probably better to just report it since it's an internal problem.
     -}
    encryptRSA :: Context -> ByteString -> IO ByteString
    encryptRSA :: Context -> ByteString -> IO ByteString
    encryptRSA Context
    ctx ByteString
    content = do
        publicKey <- Context -> HandshakeM PubKey -> IO PubKey
    forall (m :: * -> *) a. MonadIO m => Context -> HandshakeM a -> m a
    usingHState Context
    ctx HandshakeM PubKey
    getRemotePublicKey
        usingState_ ctx $ do
            v <- withRNG $ kxEncrypt publicKey content
            case v of
                Left KxError
    err -> [Char] -> TLSSt ByteString
    forall a. HasCallStack => [Char] -> a
    error ([Char]
    "rsa encrypt failed: " [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ KxError -> [Char]
    forall a. Show a => a -> [Char]
    show KxError
    err)
                Right ByteString
    econtent -> ByteString -> TLSSt ByteString
    forall a. a -> TLSSt a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ByteString
    econtent

    signPrivate :: Context -> Role -> SignatureParams -> ByteString -> IO ByteString
    signPrivate :: Context -> Role -> SignatureParams -> ByteString -> IO ByteString
    signPrivate Context
    ctx Role
    _ SignatureParams
    params ByteString
    content = do
        (publicKey, privateKey) <- Context -> HandshakeM (PubKey, PrivKey) -> IO (PubKey, PrivKey)
    forall (m :: * -> *) a. MonadIO m => Context -> HandshakeM a -> m a
    usingHState Context
    ctx HandshakeM (PubKey, PrivKey)
    getLocalPublicPrivateKeys
        usingState_ ctx $ do
            r <- withRNG $ kxSign privateKey publicKey params content
            case r of
                Left KxError
    err -> [Char] -> TLSSt ByteString
    forall a. HasCallStack => [Char] -> a
    error ([Char]
    "sign failed: " [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ KxError -> [Char]
    forall a. Show a => a -> [Char]
    show KxError
    err)
                Right ByteString
    econtent -> ByteString -> TLSSt ByteString
    forall a. a -> TLSSt a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ByteString
    econtent

    decryptRSA :: Context -> ByteString -> IO (Either KxError ByteString)
    decryptRSA :: Context -> ByteString -> IO (Either KxError ByteString)
    decryptRSA Context
    ctx ByteString
    econtent = do
        (_, privateKey) <- Context -> HandshakeM (PubKey, PrivKey) -> IO (PubKey, PrivKey)
    forall (m :: * -> *) a. MonadIO m => Context -> HandshakeM a -> m a
    usingHState Context
    ctx HandshakeM (PubKey, PrivKey)
    getLocalPublicPrivateKeys
        usingState_ ctx $ do
            let cipher = Int -> ByteString -> ByteString
    B.drop Int
    2 ByteString
    econtent
            withRNG $ kxDecrypt privateKey cipher

    verifyPublic
        :: Context -> SignatureParams -> ByteString -> ByteString -> IO Bool
    verifyPublic :: Context -> SignatureParams -> ByteString -> ByteString -> IO Bool
    verifyPublic Context
    ctx SignatureParams
    params ByteString
    econtent ByteString
    sign = do
        publicKey <- Context -> HandshakeM PubKey -> IO PubKey
    forall (m :: * -> *) a. MonadIO m => Context -> HandshakeM a -> m a
    usingHState Context
    ctx HandshakeM PubKey
    getRemotePublicKey
        return $ kxVerify publicKey params econtent sign

    generateDHE :: Context -> DHParams -> IO (DHPrivate, DHPublic)
    generateDHE :: Context -> DHParams -> IO (DHPrivate, DHPublic)
    generateDHE Context
    ctx DHParams
    dhp = Context -> TLSSt (DHPrivate, DHPublic) -> IO (DHPrivate, DHPublic)
    forall a. Context -> TLSSt a -> IO a
    usingState_ Context
    ctx (TLSSt (DHPrivate, DHPublic) -> IO (DHPrivate, DHPublic))
    -> TLSSt (DHPrivate, DHPublic) -> IO (DHPrivate, DHPublic)
    forall a b. (a -> b) -> a -> b
    $ MonadPseudoRandom StateRNG (DHPrivate, DHPublic)
    -> TLSSt (DHPrivate, DHPublic)
    forall a. MonadPseudoRandom StateRNG a -> TLSSt a
    withRNG (MonadPseudoRandom StateRNG (DHPrivate, DHPublic)
     -> TLSSt (DHPrivate, DHPublic))
    -> MonadPseudoRandom StateRNG (DHPrivate, DHPublic)
    -> TLSSt (DHPrivate, DHPublic)
    forall a b. (a -> b) -> a -> b
    $ DHParams -> MonadPseudoRandom StateRNG (DHPrivate, DHPublic)
    forall (r :: * -> *).
    MonadRandom r =>
    DHParams -> r (DHPrivate, DHPublic)
    dhGenerateKeyPair DHParams
    dhp

    generateECDHE :: Context -> Group -> IO (GroupPrivate, GroupPublic)
    generateECDHE :: Context -> Group -> IO (GroupPrivate, GroupPublic)
    generateECDHE Context
    ctx Group
    grp = Context
    -> TLSSt (GroupPrivate, GroupPublic)
    -> IO (GroupPrivate, GroupPublic)
    forall a. Context -> TLSSt a -> IO a
    usingState_ Context
    ctx (TLSSt (GroupPrivate, GroupPublic)
     -> IO (GroupPrivate, GroupPublic))
    -> TLSSt (GroupPrivate, GroupPublic)
    -> IO (GroupPrivate, GroupPublic)
    forall a b. (a -> b) -> a -> b
    $ MonadPseudoRandom StateRNG (GroupPrivate, GroupPublic)
    -> TLSSt (GroupPrivate, GroupPublic)
    forall a. MonadPseudoRandom StateRNG a -> TLSSt a
    withRNG (MonadPseudoRandom StateRNG (GroupPrivate, GroupPublic)
     -> TLSSt (GroupPrivate, GroupPublic))
    -> MonadPseudoRandom StateRNG (GroupPrivate, GroupPublic)
    -> TLSSt (GroupPrivate, GroupPublic)
    forall a b. (a -> b) -> a -> b
    $ Group -> MonadPseudoRandom StateRNG (GroupPrivate, GroupPublic)
    forall (r :: * -> *).
    MonadRandom r =>
    Group -> r (GroupPrivate, GroupPublic)
    groupGenerateKeyPair Group
    grp

    generateECDHEShared
        :: Context -> GroupPublic -> IO (Maybe (GroupPublic, GroupKey))
    generateECDHEShared :: Context -> GroupPublic -> IO (Maybe (GroupPublic, GroupKey))
    generateECDHEShared Context
    ctx GroupPublic
    pub = Context
    -> TLSSt (Maybe (GroupPublic, GroupKey))
    -> IO (Maybe (GroupPublic, GroupKey))
    forall a. Context -> TLSSt a -> IO a
    usingState_ Context
    ctx (TLSSt (Maybe (GroupPublic, GroupKey))
     -> IO (Maybe (GroupPublic, GroupKey)))
    -> TLSSt (Maybe (GroupPublic, GroupKey))
    -> IO (Maybe (GroupPublic, GroupKey))
    forall a b. (a -> b) -> a -> b
    $ MonadPseudoRandom StateRNG (Maybe (GroupPublic, GroupKey))
    -> TLSSt (Maybe (GroupPublic, GroupKey))
    forall a. MonadPseudoRandom StateRNG a -> TLSSt a
    withRNG (MonadPseudoRandom StateRNG (Maybe (GroupPublic, GroupKey))
     -> TLSSt (Maybe (GroupPublic, GroupKey)))
    -> MonadPseudoRandom StateRNG (Maybe (GroupPublic, GroupKey))
    -> TLSSt (Maybe (GroupPublic, GroupKey))
    forall a b. (a -> b) -> a -> b
    $ GroupPublic
    -> MonadPseudoRandom StateRNG (Maybe (GroupPublic, GroupKey))
    forall (r :: * -> *).
    MonadRandom r =>
    GroupPublic -> r (Maybe (GroupPublic, GroupKey))
    groupGetPubShared GroupPublic
    pub

    generateFFDHE :: Context -> Group -> IO (DHParams, DHPrivate, DHPublic)
    generateFFDHE :: Context -> Group -> IO (DHParams, DHPrivate, DHPublic)
    generateFFDHE Context
    ctx Group
    grp = Context
    -> TLSSt (DHParams, DHPrivate, DHPublic)
    -> IO (DHParams, DHPrivate, DHPublic)
    forall a. Context -> TLSSt a -> IO a
    usingState_ Context
    ctx (TLSSt (DHParams, DHPrivate, DHPublic)
     -> IO (DHParams, DHPrivate, DHPublic))
    -> TLSSt (DHParams, DHPrivate, DHPublic)
    -> IO (DHParams, DHPrivate, DHPublic)
    forall a b. (a -> b) -> a -> b
    $ MonadPseudoRandom StateRNG (DHParams, DHPrivate, DHPublic)
    -> TLSSt (DHParams, DHPrivate, DHPublic)
    forall a. MonadPseudoRandom StateRNG a -> TLSSt a
    withRNG (MonadPseudoRandom StateRNG (DHParams, DHPrivate, DHPublic)
     -> TLSSt (DHParams, DHPrivate, DHPublic))
    -> MonadPseudoRandom StateRNG (DHParams, DHPrivate, DHPublic)
    -> TLSSt (DHParams, DHPrivate, DHPublic)
    forall a b. (a -> b) -> a -> b
    $ Group -> MonadPseudoRandom StateRNG (DHParams, DHPrivate, DHPublic)
    forall (r :: * -> *).
    MonadRandom r =>
    Group -> r (DHParams, DHPrivate, DHPublic)
    dhGroupGenerateKeyPair Group
    grp

    generateFFDHEShared
        :: Context -> Group -> DHPublic -> IO (Maybe (DHPublic, DHKey))
    generateFFDHEShared :: Context -> Group -> DHPublic -> IO (Maybe (DHPublic, DHKey))
    generateFFDHEShared Context
    ctx Group
    grp DHPublic
    pub = Context
    -> TLSSt (Maybe (DHPublic, DHKey)) -> IO (Maybe (DHPublic, DHKey))
    forall a. Context -> TLSSt a -> IO a
    usingState_ Context
    ctx (TLSSt (Maybe (DHPublic, DHKey)) -> IO (Maybe (DHPublic, DHKey)))
    -> TLSSt (Maybe (DHPublic, DHKey)) -> IO (Maybe (DHPublic, DHKey))
    forall a b. (a -> b) -> a -> b
    $ MonadPseudoRandom StateRNG (Maybe (DHPublic, DHKey))
    -> TLSSt (Maybe (DHPublic, DHKey))
    forall a. MonadPseudoRandom StateRNG a -> TLSSt a
    withRNG (MonadPseudoRandom StateRNG (Maybe (DHPublic, DHKey))
     -> TLSSt (Maybe (DHPublic, DHKey)))
    -> MonadPseudoRandom StateRNG (Maybe (DHPublic, DHKey))
    -> TLSSt (Maybe (DHPublic, DHKey))
    forall a b. (a -> b) -> a -> b
    $ Group
    -> DHPublic -> MonadPseudoRandom StateRNG (Maybe (DHPublic, DHKey))
    forall (r :: * -> *).
    MonadRandom r =>
    Group -> DHPublic -> r (Maybe (DHPublic, DHKey))
    dhGroupGetPubShared Group
    grp DHPublic
    pub

    {- FOURMOLU_DISABLE -}
    isDigitalSignatureKey :: PubKey -> Bool
    isDigitalSignatureKey :: PubKey -> Bool
    isDigitalSignatureKey (PubKeyRSA PublicKey
    _)     = Bool
    True
    isDigitalSignatureKey (PubKeyDSA PublicKey
    _)     = Bool
    True
    isDigitalSignatureKey (PubKeyEC PubKeyEC
    _)      = Bool
    True
    isDigitalSignatureKey (PubKeyEd25519 PublicKey
    _) = Bool
    True
    isDigitalSignatureKey (PubKeyEd448 PublicKey
    _)   = Bool
    True
    isDigitalSignatureKey PubKey
    _                 = Bool
    False

    versionCompatible :: PubKey -> Version -> Bool
    versionCompatible :: PubKey -> Version -> Bool
    versionCompatible (PubKeyRSA PublicKey
    _) Version
    _     = Bool
    True
    versionCompatible (PubKeyDSA PublicKey
    _) Version
    v     = Version
    v Version -> Version -> Bool
    forall a. Ord a => a -> a -> Bool
    <= Version
    TLS12
    versionCompatible (PubKeyEC PubKeyEC
    _) Version
    v      = Version
    v Version -> Version -> Bool
    forall a. Ord a => a -> a -> Bool
    >= Version
    TLS10
    versionCompatible (PubKeyEd25519 PublicKey
    _) Version
    v = Version
    v Version -> Version -> Bool
    forall a. Ord a => a -> a -> Bool
    >= Version
    TLS12
    versionCompatible (PubKeyEd448 PublicKey
    _) Version
    v   = Version
    v Version -> Version -> Bool
    forall a. Ord a => a -> a -> Bool
    >= Version
    TLS12
    versionCompatible PubKey
    _ Version
    _                 = Bool
    False
    {- FOURMOLU_ENABLE -}

    -- | Test whether the argument is a public key supported for signature at the
    -- specified TLS version.  This also accepts a key for RSA encryption.  This
    -- test is performed by clients or servers before verifying a remote
    -- Certificate Verify.
    checkDigitalSignatureKey :: MonadIO m => Version -> PubKey -> m ()
    checkDigitalSignatureKey :: forall (m :: * -> *). MonadIO m => Version -> PubKey -> m ()
    checkDigitalSignatureKey Version
    usedVersion PubKey
    key = do
        Bool -> m () -> m ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless (PubKey -> Bool
    isDigitalSignatureKey PubKey
    key) (m () -> m ()) -> m () -> m ()
    forall a b. (a -> b) -> a -> b
    $
            TLSError -> m ()
    forall (m :: * -> *) a. MonadIO m => TLSError -> m a
    throwCore (TLSError -> m ()) -> TLSError -> m ()
    forall a b. (a -> b) -> a -> b
    $
                [Char] -> AlertDescription -> TLSError
    Error_Protocol [Char]
    "unsupported remote public key type" AlertDescription
    HandshakeFailure
        Bool -> m () -> m ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless (PubKey
    key PubKey -> Version -> Bool
    `versionCompatible` Version
    usedVersion) (m () -> m ()) -> m () -> m ()
    forall a b. (a -> b) -> a -> b
    $
            TLSError -> m ()
    forall (m :: * -> *) a. MonadIO m => TLSError -> m a
    throwCore (TLSError -> m ()) -> TLSError -> m ()
    forall a b. (a -> b) -> a -> b
    $
                [Char] -> AlertDescription -> TLSError
    Error_Protocol
                    (Version -> [Char]
    forall a. Show a => a -> [Char]
    show Version
    usedVersion [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    " has no support for " [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ PubKey -> [Char]
    pubkeyType PubKey
    key)
                    AlertDescription
    IllegalParameter

    -- | Test whether the argument is matching key pair supported for signature.
    -- This also accepts material for RSA encryption.  This test is performed by
    -- servers or clients before using a credential from the local configuration.
    isDigitalSignaturePair :: (PubKey, PrivKey) -> Bool
    isDigitalSignaturePair :: (PubKey, PrivKey) -> Bool
    isDigitalSignaturePair (PubKey, PrivKey)
    keyPair =
        case (PubKey, PrivKey)
    keyPair of
            (PubKeyRSA PublicKey
    _, PrivKeyRSA PrivateKey
    _) -> Bool
    True
            (PubKeyDSA PublicKey
    _, PrivKeyDSA PrivateKey
    _) -> Bool
    True
            (PubKeyEC PubKeyEC
    _, PrivKeyEC PrivKeyEC
    k) -> PrivKeyEC -> Bool
    kxSupportedPrivKeyEC PrivKeyEC
    k
            (PubKeyEd25519 PublicKey
    _, PrivKeyEd25519 SecretKey
    _) -> Bool
    True
            (PubKeyEd448 PublicKey
    _, PrivKeyEd448 SecretKey
    _) -> Bool
    True
            (PubKey, PrivKey)
    _ -> Bool
    False

    getLocalPublicKey :: MonadIO m => Context -> m PubKey
    getLocalPublicKey :: forall (m :: * -> *). MonadIO m => Context -> m PubKey
    getLocalPublicKey Context
    ctx =
        Context -> HandshakeM PubKey -> m PubKey
    forall (m :: * -> *) a. MonadIO m => Context -> HandshakeM a -> m a
    usingHState Context
    ctx ((PubKey, PrivKey) -> PubKey
    forall a b. (a, b) -> a
    fst ((PubKey, PrivKey) -> PubKey)
    -> HandshakeM (PubKey, PrivKey) -> HandshakeM PubKey
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> HandshakeM (PubKey, PrivKey)
    getLocalPublicPrivateKeys)

    -- | Test whether the public key satisfies a predicate about the elliptic curve.
    -- When the public key is not suitable for ECDSA, like RSA for instance, the
    -- predicate is not used and the result is 'True'.
    satisfiesEcPredicate :: (Group -> Bool) -> PubKey -> Bool
    satisfiesEcPredicate :: (Group -> Bool) -> PubKey -> Bool
    satisfiesEcPredicate Group -> Bool
    p (PubKeyEC PubKeyEC
    ecPub) =
        Bool -> (Group -> Bool) -> Maybe Group -> Bool
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe Bool
    False Group -> Bool
    p (Maybe Group -> Bool) -> Maybe Group -> Bool
    forall a b. (a -> b) -> a -> b
    $ PubKeyEC -> Maybe Group
    findEllipticCurveGroup PubKeyEC
    ecPub
    satisfiesEcPredicate Group -> Bool
    _ PubKey
    _ = Bool
    True

    ----------------------------------------------------------------

    class LogLabel a where
        labelAndKey :: a -> (String, ByteString)

    instance LogLabel MainSecret where
        labelAndKey :: MainSecret -> ([Char], ByteString)
    labelAndKey (MainSecret ByteString
    key) = ([Char]
    "CLIENT_RANDOM", ByteString
    key)

    instance LogLabel (ClientTrafficSecret EarlySecret) where
        labelAndKey :: ClientTrafficSecret EarlySecret -> ([Char], ByteString)
    labelAndKey (ClientTrafficSecret ByteString
    key) = ([Char]
    "CLIENT_EARLY_TRAFFIC_SECRET", ByteString
    key)

    instance LogLabel (ServerTrafficSecret HandshakeSecret) where
        labelAndKey :: ServerTrafficSecret HandshakeSecret -> ([Char], ByteString)
    labelAndKey (ServerTrafficSecret ByteString
    key) = ([Char]
    "SERVER_HANDSHAKE_TRAFFIC_SECRET", ByteString
    key)

    instance LogLabel (ClientTrafficSecret HandshakeSecret) where
        labelAndKey :: ClientTrafficSecret HandshakeSecret -> ([Char], ByteString)
    labelAndKey (ClientTrafficSecret ByteString
    key) = ([Char]
    "CLIENT_HANDSHAKE_TRAFFIC_SECRET", ByteString
    key)

    instance LogLabel (ServerTrafficSecret ApplicationSecret) where
        labelAndKey :: ServerTrafficSecret ApplicationSecret -> ([Char], ByteString)
    labelAndKey (ServerTrafficSecret ByteString
    key) = ([Char]
    "SERVER_TRAFFIC_SECRET_0", ByteString
    key)

    instance LogLabel (ClientTrafficSecret ApplicationSecret) where
        labelAndKey :: ClientTrafficSecret ApplicationSecret -> ([Char], ByteString)
    labelAndKey (ClientTrafficSecret ByteString
    key) = ([Char]
    "CLIENT_TRAFFIC_SECRET_0", ByteString
    key)

    -- NSS Key Log Format
    -- See https://developer.mozilla.org/en-US/docs/Mozilla/Projects/NSS/Key_Log_Format
    logKey :: LogLabel a => Context -> a -> IO ()
    logKey :: forall a. LogLabel a => Context -> a -> IO ()
    logKey Context
    ctx a
    logkey = do
        mhst <- Context -> IO (Maybe HandshakeState)
    forall (m :: * -> *).
    MonadIO m =>
    Context -> m (Maybe HandshakeState)
    getHState Context
    ctx
        case mhst of
            Maybe HandshakeState
    Nothing -> () -> IO ()
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ()
            Just HandshakeState
    hst -> do
                let cr :: ByteString
    cr = ClientRandom -> ByteString
    unClientRandom (ClientRandom -> ByteString) -> ClientRandom -> ByteString
    forall a b. (a -> b) -> a -> b
    $ HandshakeState -> ClientRandom
    hstClientRandom HandshakeState
    hst
                    ([Char]
    label, ByteString
    key) = a -> ([Char], ByteString)
    forall a. LogLabel a => a -> ([Char], ByteString)
    labelAndKey a
    logkey
                Context -> [Char] -> IO ()
    ctxKeyLogger Context
    ctx ([Char] -> IO ()) -> [Char] -> IO ()
    forall a b. (a -> b) -> a -> b
    $ [Char]
    label [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    " " [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ ByteString -> [Char]
    dump ByteString
    cr [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ [Char]
    " " [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ ByteString -> [Char]
    dump ByteString
    key
      where
        dump :: ByteString -> [Char]
    dump = [Char] -> [Char]
    forall a. HasCallStack => [a] -> [a]
    init ([Char] -> [Char])
    -> (ByteString -> [Char]) -> ByteString -> [Char]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Int -> [Char] -> [Char]
    forall a. Int -> [a] -> [a]
    drop Int
    1 ([Char] -> [Char])
    -> (ByteString -> [Char]) -> ByteString -> [Char]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ByteString -> [Char]
    showBytesHex