    {-# LANGUAGE OverloadedStrings #-}
    -- |
    -- Module      : Network.TLS.Packet
    -- License     : BSD-style
    -- Maintainer  : Vincent Hanquez <vincent@snarc.org>
    -- Stability   : experimental
    -- Portability : unknown
    --
    -- the Packet module contains everything necessary to serialize and deserialize things
    -- with only explicit parameters, no TLS state is involved here.
    --
    module Network.TLS.Packet
            (
            -- * params for encoding and decoding
              CurrentParams(..)
            -- * marshall functions for header messages
            , decodeHeader
            , decodeDeprecatedHeaderLength
            , decodeDeprecatedHeader
            , encodeHeader
            , encodeHeaderNoVer -- use for SSL3

            -- * marshall functions for alert messages
            , decodeAlert
            , decodeAlerts
            , encodeAlerts

            -- * marshall functions for handshake messages
            , decodeHandshakes
            , decodeHandshake
            , decodeDeprecatedHandshake
            , encodeHandshake
            , encodeHandshakes
            , encodeHandshakeHeader
            , encodeHandshakeContent

            -- * marshall functions for change cipher spec message
            , decodeChangeCipherSpec
            , encodeChangeCipherSpec

            , decodePreMasterSecret
            , encodePreMasterSecret

            -- * generate things for packet content
            , generateMasterSecret
            , generateKeyBlock
            , generateClientFinished
            , generateServerFinished

            , generateCertificateVerify_SSL
            ) where

    import Network.TLS.Struct
    import Network.TLS.Wire
    import Network.TLS.Cap
    import Data.Either (partitionEithers)
    import Data.Maybe (fromJust)
    import Data.Word
    import Data.Bits ((.|.))
    import Control.Applicative ((<$>))
    import Control.Monad
    import Data.Certificate.X509 (decodeCertificate, encodeCertificate, X509, encodeDN, decodeDN)
    import Network.TLS.Crypto
    import Network.TLS.MAC
    import Network.TLS.Cipher (CipherKeyExchangeType(..))
    import Data.ByteString (ByteString)
    import qualified Data.ByteString as B
    import qualified Data.ByteString.Char8 as BC
    import qualified Data.ByteString.Lazy as L

    import qualified Crypto.Hash.SHA1 as SHA1
    import qualified Crypto.Hash.MD5 as MD5

    data CurrentParams = CurrentParams
            { cParamsVersion     :: Version               -- ^ current protocol version
            , cParamsKeyXchgType :: CipherKeyExchangeType -- ^ current key exchange type
            , cParamsSupportNPN  :: Bool                  -- ^ support Next Protocol Negotiation extension
            } deriving (Show,Eq)

    {- marshall helpers -}
    getVersion :: Get Version
    getVersion = do
            major <- getWord8
            minor <- getWord8
            case verOfNum (major, minor) of
                    Nothing -> fail ("invalid version : " ++ show major ++ "," ++ show minor)
                    Just v  -> return v

    putVersion :: Version -> Put
    putVersion ver = putWord8 major >> putWord8 minor
            where (major, minor) = numericalVer ver

    getHeaderType :: Get ProtocolType
    getHeaderType = do
            ty <- getWord8
            case valToType ty of
                    Nothing -> fail ("invalid header type: " ++ show ty)
                    Just t  -> return t

    putHeaderType :: ProtocolType -> Put
    putHeaderType = putWord8 . valOfType

    getHandshakeType :: Get HandshakeType
    getHandshakeType = do
            ty <- getWord8
            case valToType ty of
                    Nothing -> fail ("invalid handshake type: " ++ show ty)
                    Just t  -> return t

    {-
     - decode and encode headers
     -}
    decodeHeader :: ByteString -> Either TLSError Header
    decodeHeader = runGetErr "header" $ liftM3 Header getHeaderType getVersion getWord16

    decodeDeprecatedHeaderLength :: ByteString -> Either TLSError Word16
    decodeDeprecatedHeaderLength = runGetErr "deprecatedheaderlength" $ subtract 0x8000 <$> getWord16

    decodeDeprecatedHeader :: Word16 -> ByteString -> Either TLSError Header
    decodeDeprecatedHeader size =
            runGetErr "deprecatedheader" $ do
                    1 <- getWord8
                    version <- getVersion
                    return $ Header ProtocolType_DeprecatedHandshake version size

    encodeHeader :: Header -> ByteString
    encodeHeader (Header pt ver len) = runPut (putHeaderType pt >> putVersion ver >> putWord16 len)
            {- FIXME check len <= 2^14 -}

    encodeHeaderNoVer :: Header -> ByteString
    encodeHeaderNoVer (Header pt _ len) = runPut (putHeaderType pt >> putWord16 len)
            {- FIXME check len <= 2^14 -}

    {-
     - decode and encode ALERT
     -}
    decodeAlert :: Get (AlertLevel, AlertDescription)
    decodeAlert = do
            al <- getWord8
            ad <- getWord8
            case (valToType al, valToType ad) of
                    (Just a, Just d) -> return (a, d)
                    (Nothing, _)     -> fail "cannot decode alert level"
                    (_, Nothing)     -> fail "cannot decode alert description"

    decodeAlerts :: ByteString -> Either TLSError [(AlertLevel, AlertDescription)]
    decodeAlerts = runGetErr "alerts" $ loop
            where loop = do
                    r <- remaining
                    if r == 0
                            then return []
                            else liftM2 (:) decodeAlert loop

    encodeAlerts :: [(AlertLevel, AlertDescription)] -> ByteString
    encodeAlerts l = runPut $ mapM_ encodeAlert l
            where encodeAlert (al, ad) = putWord8 (valOfType al) >> putWord8 (valOfType ad)

    {- decode and encode HANDSHAKE -}
    decodeHandshakeHeader :: Get (HandshakeType, Bytes)
    decodeHandshakeHeader = do
            ty      <- getHandshakeType
            content <- getOpaque24
            return (ty, content)

    decodeHandshakes :: ByteString -> Either TLSError [(HandshakeType, Bytes)]
    decodeHandshakes b = runGetErr "handshakes" getAll b where
            getAll = do
                    x <- decodeHandshakeHeader
                    empty <- isEmpty
                    if empty
                            then return [x]
                            else liftM ((:) x) getAll

    decodeHandshake :: CurrentParams -> HandshakeType -> ByteString -> Either TLSError Handshake
    decodeHandshake cp ty = runGetErr "handshake" $ case ty of
            HandshakeType_HelloRequest    -> decodeHelloRequest
            HandshakeType_ClientHello     -> decodeClientHello
            HandshakeType_ServerHello     -> decodeServerHello
            HandshakeType_Certificate     -> decodeCertificates
            HandshakeType_ServerKeyXchg   -> decodeServerKeyXchg cp
            HandshakeType_CertRequest     -> decodeCertRequest cp
            HandshakeType_ServerHelloDone -> decodeServerHelloDone
            HandshakeType_CertVerify      -> decodeCertVerify cp
            HandshakeType_ClientKeyXchg   -> decodeClientKeyXchg
            HandshakeType_Finished        -> decodeFinished
            HandshakeType_NPN             -> do
                    unless (cParamsSupportNPN cp) $ fail "unsupported handshake type"
                    decodeNextProtocolNegotiation

    decodeDeprecatedHandshake :: ByteString -> Either TLSError Handshake
    decodeDeprecatedHandshake b = runGetErr "deprecatedhandshake" getDeprecated b where
            getDeprecated = do
                    1 <- getWord8
                    ver <- getVersion
                    cipherSpecLen <- fromEnum <$> getWord16
                    sessionIdLen <- fromEnum <$> getWord16
                    challengeLen <- fromEnum <$> getWord16
                    ciphers <- getCipherSpec cipherSpecLen
                    session <- getSessionId sessionIdLen
                    random <- getChallenge challengeLen
                    let compressions = [0]
                    return $ ClientHello ver random session ciphers compressions [] (Just b)
            getCipherSpec len | len < 3 = return []
            getCipherSpec len = do
                    [c0,c1,c2] <- map fromEnum <$> replicateM 3 getWord8
                    ([ toEnum $ c1 * 0x100 + c2 | c0 == 0 ] ++) <$> getCipherSpec (len - 3)
            getSessionId 0 = return $ Session Nothing
            getSessionId len = Session . Just <$> getBytes len
            getChallenge len | 32 < len = getBytes (len - 32) >> getChallenge 32
            getChallenge len = ClientRandom . B.append (B.replicate (32 - len) 0) <$> getBytes len

    decodeHelloRequest :: Get Handshake
    decodeHelloRequest = return HelloRequest

    decodeClientHello :: Get Handshake
    decodeClientHello = do
            ver          <- getVersion
            random       <- getClientRandom32
            session      <- getSession
            ciphers      <- getWords16
            compressions <- getWords8
            r            <- remaining
            exts <- if hasHelloExtensions ver && r > 0
                    then fmap fromIntegral getWord16 >>= getExtensions
                    else return []
            return $ ClientHello ver random session ciphers compressions exts Nothing

    decodeServerHello :: Get Handshake
    decodeServerHello = do
            ver           <- getVersion
            random        <- getServerRandom32
            session       <- getSession
            cipherid      <- getWord16
            compressionid <- getWord8
            r             <- remaining
            exts <- if hasHelloExtensions ver && r > 0
                    then fmap fromIntegral getWord16 >>= getExtensions
                    else return []
            return $ ServerHello ver random session cipherid compressionid exts

    decodeServerHelloDone :: Get Handshake
    decodeServerHelloDone = return ServerHelloDone

    decodeCertificates :: Get Handshake
    decodeCertificates = do
        certsRaw <- getWord24 >>= \len -> getList (fromIntegral len) getCertRaw
        let (badCerts, certs) = partitionEithers $ map (decodeCertificate . L.fromChunks . (:[])) certsRaw
        if not $ null badCerts
            then fail ("error certificate parsing: " ++ show badCerts)
            else return $ Certificates certs
        where getCertRaw = getOpaque24 >>= \cert -> return (3 + B.length cert, cert)

    decodeFinished :: Get Handshake
    decodeFinished = Finished <$> (remaining >>= getBytes)

    decodeNextProtocolNegotiation :: Get Handshake
    decodeNextProtocolNegotiation = do
            opaque <- getOpaque8
            _      <- getOpaque8 -- ignore padding
            return $ HsNextProtocolNegotiation opaque

    getSignatureHashAlgorithm :: Get HashAndSignatureAlgorithm
    getSignatureHashAlgorithm = do
            h <- fromJust . valToType <$> getWord8
            s <- fromJust . valToType <$> getWord8
            return (h,s)

    decodeCertRequest :: CurrentParams -> Get Handshake
    decodeCertRequest cp = do
            certTypes <- map (fromJust . valToType . fromIntegral) <$> getWords8

            sigHashAlgs <- if cParamsVersion cp >= TLS12
                               then Just <$> (getWord16 >>= getSignatureHashAlgorithms)
                               else return Nothing
            dNameLen <- getWord16
            -- FIXME: Decide whether to remove this check completely or to make it an option.
            -- when (cParamsVersion cp < TLS12 && dNameLen < 3) $ fail "certrequest distinguishname not of the correct size"
            dNames <- getList (fromIntegral dNameLen) getDName
            return $ CertRequest certTypes sigHashAlgs dNames
        where
            getSignatureHashAlgorithms len = getList (fromIntegral len) (getSignatureHashAlgorithm >>= \sh -> return (2, sh))
            getDName = do
                dName <- getOpaque16
                when (B.length dName == 0) $ fail "certrequest: invalid DN length"
                dn <- decodeDName dName
                return (2 + B.length dName, dn)

            decodeDName d = case decodeDN (L.fromChunks [d]) of
                                Left err -> fail ("certrequest: " ++ show err)
                                Right s  -> return s

    decodeCertVerify :: CurrentParams -> Get Handshake
    decodeCertVerify cp = do
            mbHashSig <- if cParamsVersion cp >= TLS12
                         then Just <$> getSignatureHashAlgorithm
                         else return Nothing
            bs <- getOpaque16
            return $ CertVerify mbHashSig (CertVerifyData bs)

    decodeClientKeyXchg :: Get Handshake
    decodeClientKeyXchg = ClientKeyXchg <$> (remaining >>= getBytes)

    os2ip :: ByteString -> Integer
    os2ip = B.foldl' (\a b -> (256 * a) .|. (fromIntegral b)) 0

    decodeServerKeyXchg_DH :: Get ServerDHParams
    decodeServerKeyXchg_DH = do
            p <- getOpaque16
            g <- getOpaque16
            y <- getOpaque16
            return $ ServerDHParams { dh_p = os2ip p, dh_g = os2ip g, dh_Ys = os2ip y }

    decodeServerKeyXchg_RSA :: Get ServerRSAParams
    decodeServerKeyXchg_RSA = do
            modulus <- getOpaque16
            expo    <- getOpaque16
            return $ ServerRSAParams { rsa_modulus = os2ip modulus, rsa_exponent = os2ip expo }

    decodeServerKeyXchg :: CurrentParams -> Get Handshake
    decodeServerKeyXchg cp = ServerKeyXchg <$> case cParamsKeyXchgType cp of
            CipherKeyExchange_RSA     -> SKX_RSA . Just <$> decodeServerKeyXchg_RSA
            CipherKeyExchange_DH_Anon -> SKX_DH_Anon <$> decodeServerKeyXchg_DH
            CipherKeyExchange_DHE_RSA -> do
                    dhparams <- decodeServerKeyXchg_DH
                    signature <- getOpaque16
                    return $ SKX_DHE_RSA dhparams (B.unpack signature)
            CipherKeyExchange_DHE_DSS -> do
                    dhparams  <- decodeServerKeyXchg_DH
                    signature <- getOpaque16
                    return $ SKX_DHE_DSS dhparams (B.unpack signature)
            _ -> do
                    bs <- remaining >>= getBytes
                    return $ SKX_Unknown bs

    encodeHandshake :: Handshake -> ByteString
    encodeHandshake o =
            let content = runPut $ encodeHandshakeContent o in
            let len = fromIntegral $ B.length content in
            let header = case o of
                            ClientHello _ _ _ _ _ _ (Just _) -> "" -- SSLv2 ClientHello message
                            _ -> runPut $ encodeHandshakeHeader (typeOfHandshake o) len in
            B.concat [ header, content ]

    encodeHandshakes :: [Handshake] -> ByteString
    encodeHandshakes hss = B.concat $ map encodeHandshake hss

    encodeHandshakeHeader :: HandshakeType -> Int -> Put
    encodeHandshakeHeader ty len = putWord8 (valOfType ty) >> putWord24 len

    encodeHandshakeContent :: Handshake -> Put

    encodeHandshakeContent (ClientHello _ _ _ _ _ _ (Just deprecated)) = do
            putBytes deprecated
    encodeHandshakeContent (ClientHello version random session cipherIDs compressionIDs exts Nothing) = do
            putVersion version
            putClientRandom32 random
            putSession session
            putWords16 cipherIDs
            putWords8 compressionIDs
            putExtensions exts
            return ()

    encodeHandshakeContent (ServerHello version random session cipherID compressionID exts) =
            putVersion version >> putServerRandom32 random >> putSession session
                               >> putWord16 cipherID >> putWord8 compressionID
                               >> putExtensions exts >> return ()

    encodeHandshakeContent (Certificates certs) = putOpaque24 (runPut $ mapM_ putCert certs)

    encodeHandshakeContent (ClientKeyXchg content) = do
            putBytes content

    encodeHandshakeContent (ServerKeyXchg _) = do
            -- FIXME
            return ()

    encodeHandshakeContent (HelloRequest) = return ()
    encodeHandshakeContent (ServerHelloDone) = return ()

    encodeHandshakeContent (CertRequest certTypes sigAlgs certAuthorities) = do
            putWords8 (map valOfType certTypes)
            case sigAlgs of
                    Nothing -> return ()
                    Just l  -> putWords16 $ map (\(x,y) -> (fromIntegral $ valOfType x) * 256 + (fromIntegral $ valOfType y)) l
            encodeCertAuthorities certAuthorities
      where
        -- Convert a distinguished name to its DER encoding.
        encodeCA dn = return $ B.concat $ L.toChunks $ encodeDN dn

        -- Encode a list of distinguished names.
        encodeCertAuthorities certAuths = do
          enc <- mapM encodeCA certAuths
          let totLength = sum $ map (((+) 2) . B.length) enc
          putWord16 (fromIntegral totLength)
          mapM_ (\ b -> putWord16 (fromIntegral (B.length b)) >> putBytes b) enc

    encodeHandshakeContent (CertVerify mbHashSig (CertVerifyData c)) = do
            -- TLS 1.2 prepends the hash and signature algorithms to the
            -- signature.
            case mbHashSig of
              Nothing -> return ()
              Just (h, s) -> putWord16 $ (fromIntegral $ valOfType h) * 256 + (fromIntegral $ valOfType s)
            putWord16 (fromIntegral $ B.length c)
            putBytes c

    encodeHandshakeContent (Finished opaque) = putBytes opaque

    encodeHandshakeContent (HsNextProtocolNegotiation protocol) = do
            putOpaque8 protocol
            putOpaque8 $ B.replicate paddingLen 0
            where paddingLen = 32 - ((B.length protocol + 2) `mod` 32)

    {- FIXME make sure it return error if not 32 available -}
    getRandom32 :: Get Bytes
    getRandom32 = getBytes 32

    getServerRandom32 :: Get ServerRandom
    getServerRandom32 = ServerRandom <$> getRandom32

    getClientRandom32 :: Get ClientRandom
    getClientRandom32 = ClientRandom <$> getRandom32

    putRandom32 :: Bytes -> Put
    putRandom32 = putBytes

    putClientRandom32 :: ClientRandom -> Put
    putClientRandom32 (ClientRandom r) = putRandom32 r

    putServerRandom32 :: ServerRandom -> Put
    putServerRandom32 (ServerRandom r) = putRandom32 r

    getSession :: Get Session
    getSession = do
            len8 <- getWord8
            case fromIntegral len8 of
                    0   -> return $ Session Nothing
                    len -> Session . Just <$> getBytes len

    putSession :: Session -> Put
    putSession (Session Nothing)  = putWord8 0
    putSession (Session (Just s)) = putOpaque8 s

    putCert :: X509 -> Put
    putCert cert = putOpaque24 (B.concat $ L.toChunks $ encodeCertificate cert)

    getExtensions :: Int -> Get [ExtensionRaw]
    getExtensions 0   = return []
    getExtensions len = do
            extty <- getWord16
            extdatalen <- getWord16
            extdata <- getBytes $ fromIntegral extdatalen
            extxs <- getExtensions (len - fromIntegral extdatalen - 4)
            return $ (extty, extdata) : extxs

    putExtension :: ExtensionRaw -> Put
    putExtension (ty, l) = putWord16 ty >> putOpaque16 l

    putExtensions :: [ExtensionRaw] -> Put
    putExtensions [] = return ()
    putExtensions es = putOpaque16 (runPut $ mapM_ putExtension es)

    {-
     - decode and encode ALERT
     -}

    decodeChangeCipherSpec :: ByteString -> Either TLSError ()
    decodeChangeCipherSpec = runGetErr "changecipherspec" $ do
            x <- getWord8
            when (x /= 1) (fail "unknown change cipher spec content")

    encodeChangeCipherSpec :: ByteString
    encodeChangeCipherSpec = runPut (putWord8 1)

    -- rsa pre master secret
    decodePreMasterSecret :: Bytes -> Either TLSError (Version, Bytes)
    decodePreMasterSecret = runGetErr "pre-master-secret" $ do
            liftM2 (,) getVersion (getBytes 46)

    encodePreMasterSecret :: Version -> Bytes -> Bytes
    encodePreMasterSecret version bytes = runPut (putVersion version >> putBytes bytes)

    {-
     - generate things for packet content
     -}
    type PRF = Bytes -> Bytes -> Int -> Bytes

    generateMasterSecret_SSL :: Bytes -> ClientRandom -> ServerRandom -> Bytes
    generateMasterSecret_SSL premasterSecret (ClientRandom c) (ServerRandom s) =
            B.concat $ map (computeMD5) ["A","BB","CCC"]
            where
                    computeMD5  label = MD5.hash $ B.concat [ premasterSecret, computeSHA1 label ]
                    computeSHA1 label = SHA1.hash $ B.concat [ label, premasterSecret, c, s ]

    generateMasterSecret_TLS :: PRF -> Bytes -> ClientRandom -> ServerRandom -> Bytes
    generateMasterSecret_TLS prf premasterSecret (ClientRandom c) (ServerRandom s) =
            prf premasterSecret seed 48
            where
                    seed = B.concat [ "master secret", c, s ]

    generateMasterSecret :: Version -> Bytes -> ClientRandom -> ServerRandom -> Bytes
    generateMasterSecret SSL2  = generateMasterSecret_SSL
    generateMasterSecret SSL3  = generateMasterSecret_SSL
    generateMasterSecret TLS10 = generateMasterSecret_TLS prf_MD5SHA1
    generateMasterSecret TLS11 = generateMasterSecret_TLS prf_MD5SHA1
    generateMasterSecret TLS12 = generateMasterSecret_TLS prf_SHA256

    generateKeyBlock_TLS :: PRF -> ClientRandom -> ServerRandom -> Bytes -> Int -> Bytes
    generateKeyBlock_TLS prf (ClientRandom c) (ServerRandom s) mastersecret kbsize =
            prf mastersecret seed kbsize where seed = B.concat [ "key expansion", s, c ]

    generateKeyBlock_SSL :: ClientRandom -> ServerRandom -> Bytes -> Int -> Bytes
    generateKeyBlock_SSL (ClientRandom c) (ServerRandom s) mastersecret kbsize =
            B.concat $ map computeMD5 $ take ((kbsize `div` 16) + 1) labels
            where
                    labels            = [ uncurry BC.replicate x | x <- zip [1..] ['A'..'Z'] ]
                    computeMD5  label = MD5.hash $ B.concat [ mastersecret, computeSHA1 label ]
                    computeSHA1 label = SHA1.hash $ B.concat [ label, mastersecret, s, c ]

    generateKeyBlock :: Version -> ClientRandom -> ServerRandom -> Bytes -> Int -> Bytes
    generateKeyBlock SSL2  = generateKeyBlock_SSL
    generateKeyBlock SSL3  = generateKeyBlock_SSL
    generateKeyBlock TLS10 = generateKeyBlock_TLS prf_MD5SHA1
    generateKeyBlock TLS11 = generateKeyBlock_TLS prf_MD5SHA1
    generateKeyBlock TLS12 = generateKeyBlock_TLS prf_SHA256

    generateFinished_TLS :: PRF -> Bytes -> Bytes -> HashCtx -> Bytes
    generateFinished_TLS prf label mastersecret hashctx = prf mastersecret seed 12
            where
                    seed = B.concat [ label, hashFinal hashctx ]

    generateFinished_SSL :: Bytes -> Bytes -> HashCtx -> Bytes
    generateFinished_SSL sender mastersecret hashctx = B.concat [md5hash, sha1hash]
            where
                    md5hash  = MD5.hash $ B.concat [ mastersecret, pad2, md5left ]
                    sha1hash = SHA1.hash $ B.concat [ mastersecret, B.take 40 pad2, sha1left ]

                    lefthash = hashFinal $ flip hashUpdateSSL (pad1, B.take 40 pad1)
                                         $ foldl hashUpdate hashctx [sender,mastersecret]
                    (md5left,sha1left) = B.splitAt 16 lefthash
                    pad2     = B.replicate 48 0x5c
                    pad1     = B.replicate 48 0x36

    generateClientFinished :: Version -> Bytes -> HashCtx -> Bytes
    generateClientFinished ver
            | ver < TLS10 = generateFinished_SSL "CLNT"
            | ver < TLS12 = generateFinished_TLS prf_MD5SHA1 "client finished"
            | otherwise   = generateFinished_TLS prf_SHA256 "client finished"

    generateServerFinished :: Version -> Bytes -> HashCtx -> Bytes
    generateServerFinished ver
            | ver < TLS10 = generateFinished_SSL "SRVR"
            | ver < TLS12 = generateFinished_TLS prf_MD5SHA1 "server finished"
            | otherwise   = generateFinished_TLS prf_SHA256 "server finished"

    generateCertificateVerify_SSL :: Bytes -> HashCtx -> Bytes
    generateCertificateVerify_SSL = generateFinished_SSL ""