From 05bf69c604e34300b02406462d4376c86d595470 Mon Sep 17 00:00:00 2001
From: 0rxa <reiarifi@gmail.com>
Date: Thu, 1 Feb 2024 17:11:13 +0100
Subject: [PATCH] feat(aws): implement permanent state (#12018)

Co-authored-by: Rei Arifi <reiarifi@Reis-MacBook-Pro.local>
Co-authored-by: Hysen Ndregjoni <hndregjoni@hotmail.com>
---
 plugins/aws/README.md      |  5 +++++
 plugins/aws/aws.plugin.zsh | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/plugins/aws/README.md b/plugins/aws/README.md
index dbf8dcfeb..0d0773f63 100644
--- a/plugins/aws/README.md
+++ b/plugins/aws/README.md
@@ -47,6 +47,11 @@ plugins=(... aws)
   Some themes might overwrite the value of RPROMPT instead of appending to it, so they'll need to be fixed to
   see the AWS profile/region prompt.
 
+* Set `AWS_PROFILE_STATE_ENABLED=true` in your zshrc file if you want the aws profile to persist between shell sessions.
+  This option might slow down your shell startup time.
+  By default the state file path is `/tmp/.aws_current_profile`. This means that the state won't survive a reboot or otherwise GC.
+  You can control the state file path using the `AWS_STATE_FILE` environment variable.
+
 ## Theme
 
 The plugin creates an `aws_prompt_info` function that you can use in your theme, which displays
diff --git a/plugins/aws/aws.plugin.zsh b/plugins/aws/aws.plugin.zsh
index 7a2e08a7c..071dd1f0b 100644
--- a/plugins/aws/aws.plugin.zsh
+++ b/plugins/aws/aws.plugin.zsh
@@ -6,10 +6,26 @@ function agr() {
   echo $AWS_REGION
 }
 
+# Update state file if enabled
+function _aws_update_state() {
+  if [[ "$AWS_PROFILE_STATE_ENABLED" == true ]]; then
+    test -d $(dirname ${AWS_STATE_FILE}) || exit 1
+    echo "${AWS_PROFILE} ${AWS_REGION}" > "${AWS_STATE_FILE}"
+  fi
+}
+
+function _aws_clear_state() {
+  if [[ "$AWS_PROFILE_STATE_ENABLED" == true ]]; then
+    test -d $(dirname ${AWS_STATE_FILE}) || exit 1
+    echo -n > "${AWS_STATE_FILE}"
+  fi
+}
+
 # AWS profile selection
 function asp() {
   if [[ -z "$1" ]]; then
     unset AWS_DEFAULT_PROFILE AWS_PROFILE AWS_EB_PROFILE AWS_PROFILE_REGION
+    _aws_clear_state
     echo AWS profile cleared.
     return
   fi
@@ -28,6 +44,8 @@ function asp() {
 
   export AWS_PROFILE_REGION=$(aws configure get region)
 
+  _aws_update_state
+
   if [[ "$2" == "login" ]]; then
     if [[ -n "$3" ]]; then
       aws sso login --sso-session $3
@@ -43,6 +61,7 @@ function asp() {
 function asr() {
   if [[ -z "$1" ]]; then
     unset AWS_DEFAULT_REGION AWS_REGION
+    _aws_update_state
     echo AWS region cleared.
     return
   fi
@@ -56,6 +75,7 @@ function asr() {
 
   export AWS_REGION=$1
   export AWS_DEFAULT_REGION=$1
+  _aws_update_state
 }
 
 # AWS profile switch
@@ -255,6 +275,22 @@ if [[ "$SHOW_AWS_PROMPT" != false && "$RPROMPT" != *'$(aws_prompt_info)'* ]]; th
   RPROMPT='$(aws_prompt_info)'"$RPROMPT"
 fi
 
+if [[ "$AWS_PROFILE_STATE_ENABLED" == true ]]; then
+  AWS_STATE_FILE="${AWS_STATE_FILE:-/tmp/.aws_current_profile}"
+  test -s "${AWS_STATE_FILE}" || return
+
+  aws_state=($(cat $AWS_STATE_FILE))
+  
+  export AWS_DEFAULT_PROFILE="${aws_state[1]}"
+  export AWS_PROFILE="$AWS_DEFAULT_PROFILE"
+  export AWS_EB_PROFILE="$AWS_DEFAULT_PROFILE"
+
+  test -z "${aws_state[2]}" && AWS_REGION=$(aws configure get region)
+
+  export AWS_REGION=${AWS_REGION:-$aws_state[2]}
+  export AWS_DEFAULT_REGION="$AWS_REGION"
+fi
+
 # Load awscli completions
 
 # AWS CLI v2 comes with its own autocompletion. Check if that is there, otherwise fall back