{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Web.Rails4.Session (
  -- * Decoding
    decode
  , decodeEither
  -- * Decrypting
  , decrypt
  -- * Utilities
  , csrfToken
  , sessionId
  , lookupString
  , lookupFixnum
  ) where

import              Control.Applicative ((<$>))
import              Crypto.PBKDF.ByteString (sha1PBKDF2)
import              Data.ByteString (ByteString)
import              Data.Either (Either(..), either)
import              Data.Function.Compat ((&))
import              Data.Maybe (Maybe(..), fromMaybe)
import              Data.Monoid ((<>))
import              Data.Ruby.Marshal (RubyObject(..), RubyStringEncoding(..))
import              Data.String.Conv (toS)
import              Network.HTTP.Types (urlDecode)
import              Prelude (Bool(..), Eq, Int, Ord, Show, String, ($!), (.) , (==), const, error, fst, show, snd)
import              Web.Rails.Session.Types
import              Crypto.Cipher.AES (AES256)
import              Crypto.Cipher.Types (cbcDecrypt, cipherInit, makeIV)
import              Crypto.Error (CryptoFailable(CryptoFailed, CryptoPassed))
import qualified    Data.ByteString as BS
import qualified    Data.ByteString.Base64 as B64
import qualified    Data.Ruby.Marshal as Ruby
import qualified    Data.Vector as Vec

-- EXPORTS

-- | Decode a cookie encrypted by Rails.
decode :: Maybe Salt
       -> SecretKeyBase
       -> Cookie
       -> Maybe RubyObject
decode :: Maybe Salt -> SecretKeyBase -> Cookie -> Maybe RubyObject
decode Maybe Salt
mbSalt SecretKeyBase
secretKeyBase Cookie
cookie =
  (String -> Maybe RubyObject)
-> (RubyObject -> Maybe RubyObject)
-> Either String RubyObject
-> Maybe RubyObject
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe RubyObject -> String -> Maybe RubyObject
forall a b. a -> b -> a
const Maybe RubyObject
forall a. Maybe a
Nothing) RubyObject -> Maybe RubyObject
forall a. a -> Maybe a
Just (Maybe Salt -> SecretKeyBase -> Cookie -> Either String RubyObject
decodeEither Maybe Salt
mbSalt SecretKeyBase
secretKeyBase Cookie
cookie)

-- | Decode a cookie encrypted by Rails and retain some error information on failure.
decodeEither :: Maybe Salt
             -> SecretKeyBase
             -> Cookie
             -> Either String RubyObject
decodeEither :: Maybe Salt -> SecretKeyBase -> Cookie -> Either String RubyObject
decodeEither Maybe Salt
mbSalt SecretKeyBase
secretKeyBase Cookie
cookie = do
  case Maybe Salt
-> SecretKeyBase -> Cookie -> Either String DecryptedData
decrypt Maybe Salt
mbSalt SecretKeyBase
secretKeyBase Cookie
cookie of
    Left String
errorMessage ->
      String -> Either String RubyObject
forall a b. a -> Either a b
Left String
errorMessage
    Right (DecryptedData ByteString
deData) ->
      ByteString -> Either String RubyObject
Ruby.decodeEither ByteString
deData

-- | Decrypts a cookie encrypted by Rails. Use this if you are using a
-- serialisation format other than Ruby's Marshal format.
decrypt :: Maybe Salt
        -> SecretKeyBase
        -> Cookie
        -> Either String DecryptedData
decrypt :: Maybe Salt
-> SecretKeyBase -> Cookie -> Either String DecryptedData
decrypt Maybe Salt
mbSalt SecretKeyBase
secretKeyBase Cookie
cookie =
  let salt :: Salt
salt = Salt -> Maybe Salt -> Salt
forall a. a -> Maybe a -> a
fromMaybe Salt
defaultSalt Maybe Salt
mbSalt
      (SecretKey ByteString
secret) = Salt -> SecretKeyBase -> SecretKey
generateSecret Salt
salt SecretKeyBase
secretKeyBase
      (EncryptedData ByteString
encData, InitVector ByteString
initVec) = Cookie -> (EncryptedData, InitVector)
prepare Cookie
cookie
  in case ByteString -> Maybe (IV AES256)
