-- |
-- Module      : Amazonka.Auth.STS
-- Copyright   : (c) 2013-2023 Brendan Hay
-- License     : Mozilla Public License, v. 2.0.
-- Maintainer  : Brendan Hay <brendan.g.hay+amazonka@gmail.com>
-- Stability   : provisional
-- Portability : non-portable (GHC extensions)
--
-- Retrieve authentication credentials from Secure Token Service
module Amazonka.Auth.STS where

import Amazonka.Auth.Background (fetchAuthInBackground)
import Amazonka.Auth.Exception
import Amazonka.Core.Lens.Internal ((^.))
import Amazonka.Env (Env, Env' (..))
import Amazonka.Prelude
import qualified Amazonka.STS as STS
import qualified Amazonka.STS.AssumeRole as STS
import qualified Amazonka.STS.AssumeRoleWithWebIdentity as STS
import Amazonka.Send (send, sendUnsigned)
import Control.Exception (throw)
import Control.Monad.Trans.Resource (runResourceT)
import qualified Data.Text as Text
import qualified Data.Text.IO as Text
import qualified Data.UUID as UUID
import qualified Data.UUID.V4 as UUID
import qualified System.Environment as Environment

-- | Assume a role using the @sts:AssumeRole@ API.
--
-- This is a simplified interface suitable for most purposes, but if
-- you need the full functionality of the @sts:AssumeRole@ API, you
-- will need to craft your own requests using @amazonka-sts@. If you
-- do this, remember to use 'fetchAuthInBackground' so that your
-- application does not get stuck holding temporary credentials which
-- have expired.
fromAssumedRole ::
  (MonadIO m) =>
  -- | Role ARN
  Text ->
  -- | Role session name
  Text ->
  Env ->
  m Env
fromAssumedRole :: forall (m :: * -> *). MonadIO m => Text -> Text -> Env -> m Env
fromAssumedRole Text
roleArn Text
roleSessionName Env
env = do
  let getCredentials :: IO AuthEnv
getCredentials = do
        let assumeRole :: AssumeRole
assumeRole = Text -> Text -> AssumeRole
STS.newAssumeRole Text
roleArn Text
roleSessionName
        resp <- ResourceT IO AssumeRoleResponse -> IO AssumeRoleResponse
forall (m :: * -> *) a. MonadUnliftIO m => ResourceT m a -> m a
runResourceT (ResourceT IO AssumeRoleResponse -> IO AssumeRoleResponse)
-> ResourceT IO AssumeRoleResponse -> IO AssumeRoleResponse
forall a b. (a -> b) -> a -> b
$ Env -> AssumeRole -> ResourceT IO (AWSResponse AssumeRole)
forall (m :: * -> *) a.
(MonadResource m, AWSRequest a) =>
Env -> a -> m (AWSResponse a)
send Env
env AssumeRole
assumeRole
        pure $ resp ^. STS.assumeRoleResponse_credentials
  keys <- IO Auth -> m Auth
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Auth -> m Auth) -> IO Auth -> m Auth
forall a b. (a -> b) -> a -> b
$ IO AuthEnv -> IO Auth
fetchAuthInBackground IO AuthEnv
getCredentials
  pure env {auth = Identity keys}

-- | https://aws.amazon.com/blogs/opensource/introducing-fine-grained-iam-roles-service-accounts/
-- Obtain temporary credentials from @sts:AssumeRoleWithWebIdentity@.
--
-- The STS service provides an access key, secret key, session token,
-- and expiration time. Also spawns a refresh thread that will
-- periodically fetch fresh credentials before the current ones
-- expire.
--
-- The implementation is modelled on the C++ SDK:
-- https://github.com/aws/aws-sdk-cpp/blob/6d6dcdbfa377393306bf79585f61baea524ac124/aws-cpp-sdk-core/source/auth/STSCredentialsProvider.cpp#L33
fromWebIdentity ::
  (MonadIO m) =>
  -- | Path to token file
  FilePath ->
  -- | Role ARN
  Text ->
  -- | Role Session Name
  Maybe Text ->
  Env' withAuth ->
  m Env
fromWebIdentity :: forall (m :: * -> *) (withAuth :: * -> *).
MonadIO m =>
FilePath -> Text -> Maybe Text -> Env' withAuth -> m Env
fromWebIdentity FilePath
tokenFile Text
roleArn Maybe Text
mSessionName Env' withAuth
env = do
  -- Mimic the C++ SDK; fall back to a random UUID if the session name is unset.
  sessionName <-
    IO Text -> m Text
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Text -> m Text) -> IO Text -> m Text
forall a b. (a -> b) -> a -> b
$ IO Text -> (Text -> IO Text) -> Maybe Text -> IO Text
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (UUID -> Text
UUID.toText (UUID -> Text) -> IO UUID -> IO Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO UUID
UUID.nextRandom) Text -> IO Text
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Text
mSessionName

  -- We copy the behaviour of the C++ implementation: upon credential
  -- expiration, re-read the token file content but ignore any changes
  -- to environment variables.
  let getCredentials = do
        token <- FilePath -> IO Text
