#!/usr/bin/env bash
#
# Adds an S3 event trigger to wpe-log-sweep so logs leave the public bucket
# within seconds of landing instead of waiting up to 15 minutes.
#
#   ./deploy-sweep.sh            # DRY RUN - shows what would change
#   ./deploy-sweep.sh --apply    # make the changes
#   ./deploy-sweep.sh --rollback # restore the original code, drop the trigger
#
# WHAT CHANGES
#   1. Function code -> lambda_function.py (handles S3 events AND the sweep).
#      Original saved as lambda_function.original.py.
#   2. Reserved concurrency -> 50. WP Engine drops ~120 files at once and this
#      account's TOTAL concurrency limit is 400, so an unbounded burst would
#      eat 30% of it and contend with wpe-log-parser. Throttled S3 events are
#      retried automatically by Lambda's async path, so nothing is lost.
#   3. Resource policy -> allow s3.amazonaws.com to invoke, scoped to this
#      bucket and account.
#   4. Bucket notification on bsd-wpe-logs: s3:ObjectCreated:* under logs/.
#      The bucket currently has NO notification config, so nothing is
#      overwritten (verified before writing).
#
# WHAT DOES NOT CHANGE
#   - The bsd-wpe-logs bucket POLICY. It stays public, as instructed.
#   - The wpe-log-sweep-15min EventBridge rule. It stays ENABLED as a backstop.
#     See below.
#
# WHY THE 15-MINUTE RULE STAYS
#   S3 -> Lambda is an ASYNCHRONOUS invoke. Lambda retries a failing event
#   twice and then DISCARDS it. There is no DLQ on this function. So with the
#   schedule removed, a single dropped event strands a log file in a public
#   bucket forever -- unbounded exposure, which is worse than the 15-minute
#   window this is meant to shrink.
#
#   Keeping it makes the change strictly an improvement:
#     typical case   ~seconds   (event trigger does the work)
#     worst case     <=15 min   (backstop catches what events missed)
#   It costs roughly nothing: ~800ms of a 256MB Lambda every 15 minutes,
#   moving 0 objects. Cents per month.
#
#   In the healthy steady state the scheduled run logs "sweep: moved=0". If it
#   is consistently moving files, the event trigger is NOT working -- see
#   VERIFY below.
#
#   To drop the backstop anyway (not recommended):
#     aws events disable-rule --name wpe-log-sweep-15min --region us-east-1
#
# VERIFY AFTER APPLYING
#   Next WP Engine drop is ~00:17 UTC. Then:
#     aws logs tail /aws/lambda/wpe-log-sweep --region us-east-1 --since 30m
#   Healthy looks like many "event: moved=1" lines around 00:17, and the
#   scheduled "sweep: moved=0" lines afterwards.
#
set -euo pipefail

REGION="us-east-1"
FUNCTION="wpe-log-sweep"
BUCKET="bsd-wpe-logs"
RULE="wpe-log-sweep-15min"
RESERVED=50
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ACCOUNT=$(aws sts get-caller-identity --query Account --output text)

MODE="dry-run"
case "${1:-}" in
  --apply)    MODE="apply" ;;
  --rollback) MODE="rollback" ;;
  "")         MODE="dry-run" ;;
  *) echo "usage: $0 [--apply|--rollback]" >&2; exit 2 ;;
esac

pack() {
  local src="$1" out="$2"
  rm -f "$out"
  ( cd "$(dirname "$src")" && zip -qj "$out" "$(basename "$src")" )
  # Lambda expects the handler file to be named lambda_function.py.
  if [[ "$(basename "$src")" != "lambda_function.py" ]]; then
    local tmp; tmp="$(mktemp -d)"
    cp "$src" "$tmp/lambda_function.py"
    rm -f "$out"
    ( cd "$tmp" && zip -qj "$out" lambda_function.py )
    rm -rf "$tmp"
  fi
}

if [[ "$MODE" == "rollback" ]]; then
  echo "!! rolling back: removing S3 trigger, restoring original sweep code"
  read -r -p "   type YES: " c; [[ "$c" == "YES" ]] || { echo "aborted"; exit 1; }
  echo "[+] removing bucket notification"
  aws s3api put-bucket-notification-configuration --bucket "$BUCKET" --region "$REGION" \
    --notification-configuration '{}'
  echo "[+] restoring original code"
  pack "$HERE/lambda_function.original.py" "$HERE/sweep.zip"
  aws lambda update-function-code --function-name "$FUNCTION" --region "$REGION" \
    --zip-file "fileb://$HERE/sweep.zip" >/dev/null
  aws lambda wait function-updated --function-name "$FUNCTION" --region "$REGION"
  echo "[+] removing reserved concurrency"
  aws lambda delete-function-concurrency --function-name "$FUNCTION" --region "$REGION" || true
  aws events enable-rule --name "$RULE" --region "$REGION"
  rm -f "$HERE/sweep.zip"
  echo "rolled back. the 15-minute sweep is the only mechanism again."
  exit 0
