Utils For Everyday use

#!/bin/bash

function help_function() {
  echo "Usage: aws-sso-login [option...]" >&2
  echo
  echo "   -c, -config, --config                 config aws-sso-login for first time use"
  echo "   -h, -help, --help                     get help about the script"
  echo "   -p, -profiles, --profiles             list all available profiles"
  echo
  exit 1
}

function sso_init() {
  login_profile="
[profile dummy]
sso_start_url = https://d-9067f17bb4.awsapps.com/start
sso_region = us-east-1
sso_account_id = 000000000000
sso_role_name = dummy
region = us-east-2"
  filename=~/.aws/credentials.sso
  rm -f "$filename"
  echo "$login_profile" > "$filename"
  aws sso login --region us-east-1 --profile dummy
}

# Slugify the account name (lowercase, non-alnum → '-', squeeze and trim) and
# rewrite "vega" → "prod" so AWS_PROFILE=prod in tilt and CI keeps working.
function normalize_profile_name() {
  local name
  name=$(echo "$1" \
    | tr '[:upper:]' '[:lower:]' \
    | tr -c 'a-z0-9-' '-' \
    | sed -E 's/-+/-/g; s/^-//; s/-$//')
  if [ "$name" = "vega" ]; then
    name="prod"
  fi
  echo "$name"
}

# sso_config <profile_name> <account_id> <role_name>
function sso_config() {
  local profile_name="$1"
  local account_id="$2"
  local role_name="$3"

  if grep -q "^\[profile ${profile_name}\]\$" ~/.aws/credentials.sso; then
    return 1
  fi
  cat >> ~/.aws/credentials.sso <<EOF

[profile $profile_name]
sso_start_url = https://d-9067f17bb4.awsapps.com/start
sso_region = us-east-1
sso_account_id = $account_id
sso_role_name = $role_name
region = us-east-2
EOF
}

# Pre-flight: confirm the dummy profile exists in the file AWS CLI reads
# (~/.aws/config, which we symlink to credentials.sso). Doing this *before*
# calling `aws sso login` matters because that command streams a verification
# URL to the user — capturing its output to inspect for an error string would
# hide that URL until the command finishes.
function has_dummy_profile() {
  [ -r ~/.aws/config ] && grep -q '^\[profile dummy\]$' ~/.aws/config
}

function sso_login() {
  if ! has_dummy_profile; then
    echo "No SSO config found (missing dummy profile in ~/.aws/config)."
    if (: < /dev/tty) 2>/dev/null; then
      local yn=""
      read -r -p "Run aws-sso-login -c now? [Y/n]: " yn </dev/tty
      case "$yn" in
        ""|y|Y|yes|YES)
          config_script
          return $?
          ;;
      esac
      return 1
    fi
    echo "Run: aws-sso-login -c" >&2
    return 1
  fi
  aws sso login --region us-east-1 --profile dummy
}

