{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE Rank2Types                #-}

------------------------------------------------------------------------------
-- | Pre-packaged Handlers that deal with form submissions and standard
--   use-cases involving authentication.

module Snap.Snaplet.Auth.Handlers where

------------------------------------------------------------------------------
import           Control.Applicative
import           Control.Error
import           Control.Monad.State
import           Data.ByteString (ByteString)
import           Data.Serialize hiding (get)
import           Data.Time
import           Data.Text.Encoding (decodeUtf8)
import           Data.Text (Text, null, strip)
import           Prelude hiding (null)
import           Web.ClientSession
------------------------------------------------------------------------------
import           Snap.Core
import           Snap.Snaplet
import           Snap.Snaplet.Auth.AuthManager
import           Snap.Snaplet.Auth.Types
import           Snap.Snaplet.Session
------------------------------------------------------------------------------


                         ----------------------------
                         -- Higher level functions --
                         ----------------------------

------------------------------------------------------------------------------
-- | Create a new user from just a username and password
--
createUser :: Text              -- ^ Username
           -> ByteString        -- ^ Password
           -> Handler b (AuthManager b) (Either AuthFailure AuthUser)
createUser unm pwd
  | null $ strip unm = return $ Left UsernameMissing
  | otherwise = do
     uExists <- usernameExists unm
     if uExists then return $ Left DuplicateLogin
                else withBackend $ \r -> liftIO $ buildAuthUser r unm pwd


------------------------------------------------------------------------------
-- | Check whether a user with the given username exists.
--
usernameExists :: Text          -- ^ The username to be checked
               -> Handler b (AuthManager b) Bool
usernameExists username =
    withBackend $ \r -> liftIO $ isJust <$> lookupByLogin r username


------------------------------------------------------------------------------
-- | Lookup a user by her username, check given password and perform login
--
loginByUsername :: Text             -- ^ Username/login for user
                -> Password         -- ^ Should be ClearText
                -> Bool             -- ^ Set remember token?
                -> Handler b (AuthManager b) (Either AuthFailure AuthUser)
loginByUsername _ (Encrypted _) _ = return $ Left EncryptedPassword
loginByUsername unm pwd shouldRemember = do
    sk <- gets siteKey
    cn <- gets rememberCookieName
    rp <- gets rememberPeriod
    withBackend $ loginByUsername' sk cn rp

  where
    --------------------------------------------------------------------------
    loginByUsername' :: (IAuthBackend t) =>
                        Key
                     -> ByteString
                     -> Maybe Int
                     -> t
                     -> Handler b (AuthManager b) (Either AuthFailure AuthUser)
    loginByUsername' sk cn rp r =
        liftIO (lookupByLogin r unm) >>=
        maybe (return $! Left UserNotFound) found

      where
        ----------------------------------------------------------------------
        found user = checkPasswordAndLogin user pwd >>=
                     either (return . Left) matched

        ----------------------------------------------------------------------
        matched user
            | shouldRemember = do
                  token <- gets randomNumberGenerator >>=
                           liftIO . randomToken 64

                  setRememberToken sk cn rp token

                  let user' = user {
                                userRememberToken = Just (decodeUtf8 token)
                              }

                  saveUser user'
                  return $! Right user'

            | otherwise = return $ Right user


------------------------------------------------------------------------------
-- | Remember user from the remember token if possible and perform login
--
loginByRememberToken :: Handler b (AuthManager b) (Either AuthFailure AuthUser)
loginByRememberToken = withBackend $ \impl -> do
    key         <- gets siteKey
    cookieName_ <- gets rememberCookieName
    period      <- gets rememberPeriod

    runEitherT $ do
        token <- noteT (AuthError "loginByRememberToken: no remember token") $
                   MaybeT $ getRememberToken key cookieName_ period
        user  <- noteT (AuthError "loginByRememberToken: no remember token") $
                   MaybeT $ liftIO $ lookupByRememberToken impl
                                   $ decodeUtf8 token
        lift $ forceLogin user
        return user


------------------------------------------------------------------------------
-- | Logout the active user
--
logout :: Handler b (AuthManager b) ()
logout = do
    s <- gets session
    withTop s $ withSession s removeSessionUserId
    rc <- gets rememberCookieName
    forgetRememberToken rc
    modify $ \mgr -> mgr { activeUser = Nothing }


