assume-menu - dmenu launcher for AWS assume (with account grouping)

assume-menu

#!/usr/bin/env bash
# assume-menu - A dmenu-like launcher for AWS assume
# Works with fuzzel (with icons), wofi, rofi, or dmenu
#
# Two selection modes controlled by ASSUME_GROUP_ACCOUNTS (default: true):
#
# Grouped mode (ASSUME_GROUP_ACCOUNTS=true):
#   SSO profiles matching "{account}/{permset}" are grouped:
#     Page 1: Select account (or standalone profile)
#     Page 2: Select permission set (last-used shown first)
#   Non-SSO profiles (no sso_account_id/granted_sso_account_id) always
#   appear directly on page 1, even if they contain "/".
#
# Flat mode (ASSUME_GROUP_ACCOUNTS=false):
#   Page 1: Select profile (all profiles listed directly)

set -euo pipefail

VERSION="5"

# User config file (created on first run via menu)
CONFIG_FILE="${HOME}/.config/assume-menu/config"

# Source user config if it exists
# shellcheck disable=SC1090
[[ -f "$CONFIG_FILE" ]] && source "$CONFIG_FILE"

# Defaults for settings not yet in config
GROUP_ACCOUNTS="${GROUP_ACCOUNTS:-true}"

# Export GRANTED_ALIAS_CONFIGURED if enabled in config
if [[ "${GRANTED_ALIAS_CONFIGURED:-}" == "true" ]]; then
    export GRANTED_ALIAS_CONFIGURED=true
fi

# User's shell ($SHELL may be unset when launched from a desktop keybind)
USER_SHELL="${SHELL:-$(getent passwd "${USER:-$(id -un)}" | cut -d: -f7)}"

# AWS config file
AWS_CONFIG="${AWS_CONFIG_FILE:-$HOME/.aws/config}"

# History file for frequency sorting (accounts/entries)
HISTORY_FILE="${ASSUME_HISTORY_FILE:-$HOME/.cache/assume-menu-history}"

# Cache file for sorted entries
CACHE_FILE="${ASSUME_CACHE_FILE:-$HOME/.cache/assume-menu-cache}"

# Service history file (stores profile:service pairs)
SERVICE_HISTORY_FILE="${ASSUME_SERVICE_HISTORY:-$HOME/.cache/assume-menu-service-history}"

# Permission set history file (stores last-used permset per account)
PERMSET_HISTORY_FILE="${ASSUME_PERMSET_HISTORY:-$HOME/.cache/assume-menu-permset-history}"

# Ensure cache directory exists
mkdir -p "$(dirname "$HISTORY_FILE")"

if [[ ! -f "$AWS_CONFIG" ]]; then
    notify-send -u critical "assume-menu" "AWS config not found at $AWS_CONFIG" 2>/dev/null || true
    echo "AWS config not found at $AWS_CONFIG" >&2
    exit 1
fi

# Compute cache key from config and history file states
get_cache_key() {
    local config_hash history_hash
    config_hash=$(md5sum "$AWS_CONFIG" 2>/dev/null | cut -d' ' -f1) || config_hash="none"
    if [[ -f "$HISTORY_FILE" ]]; then
        history_hash=$(md5sum "$HISTORY_FILE" 2>/dev/null | cut -d' ' -f1) || history_hash="none"
    else
        history_hash="empty"
    fi
    echo "${config_hash}:${history_hash}:${GROUP_ACCOUNTS}"
}

# Check if cache is valid and return cached entries
get_cached_entries() {
    if [[ ! -f "$CACHE_FILE" ]]; then
        return 1
    fi
    local cached_key current_key
    cached_key=$(head -1 "$CACHE_FILE")
    current_key=$(get_cache_key)
    if [[ "$cached_key" == "$current_key" ]]; then
        tail -n +2 "$CACHE_FILE"
        return 0
    fi
    return 1
}

# Save entries to cache
save_to_cache() {
    local entries="$1"
    {
        get_cache_key
        echo "$entries"
    } > "$CACHE_FILE"
}