forall b c. (ByteArrayAccess b, BlockCipher c) => b -> Maybe (IV c)
makeIV ByteString
initVec of
       Maybe (IV AES256)
Nothing ->
         String -> Either String DecryptedData
forall a b. a -> Either a b
Left (String -> Either String DecryptedData)
-> String -> Either String DecryptedData
forall a b. (a -> b) -> a -> b
$! String
"Failed to build init. vector for: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ByteString -> String
forall a. Show a => a -> String
show ByteString
initVec
       Just IV AES256
initVec' -> do
         let key :: ByteString
key = Int -> ByteString -> ByteString
BS.take Int
32 ByteString
secret
         case (ByteString -> CryptoFailable AES256
forall cipher key.
(Cipher cipher, ByteArray key) =>
key -> CryptoFailable cipher
forall key. ByteArray key => key -> CryptoFailable AES256
cipherInit ByteString
key :: CryptoFailable AES256) of
           CryptoFailed CryptoError
errorMessage ->
             String -> Either String DecryptedData
forall a b. a -> Either a b
Left (CryptoError -> String
forall a. Show a => a -> String
show CryptoError
errorMessage)
           CryptoPassed AES256
cipher ->
             DecryptedData -> Either String DecryptedData
forall a b. b -> Either a b
Right (DecryptedData -> Either String DecryptedData)
-> (ByteString -> DecryptedData)
-> ByteString
-> Either String DecryptedData
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> DecryptedData
DecryptedData (ByteString -> Either String DecryptedData)
-> ByteString -> Either String DecryptedData
forall a b. (a -> b) -> a -> b
$! AES256 -> IV AES256 -> ByteString -> ByteString
forall cipher ba.
(BlockCipher cipher, ByteArray ba) =>
cipher -> IV cipher -> ba -> ba
forall ba. ByteArray ba => AES256 -> IV AES256 -> ba -> ba
cbcDecrypt AES256
cipher IV AES256
initVec' ByteString
encData
  where
    defaultSalt :: Salt
    defaultSalt :: Salt
defaultSalt = ByteString -> Salt
Salt ByteString
"encrypted cookie"

-- UTIL

-- | Helper function for looking up the csrf token in a cooie.
csrfToken :: RubyObject -> Maybe ByteString
csrfToken :: RubyObject -> Maybe ByteString
csrfToken = ByteString -> RubyStringEncoding -> RubyObject -> Maybe ByteString
lookupString ByteString
"_csrf_token" RubyStringEncoding
US_ASCII

-- | Helper function for looking up the session id in a cookie.
sessionId :: RubyObject -> Maybe ByteString
sessionId :: RubyObject -> Maybe ByteString
sessionId = ByteString -> RubyStringEncoding -> RubyObject -> Maybe ByteString
lookupString ByteString
"session_id" RubyStringEncoding
UTF_8

-- | Lookup integer for a given key.
lookupFixnum :: ByteString -> RubyStringEncoding -> RubyObject -> Maybe Int
lookupFixnum :: ByteString -> RubyStringEncoding -> RubyObject -> Maybe Int
lookupFixnum ByteString
key RubyStringEncoding
enc RubyObject
rubyObject =
  case RubyObject -> RubyObject -> Maybe RubyObject
lookup ((RubyObject, RubyStringEncoding) -> RubyObject
RIVar (ByteString -> RubyObject
RString ByteString
key, RubyStringEncoding
enc)) RubyObject
rubyObject of
    Just (RFixnum Int
val) ->
      Int -> Maybe Int
forall a. a -> Maybe a
Just Int
val
    Maybe RubyObject
_ ->
      Maybe Int
forall a. Maybe a
Nothing

-- | Lookup string for a given key and throw away encoding information.
lookupString :: ByteString
             -> RubyStringEncoding
             -> RubyObject
             -> Maybe ByteString
lookupString :: ByteString -> RubyStringEncoding -> RubyObject -> Maybe ByteString
lookupString ByteString
key RubyStringEncoding
enc RubyObject
rubyObject =
  case RubyObject -> RubyObject -> Maybe RubyObject