function config_script() {
  mkdir -p ~/.aws
  if [ -L ~/.aws/credentials ]; then
    rm ~/.aws/credentials
  fi
  if [ -f ~/.aws/credentials ]; then
    mv ~/.aws/credentials ~/.aws/credentials.login
  fi
  ln -sf credentials.sso ~/.aws/config
  loggedInUser=$(stat -f %Su "/dev/console")
  shellcheck=$(finger "$loggedInUser" | grep 'Shell:*' | cut -f3 -d ":")
  if [[ $shellcheck = *"bash"* ]];
    then
      sed -i '' '/AWS_PROFILE/d' ~/.bash_profile
  elif [[ $shellcheck = *"zsh"* ]];
    then
      sed -i '' '/AWS_PROFILE/d' ~/.zshrc
  fi
  for file in ~/.aws/sso/cache/*; do
  if [ -f "$file" ]; then
    if [[ "$file" != *"botocore"* ]]; then
      rm "$file"
    fi
  fi
  done
  sso_init
  aws_access_token=$(grep -v "botocore" $(ls -1d ~/.aws/sso/cache/*) | grep -o '"accessToken": "[^"]*' | grep -o '[^"]*$')

  # Capture the accounts list up front so a CLI failure surfaces as an
  # explicit error rather than an empty loop body that prints "Setup complete".
  local accounts_output
  if ! accounts_output=$(aws sso list-accounts \
    --profile dummy \
    --access-token "$aws_access_token" \
    --region us-east-1 \
    --query 'accountList[].[accountName,accountId]' \
    --output text 2>&1); then
    echo "Failed to list SSO accounts: $accounts_output" >&2
    return 1
  fi

  while IFS=$'\t' read -r account_name account_id; do
    [ -z "$account_id" ] && continue

    local account_profile
    account_profile=$(normalize_profile_name "$account_name")

    # Disambiguate slug collisions ("Foo Bar" and "Foo-Bar" both → foo-bar).
    # Without this the second account is silently dropped by sso_config.
    if grep -q "^\[profile ${account_profile}\]\$" ~/.aws/credentials.sso; then
      local original_slug="$account_profile"
      account_profile="${account_profile}-${account_id: -4}"
      echo "  warning: profile '$original_slug' already taken, using '$account_profile'" >&2
    fi

    local roles_output
    if ! roles_output=$(aws sso list-account-roles \
      --profile dummy \
      --account-id "$account_id" \
      --access-token "$aws_access_token" \
      --region us-east-1 \
      --query 'roleList[].[roleName]' \
      --output text 2>&1); then
      echo "Failed to list roles for $account_name ($account_id): $roles_output" >&2
      continue
    fi

    # Sorted so the role-selection prompt below shows roles in a stable order.
    local -a roles=()
    while IFS= read -r role; do
      [ -n "$role" ] && roles+=("$role")
    done <<< "$(printf '%s\n' "$roles_output" | sort)"

    if [ ${#roles[@]} -eq 0 ]; then
      echo "No roles available for $account_name ($account_id), skipping"
      continue
    fi

    # Each account maps to exactly one profile. Multi-role accounts prompt
    # the developer to pick which role backs that profile.
    local primary_role="${roles[0]}"
    if [ ${#roles[@]} -gt 1 ]; then
      echo
      echo "Account $account_name ($account_id) has multiple roles."
      echo "Pick which role should back the \"$account_profile\" profile:"
      local i=1
      for role in "${roles[@]}"; do
        echo "  $i) $role"
        i=$((i+1))
      done
      # The outer loop's stdin is bound to a here-string, so prompt against
      # /dev/tty. If there's no controlling terminal (CI, piped run),
      # silently default to the first role.
      if (: < /dev/tty) 2>/dev/null; then
        local selection=""
        while true; do
          read -r -p "Selection [1-${#roles[@]}] (default 1): " selection </dev/tty
          selection="${selection:-1}"
          if [[ "$selection" =~ ^[0-9]+$ ]] && [ "$selection" -ge 1 ] && [ "$selection" -le ${#roles[@]} ]; then
            primary_role="${roles[$((selection-1))]}"
            break
          fi
          echo "  invalid selection: $selection"
        done
      else
        echo "(no tty — defaulting to first role: ${roles[0]})"
        primary_role="${roles[0]}"
      fi
    fi

    if sso_config "$account_profile" "$account_id" "$primary_role"; then
      echo "Profile $account_profile has been added (role: $primary_role)"
    else
      echo "  warning: profile $account_profile already configured, skipping" >&2
    fi
  done <<< "$accounts_output"

  echo
  echo "Setup complete. Configured profiles:"
  grep '^\[profile ' ~/.aws/credentials.sso \
    | sed -E 's/^\[profile (.*)\]$/  - \1/' \
    | grep -v '  - dummy$' \
    | sort
  echo
  echo "Use any of the above with: AWS_PROFILE=<name> aws ..."
}

function enable_sso() {
  if [ ! -f ~/.aws/credentials.login ]; then
    mv ~/.aws/credentials ~/.aws/credentials.login
  fi
  if [ -f ~/.aws/credentials ]; then
    rm ~/.aws/credentials
  fi
  if [ -f ~/.aws/config ]; then
    echo "Error - AWS SSO is already enabled"
    return 1
  fi
  ln -s credentials.sso ~/.aws/config
  echo "AWS SSO has been enabled"
}

function disable_sso() {
  if [ -L ~/.aws/credentials ]; then
    echo "Error - AWS SSO is already disabled"
    return 1
  fi
  if [ ! -f ~/.aws/credentials.login ]; then
    mv ~/.aws/credentials ~/.aws/credentials.login
  fi
  if [ -f ~/.aws/credentials ]; then
    rm ~/.aws/credentials
  fi
  rm ~/.aws/config
  ln -s credentials.login ~/.aws/credentials
  echo "AWS SSO has been disabled"
}

function list_profiles() {
  grep "sso_role_name" ~/.aws/config | grep -v "dummy" | cut -d= -f2
}

function old_profile()
{
  /usr/bin/sed -i .old "s/profile .*-/profile /g" ~/.aws/credentials.sso
}

if ! command -v aws &> /dev/null; then
    echo "aws-cli is not installed. Please install before continuing."
    exit 1
fi

if [ $# -eq 0 ] || [  "$1" == "" ]; then
    sso_login
    exit $?
fi

while [ "$1" != "" ]; do
  case $1 in
  -c | -config | --config)
    config_script
  ;;
  -d | -disable | --disable)
    disable_sso
  ;;
  -e | -enable | --enable)
    enable_sso
  ;;
  -h | -help | --help)
    help_function
  ;;
  -p | -profiles | --profiles)
    list_profiles
  ;;
  -o | -old_profile | --old_profile)
    old_profile
  ;;
  *)
  echo "Error: Unknown option: $1" >&2
    help_function
  exit 1
  ;;
  esac
  shift
done