Text.readFile FilePath
tokenFile

        let assumeRoleWithWebIdentity =
              Text -> Text -> Text -> AssumeRoleWithWebIdentity
STS.newAssumeRoleWithWebIdentity
                Text
roleArn
                Text
sessionName
                Text
token

        resp <- runResourceT $ sendUnsigned env assumeRoleWithWebIdentity
        pure $ resp ^. STS.assumeRoleWithWebIdentityResponse_credentials

  -- As the credentials from STS are temporary, we start a thread that is able
  -- to fetch new ones automatically on expiry.
  keys <- liftIO $ fetchAuthInBackground getCredentials

  pure env {auth = Identity keys}

-- | Obtain temporary credentials from
-- @sts:AssumeRoleWithWebIdentity@, sourcing arguments from standard
-- environment variables:
--
-- * @AWS_WEB_IDENTITY_TOKEN_FILE@
-- * @AWS_ROLE_ARN@
-- * @AWS_ROLE_SESSION_NAME@ (optional)
--
-- Throws 'MissingEnvError' if a required environment variable is
-- empty or unset.
fromWebIdentityEnv ::
  (MonadIO m) =>
  Env' withAuth ->
  m Env
fromWebIdentityEnv :: forall (m :: * -> *) (withAuth :: * -> *).
MonadIO m =>
Env' withAuth -> m Env
fromWebIdentityEnv Env' withAuth
env = IO Env -> m Env
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Env -> m Env) -> IO Env -> m Env
forall a b. (a -> b) -> a -> b
$ do
  tokenFile <- IO FilePath
lookupTokenFile
  roleArn <- lookupRoleArn
  mSessionName <- lookupSessionName
  fromWebIdentity tokenFile roleArn mSessionName env
  where
    lookupTokenFile :: IO FilePath
lookupTokenFile =
      FilePath -> IO (Maybe FilePath)
nonEmptyEnv FilePath
"AWS_WEB_IDENTITY_TOKEN_FILE" IO (Maybe FilePath)
-> (Maybe FilePath -> IO FilePath) -> IO FilePath
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Just FilePath
v -> FilePath -> IO FilePath
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure FilePath
v
        Maybe FilePath
Nothing ->
          AuthError -> IO FilePath
forall a e. (HasCallStack, Exception e) => e -> a
throw (AuthError -> IO FilePath) -> AuthError -> IO FilePath
forall a b. (a -> b) -> a -> b
$
            Text -> AuthError
MissingEnvError
              Text
"Unable to read token file name from AWS_WEB_IDENTITY_TOKEN_FILE"

    lookupRoleArn :: IO Text
lookupRoleArn =
      FilePath -> IO (Maybe FilePath)
nonEmptyEnv FilePath
"AWS_ROLE_ARN" IO (Maybe FilePath) -> (Maybe FilePath -> IO Text) -> IO Text
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Just FilePath
v -> Text -> IO Text
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> IO Text) -> Text -> IO Text
forall a b. (a -> b) -> a -> b
$ FilePath -> Text
Text.pack FilePath
v
        Maybe FilePath
Nothing ->
          AuthError -> IO Text
forall a e. (HasCallStack, Exception e) => e -> a
throw (AuthError -> IO Text) -> AuthError -> IO Text
forall a b. (a -> b) -> a -> b
$
            Text -> AuthError
MissingEnvError
              Text
"Unable to read role ARN from AWS_ROLE_ARN"

    lookupSessionName :: IO (Maybe Text)
lookupSessionName = (FilePath -> Text) -> Maybe FilePath -> Maybe Text
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap FilePath -> Text
Text.pack (Maybe FilePath -> Maybe Text)
-> IO (Maybe FilePath) -> IO (Maybe Text)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FilePath -> IO (Maybe FilePath)
nonEmptyEnv FilePath
"AWS_ROLE_SESSION_NAME"

    nonEmptyEnv :: String -> IO (Maybe String)
    nonEmptyEnv :: FilePath -> IO (Maybe FilePath)
nonEmptyEnv FilePath
var =
      FilePath -> IO (Maybe FilePath)
Environment.lookupEnv FilePath
var IO (Maybe FilePath)
-> (Maybe FilePath -> Maybe FilePath) -> IO (Maybe FilePath)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
        Maybe FilePath
Nothing -> Maybe FilePath
forall a. Maybe a
Nothing
        Just FilePath
"" -> Maybe FilePath
forall a. Maybe a
Nothing
        Just FilePath
v -> FilePath -> Maybe FilePath
forall a. a -> Maybe a
Just FilePath
v