------------------------------------------------------------------------------
-- | Return the current user; trying to remember from cookie if possible.
--
currentUser :: Handler b (AuthManager b) (Maybe AuthUser)
currentUser = cacheOrLookup $ withBackend $ \r -> do
    s   <- gets session
    uid <- withTop s getSessionUserId
    case uid of
      Nothing -> hush <$> loginByRememberToken
      Just uid' -> liftIO $ lookupByUserId r uid'


------------------------------------------------------------------------------
-- | Convenience wrapper around 'rememberUser' that returns a bool result
--
isLoggedIn :: Handler b (AuthManager b) Bool
isLoggedIn = isJust <$> currentUser


------------------------------------------------------------------------------
-- | Create or update a given user
--
saveUser :: AuthUser -> Handler b (AuthManager b) (Either AuthFailure AuthUser)
saveUser u
    | null $ userLogin u = return $ Left UsernameMissing
    | otherwise = withBackend $ \r -> liftIO $ save r u


------------------------------------------------------------------------------
-- | Destroy the given user
--
destroyUser :: AuthUser -> Handler b (AuthManager b) ()
destroyUser u = withBackend $ liftIO . flip destroy u


                      -----------------------------------
                      --  Lower level helper functions --
                      -----------------------------------

------------------------------------------------------------------------------
-- | Mutate an 'AuthUser', marking failed authentication
--
-- This will save the user to the backend.
--
markAuthFail :: AuthUser
             -> Handler b (AuthManager b) (Either AuthFailure AuthUser)
markAuthFail u = withBackend $ \r -> do
    lo <- gets lockout
    incFailCtr u >>= checkLockout lo >>= liftIO . save r

  where
    --------------------------------------------------------------------------
    incFailCtr u' = return $ u' {
                      userFailedLoginCount = userFailedLoginCount u' + 1
                    }

    --------------------------------------------------------------------------
    checkLockout lo u' =
        case lo of
          Nothing          -> return u'
          Just (mx, wait)  ->
              if userFailedLoginCount u' >= mx
                then do
                  now <- liftIO getCurrentTime
                  let reopen = addUTCTime wait now
                  return $! u' { userLockedOutUntil = Just reopen }
                else return u'


------------------------------------------------------------------------------
-- | Mutate an 'AuthUser', marking successful authentication
--
-- This will save the user to the backend.
--
markAuthSuccess :: AuthUser
                -> Handler b (AuthManager b) (Either AuthFailure AuthUser)
markAuthSuccess u = withBackend $ \r ->
                        incLoginCtr u     >>=
                        updateIp          >>=
                        updateLoginTS     >>=
                        resetFailCtr      >>=
                        liftIO . save r
  where
    --------------------------------------------------------------------------
    incLoginCtr u' = return $ u' { userLoginCount = userLoginCount u' + 1 }

    --------------------------------------------------------------------------
    updateIp u' = do
        ip <- rqRemoteAddr <$> getRequest
        return $ u' { userLastLoginIp = userCurrentLoginIp u'
                    , userCurrentLoginIp = Just ip }

    --------------------------------------------------------------------------
    updateLoginTS u' = do
        now <- liftIO getCurrentTime
        return $
          u' { userCurrentLoginAt = Just now
             , userLastLoginAt = userCurrentLoginAt u' }

    --------------------------------------------------------------------------
    resetFailCtr u' = return $ u' { userFailedLoginCount = 0
                                  , userLockedOutUntil = Nothing }


------------------------------------------------------------------------------
-- | Authenticate and log the user into the current session if successful.
--
-- This is a mid-level function exposed to allow roll-your-own ways of looking
-- up a user from the database.
--
-- This function will:
--
-- 1. Check the password
--
-- 2. Login the user into the current session
--
-- 3. Mark success/failure of the authentication trial on the user record
--
checkPasswordAndLogin
  :: AuthUser               -- ^ An existing user, somehow looked up from db
  -> Password               -- ^ A ClearText password
  -> Handler b (AuthManager b) (Either AuthFailure AuthUser)
checkPasswordAndLogin u pw =
    case userLockedOutUntil u of
      Just x -> do
        now <- liftIO getCurrentTime
        if now > x
          then auth u
          else return . Left $ LockedOut x
      Nothing -> auth u

  where
    auth :: AuthUser -> Handler b (AuthManager b) (Either AuthFailure AuthUser)
    auth user =
      case authenticatePassword user pw of
        Just e -> do
          markAuthFail user
          return $ Left e

        Nothing -> do
          forceLogin user
          modify (\mgr -> mgr { activeUser = Just user })
          markAuthSuccess user