# Parse AWS config into tagged profile lines:
#   sso<TAB>profile_name  - SSO profile (has sso_account_id or granted_sso_account_id)
#   std<TAB>profile_name  - non-SSO profile
parse_config() {
    awk '
    /^\[profile / {
        if (name) printf "%s\t%s\n", (sso ? "sso" : "std"), name
        name = substr($0, 10)
        sub(/\]$/, "", name)
        sso = 0
        next
    }
    /^\[default\]/ {
        if (name) printf "%s\t%s\n", (sso ? "sso" : "std"), name
        name = "default"
        sso = 0
        next
    }
    /^\[/ {
        if (name) printf "%s\t%s\n", (sso ? "sso" : "std"), name
        name = ""
        sso = 0
        next
    }
    name && /(granted_)?sso_account_id/ {
        sso = 1
    }
    END {
        if (name) printf "%s\t%s\n", (sso ? "sso" : "std"), name
    }
    ' "$AWS_CONFIG" | sort -t$'\t' -k2
}

# Extract just profile names from parsed config
get_profile_names() {
    cut -f2 <<< "$1"
}

# Derive first-page entries from parsed config
# SSO profiles with "/" -> extract account (everything before last /)
# Everything else (non-SSO, or no "/") -> standalone entry
get_entries_from_profiles() {
    local parsed="$1"
    awk -F'\t' '{
        if ($1 == "sso" && index($2, "/") > 0) {
            profile = $2
            match(profile, /.*\//)
            print substr(profile, 1, RLENGTH - 1)
        } else {
            print $2
        }
    }' <<< "$parsed" | sort -u
}

# Get permission sets for an account from the full profile list
get_permsets() {
    local account="$1" profiles="$2"
    awk -v prefix="${account}/" \
        'index($0, prefix) == 1 { print substr($0, length(prefix) + 1) }' \
        <<< "$profiles"
}

# Get the last-used permission set for an account
get_last_permset() {
    local account="$1"
    [[ -f "$PERMSET_HISTORY_FILE" ]] || return 0
    awk -F'=' -v account="$account" \
        '$1 == account { last = substr($0, length($1) + 2) } END { if (last) print last }' \
        "$PERMSET_HISTORY_FILE"
}

# Record permission set selection for an account
record_permset() {
    echo "${1}=${2}" >> "$PERMSET_HISTORY_FILE"
}

# Sort permsets: last-used first, then alphabetical
sort_permsets() {
    local account="$1" permsets="$2"
    local last
    last=$(get_last_permset "$account") || true
    if [[ -n "${last:-}" ]] && grep -qxF "$last" <<< "$permsets"; then
        echo "$last"
        grep -vxF "$last" <<< "$permsets" | sort
    else
        sort <<< "$permsets"
    fi
}

# Sort entries by frequency (most used first), then recency, then alphabetical
sort_by_frequency() {
    local entries="$1"
    if [[ ! -f "$HISTORY_FILE" ]]; then
        echo "$entries"
        return
    fi

    awk '
    NR == FNR {
        count[$0]++
        last[$0] = NR
        next
    }
    {
        printf "%d\t%d\t%s\n", count[$0]+0, last[$0]+0, $0
    }
    ' "$HISTORY_FILE" - <<< "$entries" | sort -t$'\t' -k1,1rn -k2,2rn -k3,3 | cut -f3
}

# Record an entry selection to history
record_selection() {
    echo "$1" >> "$HISTORY_FILE"
}

# Record a service selection for a profile
record_service_selection() {
    local profile="$1" service="$2"
    echo "${profile}:${service}" >> "$SERVICE_HISTORY_FILE"
}

# Sort services by frequency for a specific profile
sort_services_for_profile() {
    local profile="$1" services="$2"
    if [[ ! -f "$SERVICE_HISTORY_FILE" ]]; then
        echo "$services"
        return
    fi

    awk -v profile="$profile" '
    NR == FNR {
        if (index($0, profile ":") == 1) {
            service = substr($0, length(profile) + 2)
            count[service]++
            last[service] = NR
        }
        next
    }
    {
        printf "%d\t%d\t%s\n", count[$0]+0, last[$0]+0, $0
    }
    ' "$SERVICE_HISTORY_FILE" - <<< "$services" | sort -t$'\t' -k1,1rn -k2,2rn -k3,3 | cut -f3
}

# Format entries with icons for fuzzel (rofi dmenu protocol)
format_with_icons() {
    local icon="${1:-cloud}"
    while IFS= read -r line; do
        printf '%s\0icon\x1f%s\n' "$line" "$icon"
    done
}

# Run menu with given prompt
run_menu() {
    local prompt="$1"
    if [[ -n "${ASSUME_MENU_CMD:-}" ]]; then
        eval "$ASSUME_MENU_CMD"
    elif command -v fuzzel &>/dev/null; then
        format_with_icons "${2:-cloud}" | fuzzel --dmenu --prompt "$prompt" --width 80
    elif command -v wofi &>/dev/null; then
        wofi --dmenu --prompt "$prompt" --insensitive
    elif command -v rofi &>/dev/null; then
        rofi -dmenu -i -p "$prompt"
    elif command -v dmenu &>/dev/null; then
        dmenu -i -p "$prompt"
    else
        echo "No menu program found. Install fuzzel, wofi, rofi, or dmenu." >&2
        exit 1
    fi
}

# --- Configuration ---

# Save a setting to the config file
save_config() {
    local key="$1" value="$2"
    mkdir -p "$(dirname "$CONFIG_FILE")"
    # Remove existing key if present, then append
    if [[ -f "$CONFIG_FILE" ]]; then
        grep -v "^${key}=" "$CONFIG_FILE" > "${CONFIG_FILE}.tmp" || true
        mv "${CONFIG_FILE}.tmp" "$CONFIG_FILE"
    fi
    echo "${key}=\"${value}\"" >> "$CONFIG_FILE"
}

# Prompt user to pick a terminal emulator
# Detects running terminals via ps, then installed ones via command -v.
# Running terminals are shown first with a (detected) hint.
# User can also type a custom terminal name not in the list.
configure_terminal() {
    local known_terminals=(alacritty foot kitty wezterm terminator ghostty gnome-terminal konsole xterm)
    local running=() installed=() menu_entries=()

    # Check what's currently running for this user
    local procs
    procs=$(ps -u "$(id -un)" -o comm= 2>/dev/null | sort -u)

    for t in "${known_terminals[@]}"; do
        if grep -qxF "$t" <<< "$procs"; then
            running+=("$t")
        elif command -v "$t" &>/dev/null; then
            installed+=("$t")
        fi
    done

    # Build menu: detected first, then installed
    for t in "${running[@]}"; do
        menu_entries+=("$t (detected)")
    done
    for t in "${installed[@]}"; do
        menu_entries+=("$t")
    done

    # Always offer a custom option at the end
    menu_entries+=("Custom...")

    local choice
    choice=$(printf '%s\n' "${menu_entries[@]}" | run_menu "terminal > " "utilities-terminal") || exit 0
    [[ -z "$choice" ]] && exit 0

    if [[ "$choice" == "Custom..." ]]; then
        choice=$(echo "Type terminal name and press Enter" | run_menu "terminal > " "utilities-terminal") || exit 0
        [[ -z "$choice" || "$choice" == "Type terminal name and press Enter" ]] && exit 0
    fi

    # Strip the (detected) hint if present
    TERMINAL="${choice% (detected)}"

    save_config "TERMINAL" "$TERMINAL"
}

# Ask whether to set GRANTED_ALIAS_CONFIGURED=true
configure_granted() {
    local choice
    choice=$(printf '%s\n' "No" "Yes" | run_menu "Set GRANTED_ALIAS_CONFIGURED? > " "dialog-question") || exit 0
    if [[ "$choice" == "Yes" ]]; then
        export GRANTED_ALIAS_CONFIGURED=true
        save_config "GRANTED_ALIAS_CONFIGURED" "true"
    else
        save_config "GRANTED_ALIAS_CONFIGURED" "false"
    fi
}

# Gist URL for self-update
GIST_RAW_URL="https://gist.githubusercontent.com/lbr88/3f0808c701fd106c9f03569ac2ff6d76/raw/assume-menu"

# Download a URL to a file (tries curl, then wget; busts GitHub CDN cache)
download() {
    local url="${1}?$(date +%s)" dest="$2"
    if command -v curl &>/dev/null; then
        curl -fsSL "$url" -o "$dest" 2>/dev/null
    elif command -v wget &>/dev/null; then
        wget -qO "$dest" "$url" 2>/dev/null
    else
        return 1
    fi
}

# Marker file: exists when a newer version is available
UPDATE_CHECK_FILE="${HOME}/.cache/assume-menu-update-available"

# Check for updates by comparing checksums (meant to run in background)
check_for_update() {
    local script_path tmp
    script_path="$(realpath "$0")"
    tmp=$(mktemp)
    if download "$GIST_RAW_URL" "$tmp"; then
        if ! cmp -s "$script_path" "$tmp"; then
            touch "$UPDATE_CHECK_FILE"
        else
            rm -f "$UPDATE_CHECK_FILE"
        fi
    fi
    rm -f "$tmp"
}

# Update script from gist
# Pass "quiet" as $1 to suppress notifications (used by auto-update)
update_script() {
    local quiet="${1:-}"
    local script_path
    script_path="$(realpath "$0")"
    local tmp="${script_path}.tmp"

    if ! download "$GIST_RAW_URL" "$tmp"; then
        [[ "$quiet" != "quiet" ]] && notify-send -u critical "assume-menu" "Update failed: could not download (install curl or wget)" 2>/dev/null || true
        rm -f "$tmp"
        return 1
    fi

    # Basic sanity check: must start with a shebang
    if ! head -1 "$tmp" | grep -q '^#!/'; then
        [[ "$quiet" != "quiet" ]] && notify-send -u critical "assume-menu" "Update failed: invalid file" 2>/dev/null || true
        rm -f "$tmp"
        return 1
    fi

    # Skip if already up to date
    if cmp -s "$script_path" "$tmp"; then
        rm -f "$tmp" "$UPDATE_CHECK_FILE"
        return 0
    fi

    chmod +x "$tmp"
    mv "$tmp" "$script_path"
    rm -f "$UPDATE_CHECK_FILE"
    [[ "$quiet" != "quiet" ]] && notify-send "assume-menu" "Updated successfully" 2>/dev/null || true
}

# Show configuration menu with all settings
run_configure() {
    local options
    options=$(printf '%s\n' \
        "Terminal ($TERMINAL)" \
        "Group accounts ($GROUP_ACCOUNTS)" \
        "GRANTED_ALIAS_CONFIGURED ($GRANTED_ALIAS_CONFIGURED)" \
        "Auto-update (${AUTO_UPDATE:-false})" \
        "---  version: $VERSION  ---" \
        "Back")
    local choice
    choice=$(echo "$options" | run_menu "configure > " "preferences-system") || return
    case "$choice" in
        Terminal*) configure_terminal ;;
        Group*) configure_group_accounts ;;
        GRANTED*) configure_granted ;;
        Auto-update*) configure_auto_update ;;
        Back|"") return ;;
    esac
}

