#!/bin/bash
set -eu
# AWSに環境変数`$AWS_PROFILE`など、aws cliが理解出来る形式でログイン出来るようになっている必要があります。
# jq, batが必要です。

# aws stsで一時的なキーを発行する。
creds_json=$(aws sts assume-role --role-arn foo --role-session-name bar)
AWS_ACCESS_KEY_ID=$(echo "$creds_json"|jq .Credentials.AccessKeyId|tr -d '"')
AWS_SECRET_ACCESS_KEY=$(echo "$creds_json"|jq .Credentials.SecretAccessKey|tr -d '"')
AWS_SESSION_TOKEN=$(echo "$creds_json"|jq .Credentials.SessionToken|tr -d '"')

AWS_DEFAULT_REGION='ap-northeast-1'
SERVICE="execute-api"

curl \
  -X POST \
  -H 'Content-Type: application/json' \
  -H "X-Amz-Security-Token: ${AWS_SESSION_TOKEN}" \
  --aws-sigv4 "aws:amz:${AWS_DEFAULT_REGION}:${SERVICE}" \
  --user "${AWS_ACCESS_KEY_ID}:${AWS_SECRET_ACCESS_KEY}" \
  -d "{\"source\": \"$*\"}" \
  hoge-endpoint|
  jq|
  bat --language json

{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE NamedFieldPuns    #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE OverloadedStrings #-}
module AWS.ApiGatewayExecute.Sign (awsSignToHeader) where

import           Amazonka                               as AWS
import           Amazonka.APIGatewayManagementAPI.Types as AWS.ApiGwMaApi
import qualified Network.HTTP.Client                    as N
import qualified Network.HTTP.Types.Method              as N
import           RIO
import qualified RIO.ByteString                         as B
import           Servant.Client

-- | 純粋空間で引数から認証済みの`Request`を生成する。
awsSignToHeader
  :: ToBody body
  => AuthEnv -> UTCTime -> BaseUrl -> Region -> body -> N.Request -> N.Request
awsSignToHeader authEnv signingTime baseUrl region body req =
  let awsRequest = httpRequestToAwsRequest baseUrl region body req
      Signed _meta signedReq = requestSign awsRequest authEnv region signingTime
  in signedReq

-- | `Network.HTTP.Client.Request`のデータを変換して`Amazonka.Request`にして`requestSign`で署名する準備を整える。
-- `Network.HTTP.Client.Request`が`Network.HTTP.RequestBody`を含んでいるのに、
-- `body`を引数で別に取っているのは重複しているように見えるが、
-- `Network.HTTP.RequestBody`はバリアントによって純粋空間で取得できるとは限らないため、
-- 呼び出し側でどうにか出来る余地を残す。
httpRequestToAwsRequest
  :: ToBody body
  => BaseUrl -> Region -> body -> N.Request -> AWS.Request a
httpRequestToAwsRequest baseUrl region body req =
  AWS.Request
  { _requestService = baseUrlExecuteApiService baseUrl region
  , _requestMethod = fromRight (error "Could not parse request method.") $ N.parseMethod $ N.method req
  , _requestPath = rawPath $ N.path req
  , _requestQuery = parseQueryString $ fromMaybe "" $ B.stripPrefix "?" $ N.queryString req
  , _requestHeaders = N.requestHeaders req
  , _requestBody = toBody body
  }

-- | `BaseUrl`からある程度導出出来るカスタムドメインでのサービス定義。
baseUrlExecuteApiService :: BaseUrl -> Region -> Service
baseUrlExecuteApiService BaseUrl{baseUrlScheme, baseUrlHost, baseUrlPort} =
  customDomainExecuteApiService (fromString baseUrlHost) (baseUrlScheme == Https) baseUrlPort

-- | カスタムドメインでのAPI Gatewayのサービス定義。
customDomainExecuteApiService :: ByteString -> Bool -> Int -> Region -> Service
customDomainExecuteApiService customEndpointHost customEndpointSecure customEndpointPort region =
  executeApiService $ const $ Endpoint
  { _endpointHost = customEndpointHost
  , _endpointSecure = customEndpointSecure
  , _endpointPort = customEndpointPort
  , _endpointScope = toBS region
  }

-- | API GatewayのIAM認証突破を行うためのサービス定義。
-- API Gatewayではカスタムドメインが定義できるため、`Endpoint`は引数で定義出来るようにする。
executeApiService :: (Region -> Endpoint) -> Service
executeApiService endpoint =
  Service
  { _serviceAbbrev = "APIGatewayExecuteAPI"
  , _serviceSigner = _serviceSigner AWS.ApiGwMaApi.defaultService
  , _serviceEndpointPrefix = "execute-api"
  , _serviceSigningName = "execute-api"
  , _serviceVersion = _serviceVersion AWS.ApiGwMaApi.defaultService
  , _serviceEndpoint = endpoint
  , _serviceTimeout = Just 70
  , _serviceCheck = statusSuccess
  , _serviceError = parseJSONError "APIGatewayExecuteAPI"
  , _serviceRetry = _serviceRetry AWS.ApiGwMaApi.defaultService
  }

-- | AWSへ認証ありでリクエストを送る。
toRIOAws :: (HasLogFunc env, HasManager env, HasAwsEnv env) => BaseUrl -> ClientM a -> RIO env a
toRIOAws baseUrl clientM = do
  awsEnv <- view awsEnvL
  let auth = runIdentity $ envAuth awsEnv
  signingTime <- getCurrentTime
  let region = envRegion awsEnv
  withAuth auth $ \authEnv -> do
    let makeClientRequest _baseUrl servantRequest =
          let body = case S.requestBody servantRequest of
                Nothing                            -> ""
                Just (requestBodyMain, _mediaType) -> getRequestBody requestBodyMain
              httpRequest = defaultMakeClientRequest baseUrl servantRequest
          in awsSignToHeader authEnv signingTime baseUrl region body httpRequest
    manager' <- view managerL
    let policy = fullJitterBackoffOfCustomPolicy
        clientEnv = ClientEnv
          { manager = manager'
          , baseUrl
          , cookieJar = Nothing
          , makeClientRequest
          }
    res <- retryingForNetwork policy $ const $ liftIO $ runClientM clientM clientEnv
    either logErrorAndThrowM return res

-- | `RequestBody`を`ByteString`にする。
-- `makeClientRequest`内部で動かないといけないので、純粋関数にする必要がある。
-- よって`IO`が扱えない。
-- `unsafePerformIO`でごまかしている。
getRequestBody :: S.RequestBody -> ByteString
getRequestBody (S.RequestBodyLBS x)    = convert x
getRequestBody (S.RequestBodyBS  x)    = x
-- この`unsafePerformIO`めちゃくちゃ消したいし、ちゃんと動くかどうかの検証もしてない。
-- ただ部分関数にして単にクラッシュするのと、どちらがマシかは微妙なラインだと思うので、とりあえずやっつけ実装を書いている。
getRequestBody (S.RequestBodySource x) = unsafePerformIO $ do
  runSimpleApp $ logWarn "call getRequestBody/unsafePerformIO"
  e <- runExceptT $ S.runSourceT x
  case e of
    Left  l -> error l
    Right r -> return $ convert $ mconcat r