------------------------------------------------------------------------------
-- | Login and persist the given 'AuthUser' in the active session
--
-- Meant to be used if you have other means of being sure that the person is
-- who she says she is.
--
forceLogin :: AuthUser       -- ^ An existing user, somehow looked up from db
           -> Handler b (AuthManager b) (Either AuthFailure ())
forceLogin u = do
    s <- gets session
    withSession s $
        case userId u of
          Just x -> do
            withTop s (setSessionUserId x)
            return $ Right ()
          Nothing -> return . Left $
                     AuthError $ "forceLogin: Can't force the login of a user "
                                   ++ "without userId"


                     ------------------------------------
                     -- Internal, non-exported helpers --
                     ------------------------------------


------------------------------------------------------------------------------
getRememberToken :: (Serialize t, MonadSnap m)
                 => Key
                 -> ByteString
                 -> Maybe Int
                 -> m (Maybe t)
getRememberToken sk rc rp = getSecureCookie rc sk rp


------------------------------------------------------------------------------
setRememberToken :: (Serialize t, MonadSnap m)
                 => Key
                 -> ByteString
                 -> Maybe Int
                 -> t
                 -> m ()
setRememberToken sk rc rp token = setSecureCookie rc sk rp token


------------------------------------------------------------------------------
forgetRememberToken :: MonadSnap m => ByteString -> m ()
forgetRememberToken rc = expireCookie rc (Just "/")


------------------------------------------------------------------------------
-- | Set the current user's 'UserId' in the active session
--
setSessionUserId :: UserId -> Handler b SessionManager ()
setSessionUserId (UserId t) = setInSession "__user_id" t


------------------------------------------------------------------------------
-- | Remove 'UserId' from active session, effectively logging the user out.
removeSessionUserId :: Handler b SessionManager ()
removeSessionUserId = deleteFromSession "__user_id"


------------------------------------------------------------------------------
-- | Get the current user's 'UserId' from the active session
--
getSessionUserId :: Handler b SessionManager (Maybe UserId)
getSessionUserId = do
  uid <- getFromSession "__user_id"
  return $ liftM UserId uid


------------------------------------------------------------------------------
-- | Check password for a given user.
--
-- Returns "Nothing" if check is successful and an "IncorrectPassword" error
-- otherwise
--
authenticatePassword :: AuthUser        -- ^ Looked up from the back-end
                     -> Password        -- ^ Check against this password
                     -> Maybe AuthFailure
authenticatePassword u pw = auth
  where
    auth    = case userPassword u of
                Nothing -> Just PasswordMissing
                Just upw -> check $ checkPassword pw upw

    check b = if b then Nothing else Just IncorrectPassword


------------------------------------------------------------------------------
-- | Wrap lookups around request-local cache
--
cacheOrLookup
  :: Handler b (AuthManager b) (Maybe AuthUser)
      -- ^ Lookup action to perform if request local cache is empty
  -> Handler b (AuthManager b) (Maybe AuthUser)