# Ask whether to group SSO profiles by account
configure_group_accounts() {
    local choice
    choice=$(printf '%s\n' "Yes" "No" | run_menu "Group SSO profiles by account? > " "dialog-question") || return
    if [[ "$choice" == "Yes" ]]; then
        GROUP_ACCOUNTS="true"
    else
        GROUP_ACCOUNTS="false"
    fi
    save_config "GROUP_ACCOUNTS" "$GROUP_ACCOUNTS"
}

# Ask whether to enable auto-update
configure_auto_update() {
    local choice
    choice=$(printf '%s\n' "No" "Yes" | run_menu "Auto-update on launch? > " "dialog-question") || return
    if [[ "$choice" == "Yes" ]]; then
        AUTO_UPDATE="true"
    else
        AUTO_UPDATE="false"
    fi
    save_config "AUTO_UPDATE" "$AUTO_UPDATE"
}

# Run setup for any unconfigured settings
first_run=false
[[ -z "${TERMINAL:-}" ]] && { configure_terminal; first_run=true; }
[[ -z "${GRANTED_ALIAS_CONFIGURED:-}" ]] && { configure_granted; first_run=true; }
[[ -z "${AUTO_UPDATE:-}" ]] && { configure_auto_update; first_run=true; }

# Handle updates
if [[ "${AUTO_UPDATE:-false}" == "true" ]]; then
    # Auto-update: silently replace in background
    (update_script quiet &>/dev/null &)
