{-# LANGUAGE OverloadedStrings #-}
-- | Basic session support. Sessions are (currently) just maps of strings to
--   strings which are serialised and sent to the client using a Cookie.
--   Serialised, they should be less than 4K, so the sum of all your strings in
--   the map should be < 3900 bytes to be safe.
--
--   The cookies are HMACed and encrypted so that the client can't inspect nor
--   alter them. The key is, by default, generated randomly every time the
--   server starts. If you want the cookies to be reusable across restarts or
--   servers you need to set the key yourself.
--
--   Also, by the default, the cookies are set to expire when the browser
--   session ends.
module Network.MiniHTTP.Session
  ( getSession
  , putSession
  , addSession
  , setSessionSecretKey
  ) where

import           Control.Concurrent.STM
import           Control.Monad (forM_)
import           Control.Monad.Trans (liftIO)

import qualified Data.Map as Map
import qualified Data.Binary.Put           as P
import qualified Data.Binary.Strict.Get           as G
import           Data.Bits
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 ()
import qualified Data.ByteString.Lazy as BL
import           Data.ByteString.Internal (w2c)
import           Data.Word (Word32)

import           Network.MiniHTTP.Server
import           Network.MiniHTTP.Marshal

import           System.IO.Unsafe (unsafePerformIO)

import qualified OpenSSL.EVP.Base64 as Base64
import qualified OpenSSL.EVP.Cipher as Cipher
import qualified OpenSSL.EVP.Digest           as Digest
import           OpenSSL.Random (randBytes)

toString :: B.ByteString -> String
toString = map w2c . B.unpack

secretKey :: TVar B.ByteString
secretKey = unsafePerformIO $ do
  rand <- randBytes 16
  newTVarIO rand

-- | Set the secret key used to HMAC and encrypt the session cookies.
setSessionSecretKey :: B.ByteString -> STM ()
setSessionSecretKey = writeTVar secretKey

sha256 :: Digest.Digest
Just sha256 = unsafePerformIO $ Digest.getDigestByName "SHA256"
aes128 :: Cipher.Cipher
Just aes128 = unsafePerformIO $ Cipher.getCipherByName "AES128"

parseMap :: G.Get [(B.ByteString, B.ByteString)]
parseMap = do
  emptyp <- G.isEmpty
  if emptyp
     then return []
     else do
       keylen <- varInt32
       key <- G.getByteString keylen
       valuelen <- varInt32
       value <- G.getByteString valuelen
       rest <- parseMap
       rest `seq` (return $ (key, value) : rest)

putMap :: [(B.ByteString, B.ByteString)] -> P.Put
putMap vs = do
  forM_ vs $ \(key, value) -> do
    putVarInt32 $ B.length key
    P.putByteString key
    putVarInt32 $ B.length value
    P.putByteString value

decode :: B.ByteString -> IO (Map.Map B.ByteString B.ByteString)
decode bs = do
  secret <- atomically $ readTVar secretKey
  let bs' = Base64.decodeBase64BS bs
      (mac, bs'') = B.splitAt 32 bs'
  if Digest.hmacBS sha256 secret bs'' /= mac
     then return Map.empty
     else do
       let (iv, ciphertext) = B.splitAt 16 bs''
       plaintext <- Cipher.cipherBS aes128 (toString secret) (toString iv) Cipher.Decrypt ciphertext
       case G.runGet parseMap plaintext of
            (Left _, _) -> return Map.empty
            (Right m, _) -> return $ Map.fromList m

encode :: Map.Map B.ByteString B.ByteString -> IO B.ByteString
encode map = do
  secret <- atomically $ readTVar secretKey
  iv <- randBytes 16
  let plaintext = B.concat $ BL.toChunks $ P.runPut $ putMap $ Map.toList map
  ciphertext <- Cipher.cipherBS aes128 (toString secret) (toString iv) Cipher.Encrypt plaintext
  let bs = iv `B.append` ciphertext
      mac = Digest.hmacBS sha256 secret bs
  return $ Base64.encodeBase64BS $ mac `B.append` bs

-- | Return the current session. If the user didn't present a cookie, or the
--   cookie is invalid, an empty map is returned.
getSession :: WebMonad (Map.Map B.ByteString B.ByteString)
getSession = do
  req <- getRequest
  case filter (\cookie -> cookieName cookie == "nmhs") $ httpCookie $ reqHeaders req of
       [] -> return Map.empty
       x:_ -> liftIO $ decode $ cookieValue x

-- | Set the current session. This alters the headers of the current request,
--   so future actions which reset the headers (like 'setReply') will undo
--   this.
putSession :: Map.Map B.ByteString B.ByteString -> WebMonad ()
putSession m = do
  v <- liftIO $ encode m
  setCookie $ emptyCookie { cookieName = "nmhs"
                          , cookieValue = v
                          , cookiePath = Just "/" }

-- | Add a key value pair to the session
addSession :: B.ByteString -> B.ByteString -> WebMonad ()
addSession key value = do
  s <- getSession
  putSession $ Map.insert key value s

-- | This puts a generic Integral type, but first it casts to a Word32 because
--   it makes a different when it comes to negative numbers.
varInt32 :: (Integral i) => G.Get i
varInt32 = varIntWord32 >>= return . fromIntegral

varIntWord32 :: G.Get Word32
varIntWord32 = do
  byte <- G.getWord8
  if byte `testBit` 7
     then do rest <- varIntWord32
             return $ (rest `shiftL` 7) .|. (fromIntegral $ byte .&. 0x7f)
     else return $ fromIntegral byte

-- | This puts a generic Integral type, but first it casts to a Word32 because
--   it makes a difference with negative numbers
putVarInt32 :: (Integral i) => i -> P.Put
putVarInt32 = putVarWord32 . fromIntegral

putVarWord32 :: Word32 -> P.Put
putVarWord32 value
  | value < 128 = P.putWord8 $ fromIntegral value
  | otherwise = do
      P.putWord8 $ 0x80 .|. (fromIntegral $ value .&. 0x7f)
      putVarWord32 $ value `shiftR` 7