lookup ((RubyObject, RubyStringEncoding) -> RubyObject
RIVar (ByteString -> RubyObject
RString ByteString
key, RubyStringEncoding
enc)) RubyObject
rubyObject of
    Just (RIVar (RString ByteString
val, RubyStringEncoding
_)) ->
      ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
val
    Maybe RubyObject
_ ->
      Maybe ByteString
forall a. Maybe a
Nothing

-- PRIVATE

-- | Generate secret key using same cryptographic routines as Rails.
generateSecret :: Salt -> SecretKeyBase -> SecretKey
generateSecret :: Salt -> SecretKeyBase -> SecretKey
generateSecret (Salt ByteString
salt) (SecretKeyBase ByteString
secret) =
  ByteString -> SecretKey
SecretKey (ByteString -> SecretKey) -> ByteString -> SecretKey
forall a b. (a -> b) -> a -> b
$! ByteString -> ByteString -> Int -> Int -> ByteString
sha1PBKDF2 ByteString
secret ByteString
salt Int
1000 Int
64

-- | Prepare a cookie for decryption.
prepare :: Cookie -> (EncryptedData, InitVector)
prepare :: Cookie -> (EncryptedData, InitVector)
prepare (Cookie ByteString
cookie) =
  Bool -> ByteString -> ByteString
urlDecode Bool
True ByteString
cookie
  ByteString -> (ByteString -> ByteString) -> ByteString
forall a b. a -> (a -> b) -> b
& ((ByteString, ByteString) -> ByteString
forall a b. (a, b) -> a
fst ((ByteString, ByteString) -> ByteString)
-> (ByteString -> (ByteString, ByteString))
-> ByteString
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> (ByteString, ByteString)
split)
  ByteString -> (ByteString -> ByteString) -> ByteString
forall a b. a -> (a -> b) -> b
& ByteString -> ByteString
base64decode
  ByteString
-> (ByteString -> (ByteString, ByteString))
-> (ByteString, ByteString)
forall a b. a -> (a -> b) -> b
& ByteString -> (ByteString, ByteString)
split
  (ByteString, ByteString)
-> ((ByteString, ByteString) -> (EncryptedData, InitVector))
-> (EncryptedData, InitVector)
forall a b. a -> (a -> b) -> b
& (\(ByteString
x, ByteString
y) -> (ByteString -> EncryptedData
EncryptedData (ByteString -> ByteString
base64decode ByteString
x), ByteString -> InitVector
InitVector (ByteString -> ByteString
base64decode ByteString
y)))
  where
    base64decode :: ByteString -> ByteString
    base64decode :: ByteString -> ByteString
base64decode = ByteString -> ByteString
B64.decodeLenient

    separator :: ByteString
    separator :: ByteString
separator = ByteString
"--"

    split :: ByteString -> (ByteString, ByteString)
    split :: ByteString -> (ByteString, ByteString)
split = ByteString -> ByteString -> (ByteString, ByteString)
BS.breakSubstring ByteString
separator

-- | Lookup value for a given key.
lookup :: RubyObject -> RubyObject -> Maybe RubyObject
lookup :: RubyObject -> RubyObject -> Maybe RubyObject
lookup RubyObject
key (RHash Vector (RubyObject, RubyObject)
vec) = (RubyObject, RubyObject) -> RubyObject
forall a b. (a, b) -> b
snd ((RubyObject, RubyObject) -> RubyObject)
-> Maybe (RubyObject, RubyObject) -> Maybe RubyObject
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((RubyObject, RubyObject) -> Bool)
-> Vector (RubyObject, RubyObject)
-> Maybe (RubyObject, RubyObject)
forall a. (a -> Bool) -> Vector a -> Maybe a
Vec.find (\(RubyObject, RubyObject)
element -> (RubyObject, RubyObject) -> RubyObject
forall a b. (a, b) -> a
fst (RubyObject, RubyObject)
element RubyObject -> RubyObject -> Bool
forall a. Eq a => a -> a -> Bool
== RubyObject
key) Vector (RubyObject, RubyObject)
vec
lookup RubyObject
_ RubyObject
_ = Maybe RubyObject
forall a. Maybe a
Nothing