Session.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections #-}
{-# OPTIONS_HADDOCK prune #-}

module Network.HStratus.Internal.Session
  ( -- * Credentials
    Credentials (..)

    -- ** paths related to @Credentials@
  , cookiePath
  , clientIdPath
  , credentialsPath
  , savedHeadersPath
  , loginMsgPath

    -- * Session
  , Session (..)
  , SavedHeaders (..)
  , loadSession
  , saveCredentials
  , saveCredentialsTo
  , loadSavedHeaders
  , updateSessionSavedHeaders
  , updateSavedHeaders
  , pristine
  , saveLoginMsg

    -- * AccountData
  , Webservice (..)
  , AccountData (..)
  , accountDataRequires2FA
  , accountDataRequires2SA
  , unknownAccountData
  , accountDataPath
  , saveAccountData
  , loadAccountData

    -- * path components
  , appBase
  , (</>)

    -- * Utilities
  , encodeFileAtomic

    -- * File security
  , checkSecureMode
  , requireSecureFile
  , checkSessionFiles
  )
where

import Control.Applicative ((<|>))
import Control.Exception (bracketOnError, throwIO)
import Control.Monad (forM, when, (>=>))
import Data.Aeson
  ( FromJSON (..)
  , KeyValue (..)
  , Options (..)
  , ToJSON (..)
  , Value
  , eitherDecodeFileStrict
  , encode
  , genericParseJSON
  , genericToEncoding
  , genericToJSON
  , object
  , withObject
  , (.:)
  , (.:?)
  )
import Data.Aeson.Casing (aesonPrefix, snakeCase)
import qualified Data.Aeson.Key as AesonKey
import qualified Data.Aeson.KeyMap as KeyMap
import Data.Bits ((.&.))
import qualified Data.ByteString.Lazy as LBS
import Data.Char (isAlphaNum)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe (catMaybes)
import Data.String.Conv (toS)
import Data.Text (Text)
import qualified Data.Text as Text
import qualified Data.Text.IO as Text
import Data.UUID (toText)
import Data.UUID.V4 (nextRandom)
import GHC.Generics (Generic)
import Network.HStratus.Internal.Http
  ( hCounter
  , hCountry
  , hSessionId
  , hSessionToken
  , hTrustToken
  )
import Network.HTTP.Types.Header (Header)
import Numeric (showOct)
import System.Directory (createDirectoryIfMissing, doesFileExist, removeFile, renameFile)
import System.Environment.XDG.BaseDir (getUserConfigDir)
import System.FilePath (takeDirectory, (</>))
import System.IO (hClose, openTempFile)
import System.Posix.Files (fileMode, getFileStatus, setFileMode)
import System.Posix.Types (FileMode)


-- | Update the @SavedHeaders@ using some response headers
updateSavedHeaders :: [Header] -> SavedHeaders -> SavedHeaders
updateSavedHeaders hs sd =
  sd
    { shCountry = (toS <$> lookup hCountry hs) <|> shCountry sd
    , shSessionId = (toS <$> lookup hSessionId hs) <|> shSessionId sd
    , shSessionToken = (toS <$> lookup hSessionToken hs) <|> shSessionToken sd
    , shTrustToken = (toS <$> lookup hTrustToken hs) <|> shTrustToken sd
    , shCounter = (toS <$> lookup hCounter hs) <|> shCounter sd
    }


data Session = Session
  { sessionCreds :: !Credentials
  -- ^ the credentials used to authenticate
  , sessionTopDir :: !FilePath
  -- ^ directory where session files (cookies, headers, account data) are stored
  , sessionClientId :: !Text
  -- ^ per-client OAuth state identifier sent with each request
  }
  deriving
    ( Eq
      -- ^ don't derive Show to avoid the risk of logging a password
    )


-- | Generates a new client ID.
newClientId :: IO Text
newClientId = ("auth-" <>) . toText <$> nextRandom


{- | Determine the path of file containing the HTTP response headers to be
preserved to maintain a user's authentication state
-}
savedHeadersPath :: FilePath -> Credentials -> FilePath
savedHeadersPath topDir creds = topDir </> Text.unpack (sessionBase creds)


-- | Determine the Cookie Jar file for user with the given credentials
cookiePath :: FilePath -> Credentials -> FilePath
cookiePath topDir creds = topDir </> Text.unpack (cookieBase creds)


{- | Determine the path of file containing the client ID for user with the given
credentials
-}
clientIdPath :: FilePath -> Credentials -> FilePath
clientIdPath topDir creds = topDir </> Text.unpack (clientIdBase creds)


-- | Determine the path of file to save the message the api returns on logon
loginMsgPath :: FilePath -> Credentials -> FilePath
loginMsgPath topDir creds = topDir </> Text.unpack (loginMsgBase creds)


-- | Save the login message to user specific filepath
saveLoginMsg :: Session -> Value -> IO ()
saveLoginMsg Session{sessionCreds = creds, sessionTopDir = topDir} = saveValue (loginMsgPath topDir creds)


-- | Metadata for a single iCloud webservice entry.
data Webservice = Webservice
  { wsUrl :: !Text
  -- ^ the base URL for the service
  , wsStatus :: !(Maybe Text)
  -- ^ service status, e.g. @"active"@ or @"inactive"@; @Nothing@ if absent
  }
  deriving (Eq, Show)


data AccountData = AccountData
  { adHsaVersion :: !Int
  -- ^ HSA protocol version; drives the two-factor flow selection
  , adHsaChallengeRequired :: !Bool
  -- ^ @True@ when a 2FA challenge must be completed before access is granted
  , adHsaTrustedBrowser :: !(Maybe Bool)
  {- ^ @Just True@ when this session is already trusted; @Just False@ when explicitly
  untrusted; @Nothing@ when the key was absent from Apple's response (treated as trusted)
  -}
  , adWebservices :: !(Map Text Webservice)
  -- ^ map of webservice name to service info; use 'lookupWebservice' to resolve a URL
  , adRaw :: !Value
  -- ^ the original JSON value from Apple; preserved so serialisation round-trips losslessly
  }
  deriving (Eq, Show)


instance FromJSON AccountData where
  parseJSON v = withObject "AccountData" go v
   where
    go o = do
      dsInfo <- o .: "dsInfo"
      adHsaVersion <- withObject "dsInfo" (.: "hsaVersion") dsInfo
      adHsaChallengeRequired <- o .:? "hsaChallengeRequired" >>= maybe (pure False) pure
      adHsaTrustedBrowser <- o .:? "hsaTrustedBrowser"
      adWebservices <- do
        mbWs <- o .:? "webservices"
        maybe (pure Map.empty) (withObject "webservices" parseWebservices) mbWs
      pure AccountData{adHsaVersion, adHsaChallengeRequired, adHsaTrustedBrowser, adWebservices, adRaw = v}
    parseWebservices obj = do
      let pairs = KeyMap.toAscList obj
      wsPairs <- forM pairs $ \(k, wsVal) ->
        withObject
          "webservice"
          ( \sv -> do
              mbUrl <- sv .:? "url"
              mbStatus <- sv .:? "status"
              pure $ fmap (\u -> (AesonKey.toText k, Webservice u mbStatus)) mbUrl
          )
          wsVal
      pure $ Map.fromList $ catMaybes wsPairs


instance ToJSON AccountData where
  toJSON AccountData{adRaw} = adRaw
  toEncoding AccountData{adRaw} = toEncoding adRaw


-- | True when full 2FA (auth-endpoint) challenge is required
accountDataRequires2FA :: AccountData -> Bool
accountDataRequires2FA ad =
  adHsaVersion ad == 2
    && (adHsaChallengeRequired ad || adHsaTrustedBrowser ad == Just False)


-- | True when legacy 2SA (setup-endpoint) challenge is required
accountDataRequires2SA :: AccountData -> Bool
accountDataRequires2SA ad = adHsaVersion ad == 1


-- | Sentinel used when no saved @AccountData@ is available
unknownAccountData :: AccountData
unknownAccountData =
  AccountData
    { adHsaVersion = 0
    , adHsaChallengeRequired = False
    , adHsaTrustedBrowser = Nothing
    , adWebservices = Map.empty
    , adRaw = object []
    }


accountDataBase :: Credentials -> Text
accountDataBase = (<> ".account-data.json") . sprucedName


-- | Determine the path of the saved account-data file for the given credentials
accountDataPath :: FilePath -> Credentials -> FilePath
accountDataPath topDir creds = topDir </> Text.unpack (accountDataBase creds)


-- | Persist @AccountData@ to the session's filesystem location
saveAccountData :: Session -> AccountData -> IO ()
saveAccountData Session{sessionCreds = creds, sessionTopDir = topDir} =
  secureEncodeFileAtomic (accountDataPath topDir creds)


-- | Load persisted @AccountData@; returns @Nothing@ if the file is absent
loadAccountData :: Session -> IO (Maybe AccountData)
loadAccountData Session{sessionCreds = creds, sessionTopDir = topDir} = do
  let path = accountDataPath topDir creds
  requireSecureFile path
  exists <- doesFileExist path
  if not exists
    then pure Nothing
    else eitherDecodeFileStrict path >>= either (const (pure Nothing)) (pure . Just)


{- | Determine the path of file containing the credentials in the configuration
   directory
-}
credentialsPath :: FilePath -> FilePath
credentialsPath topDir = topDir </> "credentials.json"


saveCredentials :: Credentials -> IO ()
saveCredentials creds = getUserConfigDir appBase >>= (`saveCredentialsTo` creds)


-- | Write 'Credentials' to @credentials.json@ inside @topDir@, creating @topDir@ if absent.
saveCredentialsTo :: FilePath -> Credentials -> IO ()
saveCredentialsTo topDir creds = do
  createDirectoryIfMissing True topDir
  secureEncodeFileAtomic (credentialsPath topDir) creds


data Credentials = Credentials
  { credAccountName :: !Text
  -- ^ the account ID; typically an email address
  , credPassword :: !Text
  -- ^ the iCloud account password
  }
  deriving
    ( Eq
      -- ^ don't derive Show to avoid the risk of logging a password
    )


instance FromJSON Credentials where
  parseJSON = withObject "Credentials" $ \o ->
    let accountName = o .: "accountName"
        password = o .: "password"
     in Credentials <$> accountName <*> password


instance ToJSON Credentials where
  toJSON c =
    object
      [ "password" .= credPassword c
      , "accountName" .= credAccountName c
      ]


sprucedName :: Credentials -> Text
sprucedName =
  let p aChar = isAlphaNum aChar || aChar == '@'
      replaceAt = Text.replace "@" "-"
   in replaceAt . Text.filter p . credAccountName


cookieBase :: Credentials -> Text
cookieBase = (<> ".cookies.txt") . sprucedName


sessionBase :: Credentials -> Text
sessionBase = (<> ".session.json") . sprucedName


clientIdBase :: Credentials -> Text
clientIdBase = (<> ".client-id.txt") . sprucedName


loginMsgBase :: Credentials -> Text
loginMsgBase = (<> ".last-logon.json") . sprucedName


-- | Data obtained from HTTP response headers that define a user session
data SavedHeaders = SavedHeaders
  { shCountry :: !(Maybe Text)
  -- ^ X-Apple-ID-Country value from the last response
  , shSessionId :: !(Maybe Text)
  -- ^ X-Apple-ID-Session-Id value from the last response
  , shSessionToken :: !(Maybe Text)
  -- ^ X-Apple-Session-Token value from the last response
  , shTrustToken :: !(Maybe Text)
  -- ^ X-Apple-TwoSV-Trust-Token value from the last response
  , shCounter :: !(Maybe Text)
  -- ^ X-Apple-HC-Bits value from the last response; used to derive hashcash proofs
  }
  deriving (Eq, Show, Generic)


instance FromJSON SavedHeaders where
  parseJSON = genericParseJSON simpleOptions


instance ToJSON SavedHeaders where
  toJSON = genericToJSON simpleOptions
  toEncoding = genericToEncoding simpleOptions


-- | A @SavedHeaders@ with nothing set
pristine :: SavedHeaders
pristine = SavedHeaders Nothing Nothing Nothing Nothing Nothing


{- | Update the stored saved headers

if the sessionData file exists
then
  load it.
  update the session data from the headers
  save the updated data
else
  ensure its parent directory exists
  create the session data from the headers
  save it

not handled (thrown as IOException):
  cannot create directory
  cannot write due to permissions
  file exists, but data cannot be parsed
-}
updateSessionSavedHeaders
  :: Session
  -> (SavedHeaders -> SavedHeaders)
  -- ^ a function that modifies the session's saved headers
  -> IO ()
updateSessionSavedHeaders s modSavedHeaders = do
  let dataPath = savedHeadersPath (sessionTopDir s) (sessionCreds s)
      updateAndSave = secureEncodeFileAtomic dataPath . modSavedHeaders
      loadLast False = pure pristine
      loadLast True = eitherDecodeFileStrict dataPath >>= either (fail . show) pure

  doesFileExist dataPath >>= loadLast >>= updateAndSave


loadSession :: IO Session
loadSession = do
  sessionTopDir <- getUserConfigDir appBase
  createDirectoryIfMissing True sessionTopDir
  s <-
    loadSessionOr sessionTopDir
      >>= orFail "Credentials are missing or corrupt; run 'hstratus auth login' to authenticate"
  checkSessionFiles s
  pure s


-- | Saves a JSON @Value@ to @filepath@
saveValue :: FilePath -> Value -> IO ()
saveValue fp v = LBS.writeFile fp $ encode v


orFail :: String -> Either String a -> IO a
orFail hint = either (\e -> fail (hint <> " (" <> e <> ")")) pure


-- | Write a JSON-encodable value to @path@ atomically via a temp file and rename.
encodeFileAtomic :: (ToJSON a) => FilePath -> a -> IO ()
encodeFileAtomic path value =
  bracketOnError
    (openTempFile (takeDirectory path) ".tmp")
    (\(tmpPath, h) -> hClose h >> removeFile tmpPath)
    ( \(tmpPath, h) -> do
        LBS.hPut h (encode value)
        hClose h
        renameFile tmpPath path
    )


{- | Like 'encodeFileAtomic' but sets mode @0o600@ on the temp file before
renaming, so the file is never visible at a more permissive mode.
-}
secureEncodeFileAtomic :: (ToJSON a) => FilePath -> a -> IO ()
secureEncodeFileAtomic path value =
  bracketOnError
    (openTempFile (takeDirectory path) ".tmp")
    (\(tmpPath, h) -> hClose h >> removeFile tmpPath)
    ( \(tmpPath, h) -> do
        LBS.hPut h (encode value)
        hClose h
        setFileMode tmpPath 0o600
        renameFile tmpPath path
    )


-- | Write text to @path@ atomically, setting mode @0o600@ before renaming.
secureWriteTextFileAtomic :: FilePath -> Text -> IO ()
secureWriteTextFileAtomic path content =
  bracketOnError
    (openTempFile (takeDirectory path) ".tmp")
    (\(tmpPath, h) -> hClose h >> removeFile tmpPath)
    ( \(tmpPath, h) -> do
        Text.hPutStr h content
        hClose h
        setFileMode tmpPath 0o600
        renameFile tmpPath path
    )


{- | Pure permission check using the SSH convention: no group or world bits may
be set.  Returns @Left@ with a descriptive message (including a @chmod 600@ hint)
when the mode is too permissive.
-}
checkSecureMode :: FileMode -> FilePath -> Either String ()
checkSecureMode mode path
  | mode .&. 0o077 /= 0 =
      Left $
        path
          <> " has unsafe permissions ("
          <> showOct (fromIntegral mode :: Int) ""
          <> "); fix with: chmod 600 "
          <> path
  | otherwise = Right ()


{- | Verify that a file has secure permissions before it is read.  Does nothing
if the file does not exist; absence is handled by the caller.  Throws an
'IOError' when the file exists but its mode has group or world bits set.
-}
requireSecureFile :: FilePath -> IO ()
requireSecureFile path = do
  exists <- doesFileExist path
  when exists $ do
    mode <- fileMode <$> getFileStatus path
    either (throwIO . userError) pure (checkSecureMode mode path)


{- | Check that every session file that exists has secure permissions.  Absent
files are skipped silently.  Throws an 'IOError' for the first file found with
group or world bits set.

Covers all five session paths: credentials, saved headers, client ID, account
data, and the cookie jar.
-}
checkSessionFiles :: Session -> IO ()
checkSessionFiles sess = do
  let topDir = sessionTopDir sess
      creds = sessionCreds sess
  mapM_
    requireSecureFile
    [ credentialsPath topDir
    , savedHeadersPath topDir creds
    , clientIdPath topDir creds
    , accountDataPath topDir creds
    , cookiePath topDir creds
    ]


loadCredentials :: FilePath -> IO (Either String Credentials)
loadCredentials topDir = do
  let path = credentialsPath topDir
  requireSecureFile path
  eitherDecodeFileStrict path


loadCredentials' :: FilePath -> IO (Either String (FilePath, Credentials))
loadCredentials' topDir = fmap (topDir,) <$> loadCredentials topDir


loadSession' :: Either String (FilePath, Credentials) -> IO (Either String Session)
loadSession' (Left err) = pure $ Left err
loadSession' (Right (sessionTopDir, sessionCreds)) = do
  sessionClientId <- loadClientId sessionTopDir sessionCreds
  pure $ Right Session{sessionClientId, sessionCreds, sessionTopDir}


loadSessionOr :: FilePath -> IO (Either String Session)
loadSessionOr = loadCredentials' >=> loadSession'


-- | Load the @SavedHeaders@ for this session
loadSavedHeaders :: Session -> IO SavedHeaders
loadSavedHeaders Session{sessionTopDir, sessionCreds} =
  loadSavedHeaders' sessionTopDir sessionCreds
    >>= orFail "Session state is corrupt; run 'hstratus auth login' to re-authenticate"


loadSavedHeaders' :: FilePath -> Credentials -> IO (Either String SavedHeaders)
loadSavedHeaders' topDir creds = do
  let dataPath = savedHeadersPath topDir creds
  requireSecureFile dataPath
  pathExists <- doesFileExist dataPath
  if not pathExists
    then pure $ Right pristine
    else eitherDecodeFileStrict dataPath


loadClientId :: FilePath -> Credentials -> IO Text
loadClientId topDir creds = do
  let dataPath = clientIdPath topDir creds
  requireSecureFile dataPath
  pathExists <- doesFileExist dataPath
  if pathExists
    then Text.readFile dataPath
    else do
      anId <- newClientId
      secureWriteTextFileAtomic dataPath anId
      pure anId


simpleOptions :: Options
simpleOptions = aesonPrefix snakeCase


appBase :: FilePath
appBase = "hstratus"