elif [[ "$first_run" == "true" ]]; then
    # First run: check synchronously so update shows immediately
    check_for_update &>/dev/null
else
    # Check for updates in background (sets marker for next launch)
    (check_for_update &>/dev/null &)
fi

# Detect if an update is available from a previous background check
update_available=false
if [[ -f "$UPDATE_CHECK_FILE" ]]; then
    update_available=true
    notify-send "assume-menu" "A new version is available" 2>/dev/null || true
fi

# --- Main flow ---

# Parse config once (tagged sso/std lines)
parsed_config=$(parse_config)
all_profiles=$(get_profile_names "$parsed_config")

if [[ -z "$all_profiles" ]]; then
    notify-send -u critical "assume-menu" "No AWS profiles found" 2>/dev/null || true
    echo "No AWS profiles found in $AWS_CONFIG" >&2
    exit 1
fi

# Step 1: Select profile
if [[ "$GROUP_ACCOUNTS" == "true" || "$GROUP_ACCOUNTS" == "1" ]]; then
    # Grouped mode: select account first, then permission set
    if ! entries=$(get_cached_entries); then
        entries=$(get_entries_from_profiles "$parsed_config")
        entries=$(sort_by_frequency "$entries")
        save_to_cache "$entries"
    fi

    if [[ "$update_available" == "true" ]]; then
        menu_list=$(printf '%s\n%s\n%s\n' "Update..." "$entries" "Configure...")
    else
        menu_list=$(printf '%s\n%s\n' "$entries" "Configure...")
    fi
    selected_entry=$(echo "$menu_list" | run_menu "account > " "cloud") || exit 0
    [[ -z "$selected_entry" ]] && exit 0

    if [[ "$selected_entry" == "Configure..." ]]; then
        run_configure
        exit 0
    fi
    if [[ "$selected_entry" == "Update..." ]]; then
        update_script
        exit 0
    fi

    record_selection "$selected_entry"

    # Check if entry has permission sets (is an account) or is standalone
    permsets=$(get_permsets "$selected_entry" "$all_profiles")

    if [[ -n "$permsets" ]]; then
        # Step 1b: Select permission set (last-used shown first)
        sorted_permsets=$(sort_permsets "$selected_entry" "$permsets")
        selected_permset=$(echo "$sorted_permsets" | run_menu "role > " "dialog-password") || exit 0
        [[ -z "$selected_permset" ]] && exit 0

        record_permset "$selected_entry" "$selected_permset"
        selected="${selected_entry}/${selected_permset}"
    else
        # Standalone profile - use entry directly as profile name
        selected="$selected_entry"
    fi
