{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE TemplateHaskell #-}
---------------------------------------------------------
--
-- |
--
-- Module        : Web.ClientSession
-- Copyright     : Michael Snoyman
-- License       : BSD3
--
-- Maintainer    : Michael Snoyman <michael@snoyman.com>
-- Stability     : Stable
-- Portability   : portable
--
-- Stores session data in a client cookie.  In order to do so,
-- we:
--
-- * Encrypt the cookie data using AES in CTR mode.  This allows
-- you to store sensitive information on the client side without
-- worrying about eavesdropping.
--
-- * Sign the encrypted cookie data using Skein-MAC-512-256.
-- Besides detecting potential errors in storage or transmission
-- of the cookies (integrity), the MAC also avoids malicious
-- modifications of the cookie data by assuring you that the
-- cookie data really was generated by this server
-- (authenticity).
--
-- * Encode everything using Base64.  Thus we avoid problems with
-- non-printable characters by giving the browser a simple
-- string.
--
-- Simple usage of the library involves just calling
-- 'getDefaultKey' on the startup of your server, 'encryptIO'
-- when serializing cookies and 'decrypt' when parsing then back.
--
---------------------------------------------------------
module Web.ClientSession
    ( -- * Automatic key generation
      Key(..)
    , IV
    , randomIV
    , mkIV
    , getKey
    , defaultKeyFile
    , getDefaultKey
    , initKey
      -- * Actual encryption/decryption
    , encrypt
    , encryptIO
    , decrypt
    ) where

-- from base
import Control.Monad (guard)

-- from directory
import System.Directory (doesFileExist)

-- from bytestring
import qualified Data.ByteString as S
import qualified Data.ByteString.Base64 as B

-- from cereal
import Data.Serialize (encode, decode)

-- from crypto-api
import Crypto.Classes (buildKey)
import Crypto.Random (newGenIO, genBytes, SystemRandom)
import qualified Crypto.Modes as Modes

-- from cryptocipher
import qualified Crypto.Cipher.AES as A

-- from skein
import Crypto.Skein (skeinMAC', Skein_512_256)

-- | The keys used to store the cookies.  We have an AES key used
-- to encrypt the cookie and a Skein-MAC-512-256 key used verify
-- the authencity and integrity of the cookie.  The AES key needs
-- to have exactly 32 bytes (256 bits) while Skein-MAC-512-256
-- should have 64 bytes (512 bits).
--
-- See also 'getDefaultKey' and 'initKey'.
data Key = Key { aesKey :: A.AES256
                 -- ^ AES key with 32 bytes.
               , macKey :: S.ByteString -> Skein_512_256
                 -- ^ Skein-MAC key.  Instead of storing the key
                 -- data, we store a partially applied function
                 -- for calculating the MAC (see 'skeinMAC'').
               }

-- | Dummy 'Show' instance.
instance Show Key where
    show _ = "<Web.ClientSession.Key>"

-- | The initialization vector used by AES.  Should be exactly 16
-- bytes long.
type IV = Modes.IV A.AES256

-- | Construct an initialization vector from a 'S.ByteString'.
-- Fails if there isn't exactly 16 bytes.
mkIV :: S.ByteString -> Maybe IV
mkIV bs = case (S.length bs, decode bs) of
            (16, Right iv) -> Just iv
            _              -> Nothing

-- | Randomly construct a fresh initialization vector.  You
-- /should not/ reuse initialization vectors.
randomIV :: IO IV
randomIV = Modes.getIVIO

-- | The default key file.
defaultKeyFile :: FilePath
defaultKeyFile = "client_session_key.aes"

-- | Simply calls 'getKey' 'defaultKeyFile'.
getDefaultKey :: IO Key
getDefaultKey = getKey defaultKeyFile

-- | Get a key from the given text file.
--
-- If the file does not exist or is corrupted a random key will
-- be generated and stored in that file.
getKey :: FilePath     -- ^ File name where key is stored.
       -> IO Key       -- ^ The actual key.
getKey keyFile = do
    exists <- doesFileExist keyFile
    if exists
        then S.readFile keyFile >>= either (const newKey) return . initKey
        else newKey
  where
    newKey = do
        (bs, key') <- randomKey
        S.writeFile keyFile bs
        return key'

-- | Generate the given number of random bytes.
randomBytes :: Int -> IO S.ByteString
randomBytes len = do
    g <- newGenIO
    either (error . show) (return . fst) $ genBytes len (g :: SystemRandom)

-- | Generate a random 'Key'.  Besides the 'Key', the
-- 'ByteString' passed to 'initKey' is returned so that it can be
-- saved for later use.
randomKey :: IO (S.ByteString, Key)
randomKey = do
    bs <- randomBytes 96
    case initKey bs of
        Left e -> error $ "Web.ClientSession.randomKey: never here, " ++ e
        Right key -> return (bs, key)

-- | Initializes a 'Key' from a random 'S.ByteString'.  Fails if
-- there isn't exactly 96 bytes (256 bits for AES and 512 bits
-- for Skein-MAC-512-512).
initKey :: S.ByteString -> Either String Key
initKey bs | S.length bs /= 96 = Left $ "Web.ClientSession.initKey: length of " ++
                                         show (S.length bs) ++ " /= 96."
initKey bs = case buildKey preAesKey of
               Nothing -> Left $ "Web.ClientSession.initKey: unknown error with buildKey."
               Just k  -> Right $ Key { aesKey = k
                                      , macKey = skeinMAC' preMacKey }
    where
      (preMacKey, preAesKey) = S.splitAt 64 bs

-- | Same as 'encrypt', however randomly generates the
-- initialization vector for you.
encryptIO :: Key -> S.ByteString -> IO S.ByteString
encryptIO key x = do
    iv <- randomIV
    return $ encrypt key iv x

-- | Encrypt (AES-CTR), sign (Skein-MAC-512-256) and encode
-- (Base64) the given cookie data.  The returned byte string is
-- ready to be used in a response header.
encrypt :: Key          -- ^ Key of the server.
        -> IV           -- ^ New, random initialization vector (see 'randomIV').
        -> S.ByteString -- ^ Serialized cookie data.
        -> S.ByteString -- ^ Encoded cookie data to be given to
                        -- the client browser.
encrypt key iv x = B.encode final
  where
    (encrypted, _) = Modes.ctr' Modes.incIV (aesKey key) iv x
    toBeAuthed     = encode iv `S.append` encrypted
    auth           = macKey key toBeAuthed
    final          = encode auth `S.append` toBeAuthed

-- | Decode (Base64), verify the integrity and authenticity
-- (Skein-MAC-512-256) and decrypt (AES-CTR) the given encoded
-- cookie data.  Returns the original serialized cookie data.
-- Fails if the data is corrupted.
decrypt :: Key                -- ^ Key of the server.
        -> S.ByteString       -- ^ Encoded cookie data given by the browser.
        -> Maybe S.ByteString -- ^ Serialized cookie data.
decrypt key dataBS64 = do
    dataBS <- either (const Nothing) Just $ B.decode dataBS64
    guard (S.length dataBS >= 48) -- 16 bytes of IV + 32 bytes of Skein-MAC-512-256
    let (auth, toBeAuthed) = S.splitAt 32 dataBS
        auth' = macKey key toBeAuthed
    guard (encode auth' == auth)
    let (iv_e, encrypted) = S.splitAt 16 toBeAuthed
    iv <- either (const Nothing) Just $ decode iv_e
    let (x, _) = Modes.unCtr' Modes.incIV (aesKey key) iv encrypted
    return x