fi

echo "function: $FUNCTION"
echo "bucket:   $BUCKET  (policy NOT touched - stays public, as instructed)"
echo "backstop: $RULE stays ENABLED"
echo
echo "current notification config on $BUCKET:"
CUR=$(aws s3api get-bucket-notification-configuration --bucket "$BUCKET" --region "$REGION" --output json)
if [[ -z "$CUR" || "$CUR" == "{}" || "$CUR" == "null" ]]; then
  echo "  (none - nothing will be overwritten)"
else
  echo "$CUR" | sed 's/^/  /'
  echo
  echo "  !! WARNING: put-bucket-notification-configuration REPLACES the whole"
  echo "     config. Merge the above into the new config before applying."
fi
echo

if [[ "$MODE" == "dry-run" ]]; then
  cat <<EOF
DRY RUN - nothing changed. --apply would:

  1. update-function-code   $FUNCTION  <- lambda_function.py
  2. put-function-concurrency          <- reserved=$RESERVED (limit is 400 account-wide)
  3. add-permission                    <- allow s3.amazonaws.com to invoke
  4. put-bucket-notification-configuration on $BUCKET:
       s3:ObjectCreated:*  prefix=logs/  -> $FUNCTION

  NOT changed: the $BUCKET bucket policy, the $RULE schedule.

Run tests first:  python3 test_sweep.py
Then:             $0 --apply
EOF
  exit 0
fi

# --- apply -----------------------------------------------------------------
read -r -p "apply? this modifies the live $FUNCTION and $BUCKET notifications. type YES: " c
[[ "$c" == "YES" ]] || { echo "aborted"; exit 1; }

echo "[+] 1/4 updating function code"
pack "$HERE/lambda_function.py" "$HERE/sweep.zip"
aws lambda update-function-code --function-name "$FUNCTION" --region "$REGION" \
  --zip-file "fileb://$HERE/sweep.zip" >/dev/null
aws lambda wait function-updated --function-name "$FUNCTION" --region "$REGION"

echo "[+] 2/4 reserving concurrency ($RESERVED)"
aws lambda put-function-concurrency --function-name "$FUNCTION" --region "$REGION" \
  --reserved-concurrent-executions "$RESERVED" >/dev/null

# Permission MUST exist before the notification: S3 validates that it can
# invoke the function at put-notification time and rejects the call otherwise.
echo "[+] 3/4 allowing S3 to invoke the function"
aws lambda add-permission --function-name "$FUNCTION" --region "$REGION" \
  --statement-id "s3-invoke-${BUCKET}" \
  --action lambda:InvokeFunction --principal s3.amazonaws.com \
  --source-arn "arn:aws:s3:::${BUCKET}" --source-account "$ACCOUNT" >/dev/null 2>&1 || \
  echo "    (permission already present)"

echo "[+] 4/4 adding bucket notification (ObjectCreated under logs/)"
aws s3api put-bucket-notification-configuration --bucket "$BUCKET" --region "$REGION" \
  --notification-configuration "$(cat <<EOF
{
  "LambdaFunctionConfigurations": [
    {
      "Id": "move-logs-off-public-bucket-immediately",
      "LambdaFunctionArn": "arn:aws:lambda:${REGION}:${ACCOUNT}:function:${FUNCTION}",
      "Events": ["s3:ObjectCreated:*"],
      "Filter": {"Key": {"FilterRules": [{"Name": "prefix", "Value": "logs/"}]}}
    }
  ]
}
EOF
)"
rm -f "$HERE/sweep.zip"

echo
echo "done. verify at the next drop (~00:17 UTC):"
echo "  aws logs tail /aws/lambda/$FUNCTION --region $REGION --since 30m"
echo "expect: 'event: moved=1' lines at ~00:17, then 'sweep: moved=0' from the backstop."
echo "if the scheduled sweep keeps reporting moved>0, the trigger is not firing."
echo
echo "rollback: $0 --rollback"