cacheOrLookup f = do
    au <- gets activeUser
    if isJust au
      then return au
      else do
        au' <- f
        modify (\mgr -> mgr { activeUser = au' })
        return au'


------------------------------------------------------------------------------
-- | Register a new user by specifying login and password 'Param' fields
--
registerUser
  :: ByteString            -- ^ Login field
  -> ByteString            -- ^ Password field
  -> Handler b (AuthManager b) (Either AuthFailure AuthUser)
registerUser lf pf = do
    l <- fmap decodeUtf8 <$> getParam lf
    p <- getParam pf

    let l' = note UsernameMissing l
    let p' = note PasswordMissing p

    -- In case of multiple AuthFailure, the first available one
    -- will be propagated.
    case liftM2 (,) l' p' of
      Left e           -> return $ Left e
      Right (lgn, pwd) -> createUser lgn pwd


------------------------------------------------------------------------------
-- | A 'MonadSnap' handler that processes a login form.
--
-- The request paremeters are passed to 'performLogin'
--
-- To make your users stay logged in for longer than the session replay
-- prevention timeout, you must pass a field name as the third parameter and
-- that field must be set to a value of \"1\" by the submitting form.  This
-- lets you use a user selectable check box.  Or if you want user remembering
-- always turned on, you can use a hidden form field.
loginUser
  :: ByteString
      -- ^ Username field
  -> ByteString
      -- ^ Password field
  -> Maybe ByteString
      -- ^ Remember field; Nothing if you want no remember function.
  -> (AuthFailure -> Handler b (AuthManager b) ())
      -- ^ Upon failure
  -> Handler b (AuthManager b) ()
      -- ^ Upon success
  -> Handler b (AuthManager b) ()
loginUser unf pwdf remf loginFail loginSucc =
    runEitherT (loginUser' unf pwdf remf)
    >>= either loginFail (const loginSucc)


------------------------------------------------------------------------------
loginUser' :: ByteString
           -> ByteString
           -> Maybe ByteString
           -> EitherT AuthFailure (Handler b (AuthManager b)) AuthUser
loginUser' unf pwdf remf = do
    mbUsername <- lift $ getParam unf
    mbPassword <- lift $ getParam pwdf
    remember   <- lift $ liftM (fromMaybe False)
                    (runMaybeT $
                    do field <- MaybeT $ return remf
                       value <- MaybeT $ getParam field
                       return $ value == "1" || value == "on")

    password <- noteT PasswordMissing $ hoistMaybe mbPassword
    username <- noteT UsernameMissing $ hoistMaybe mbUsername

    EitherT $ loginByUsername (decodeUtf8 username)
                              (ClearText password) remember


------------------------------------------------------------------------------
-- | Simple handler to log the user out. Deletes user from session.
--
logoutUser :: Handler b (AuthManager b) ()   -- ^ What to do after logging out
           -> Handler b (AuthManager b) ()
logoutUser target = logout >> target


------------------------------------------------------------------------------
-- | Require that an authenticated 'AuthUser' is present in the current
-- session.
--
-- This function has no DB cost - only checks to see if a user_id is present
-- in the current session.
--
requireUser :: SnapletLens b (AuthManager b)
                -- ^ Lens reference to an "AuthManager"
            -> Handler b v a
                -- ^ Do this if no authenticated user is present.
            -> Handler b v a
                -- ^ Do this if an authenticated user is present.
            -> Handler b v a
requireUser auth bad good = do
    loggedIn <- withTop auth isLoggedIn
    if loggedIn then good else bad


------------------------------------------------------------------------------
-- | Run a function on the backend, and return the result.
--
-- This uses an existential type so that the backend type doesn't
-- 'escape' AuthManager.  The reason that the type is Handler b
-- (AuthManager v) a and not a is because anything that uses the
-- backend will return an IO something, which you can liftIO, or a
-- Handler b (AuthManager v) a if it uses other handler things.
--
withBackend ::
    (forall r. (IAuthBackend r) => r -> Handler b (AuthManager v) a)
      -- ^ The function to run with the handler.
  -> Handler b (AuthManager v) a
withBackend f = join $ do
  (AuthManager backend_ _ _ _ _ _ _ _ _) <- get
  return $ f backend_


------------------------------------------------------------------------------
-- | This function generates a random password reset token and stores it in
-- the database for the user.  Call this function when a user forgets their
-- password.  Then use the token to autogenerate a link that the user can
-- visit to reset their password.  This function also sets a timestamp so the
-- reset token can be expired.
setPasswordResetToken :: Text -> Handler b (AuthManager b) (Maybe Text)
setPasswordResetToken login = do
  tokBS <- liftIO . randomToken 40 =<< gets randomNumberGenerator
  let token = decodeUtf8 tokBS
  now <- liftIO getCurrentTime
  success <- modPasswordResetToken login (Just token) (Just now)
  return $ if success then Just token else Nothing


------------------------------------------------------------------------------
-- | Clears a user's password reset token.  Call this when the user
-- successfully changes their password to ensure that the password reset link
-- cannot be used again.
clearPasswordResetToken :: Text -> Handler b (AuthManager b) Bool
clearPasswordResetToken login = modPasswordResetToken login Nothing Nothing


------------------------------------------------------------------------------
-- | Helper function used for setting and clearing the password reset token
-- and associated timestamp.
modPasswordResetToken :: Text
                      -> Maybe Text
                      -> Maybe UTCTime
                      -> Handler v (AuthManager v) Bool
modPasswordResetToken login token timestamp = do
  res <- runMaybeT $ do
      u <- MaybeT $ withBackend $ \b -> liftIO $ lookupByLogin b login
      lift $ saveUser $ u
        { userResetToken = token
        , userResetRequestedAt = timestamp
        }
      return ()
  return $ maybe False (\_ -> True) res