module Network.TLS.Sending (
    encodePacket12,
    encodePacket13,
    updateHandshake12,
    updateHandshake13,
) where

import Control.Concurrent.MVar
import Control.Monad.State.Strict
import qualified Data.ByteString as B
import Data.IORef

import Network.TLS.Cipher
import Network.TLS.Context.Internal
import Network.TLS.Handshake.Random
import Network.TLS.Handshake.State
import Network.TLS.Handshake.State13
import Network.TLS.Imports
import Network.TLS.Packet
import Network.TLS.Packet13
import Network.TLS.Parameters
import Network.TLS.Record
import Network.TLS.State
import Network.TLS.Struct
import Network.TLS.Struct13
import Network.TLS.Types (Role (..))
import Network.TLS.Util

-- | encodePacket transform a packet into marshalled data related to current state
-- and updating state on the go
encodePacket12
    :: Monoid bytes
    => Context
    -> RecordLayer bytes
    -> Packet
    -> IO (Either TLSError bytes)
encodePacket12 ctx recordLayer pkt = do
    (ver, _) <- decideRecordVersion ctx
    let pt = packetType pkt
        mkRecord bs = Record pt ver (fragmentPlaintext bs)
        len = ctxFragmentSize ctx
    records <- map mkRecord <$> packetToFragments12 ctx len pkt
    bs <- fmap mconcat <$> forEitherM records (recordEncode recordLayer ctx)
    when (pkt == ChangeCipherSpec) $ switchTxEncryption ctx
    return bs

-- Decompose handshake packets into fragments of the specified length.  AppData
-- packets are not fragmented here but by callers of sendPacket, so that the
-- empty-packet countermeasure may be applied to each fragment independently.
packetToFragments12 :: Context -> Maybe Int -> Packet -> IO [ByteString]
packetToFragments12 ctx len (Handshake hss) =
    getChunks len . B.concat <$> mapM (updateHandshake12 ctx) hss
packetToFragments12 _ _ (Alert a) = return [encodeAlerts a]
packetToFragments12 _ _ ChangeCipherSpec = return [encodeChangeCipherSpec]
packetToFragments12 _ _ (AppData x) = return [x]

switchTxEncryption :: Context -> IO ()
switchTxEncryption ctx = do
    tx <- usingHState ctx (fromJust <$> gets hstPendingTxState)
    (ver, role) <- usingState_ ctx $ do
        v <- getVersion
        r <- getRole
        return (v, r)
    liftIO $ modifyMVar_ (ctxTxRecordState ctx) (\_ -> return tx)
    -- set empty packet counter measure if condition are met
    when
        ( ver <= TLS10
            && role == ClientRole
            && isCBC tx
            && supportedEmptyPacket (ctxSupported ctx)
        )
        $ liftIO
        $ writeIORef (ctxNeedEmptyPacket ctx) True
  where
    isCBC tx = maybe False (\c -> bulkBlockSize (cipherBulk c) > 0) (stCipher tx)

updateHandshake12 :: Context -> Handshake -> IO ByteString
updateHandshake12 ctx hs = do
    usingHState ctx $ do
        when (certVerifyHandshakeMaterial hs) $ addHandshakeMessage encoded
        when (finishedHandshakeMaterial hs) $ updateHandshakeDigest encoded
    return encoded
  where
    encoded = encodeHandshake hs

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

encodePacket13
    :: Monoid bytes
    => Context
    -> RecordLayer bytes
    -> Packet13
    -> IO (Either TLSError bytes)
encodePacket13 ctx recordLayer pkt = do
    let pt = contentType pkt
        mkRecord bs = Record pt TLS12 (fragmentPlaintext bs)
        len = ctxFragmentSize ctx
    records <- map mkRecord <$> packetToFragments13 ctx len pkt
    fmap mconcat <$> forEitherM records (recordEncode13 recordLayer ctx)

packetToFragments13 :: Context -> Maybe Int -> Packet13 -> IO [ByteString]
packetToFragments13 ctx len (Handshake13 hss) =
    getChunks len . B.concat <$> mapM (updateHandshake13 ctx) hss
packetToFragments13 _ _ (Alert13 a) = return [encodeAlerts a]
packetToFragments13 _ _ (AppData13 x) = return [x]
packetToFragments13 _ _ ChangeCipherSpec13 = return [encodeChangeCipherSpec]

updateHandshake13 :: Context -> Handshake13 -> IO ByteString
updateHandshake13 ctx hs
    | isIgnored hs = return encoded
    | otherwise = usingHState ctx $ do
        when (isHRR hs) wrapAsMessageHash13
        updateHandshakeDigest encoded
        addHandshakeMessage encoded
        return encoded
  where
    encoded = encodeHandshake13 hs

    isHRR (ServerHello13 srand _ _ _) = isHelloRetryRequest srand
    isHRR _ = False

    isIgnored NewSessionTicket13{} = True
    isIgnored KeyUpdate13{} = True
    isIgnored _ = False