else
    # Flat mode: show all profiles directly
    if ! entries=$(get_cached_entries); then
        entries=$(sort_by_frequency "$all_profiles")
        save_to_cache "$entries"
    fi

    if [[ "$update_available" == "true" ]]; then
        menu_list=$(printf '%s\n%s\n%s\n' "Update..." "$entries" "Configure...")
    else
        menu_list=$(printf '%s\n%s\n' "$entries" "Configure...")
    fi
    selected=$(echo "$menu_list" | run_menu "profile > " "cloud") || exit 0
    [[ -z "$selected" ]] && exit 0

    if [[ "$selected" == "Configure..." ]]; then
        run_configure
        exit 0
    fi
    if [[ "$selected" == "Update..." ]]; then
        update_script
        exit 0
    fi

    record_selection "$selected"
fi

# Step 2: Select action
actions="Console (browser)
Terminal (shell)"

action=$(echo "$actions" | run_menu "action > " "utilities-terminal") || exit 0
[[ -z "$action" ]] && exit 0

shopt -s nocasematch
case "$action" in
    "Console (browser)")
        # Step 3: Select AWS service (names match granted's service_map.go)
        all_services="console
acm
apigateway
appsync
athena
backup
bedrock
billing
cloudformation
cloudfront
cloudmap
cloudwatch
codeartifact
codecommit
codedeploy
codepipeline
codesuite
cognito
config
controltower
directconnect
dms
dynamodb
ebs
ec2
ecr
ecs
eks
elasticache
elasticbeanstalk
eventbridge
globalaccelerator
grafana
iam
kms
lambda
mwaa
organizations
rds
redshift
route53
s3
sagemaker
secretsmanager
securityhub
ses
sns
sqs
ssm
stepfunctions
vpc
wafv2"

        # Sort services by history, but keep "console" always first
        sorted_services=$(sort_services_for_profile "$selected" "$all_services" | grep -v '^console$')
        services="console"$'\n'"$sorted_services"
        service=$(echo "$services" | run_menu "service > " "applications-internet") || exit 0
        [[ -z "$service" ]] && exit 0

        record_service_selection "$selected" "$service"

        if [[ "$service" == "console" ]]; then
            exec assume -c "$selected"
        else
            exec assume -c "$selected" -s "$service"
        fi
        ;;
    "Terminal (shell)")
        # Launch terminal with user's shell, run assume, then spawn interactive shell with creds
        exec "$TERMINAL" -e "$USER_SHELL" -ic "assume '$selected'; $USER_SHELL"
        ;;
esac