Add config option to control shell prompt

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "awscredx"
-version = "0.7.1"
+version = "0.7.2"
 authors = ["Alexei Samokvalov <[email protected]>"]
 edition = "2018"
 

src/assume/mod.rs

@@ -40,7 +40,7 @@ fn run_raw(profile: &str, config: &Config) -> Result<(), String> {
         config,
     );
     assumer.assume(profile)?;
-    print_profile(profile);
+    print_profile(profile, config);
     if Utc::now() - state.last_version_check_time > Duration::days(config.check_new_version_interval_days as i64) {
         check_newer_version();
         state.last_version_check_time = Utc::now();
@@ -58,15 +58,21 @@ fn check_newer_version() {
     }
 }
 
-fn print_profile(profile_name: &str) {
+fn print_profile(profile_name: &str, config: &Config) {
     match env::var_os("SHELL") {
         Some(shell) => {
             let file = Path::new(&shell)
                 .file_name().unwrap()
                 .to_str().unwrap();
             match file {
-                "fish" => println!("set -xg AWS_PROFILE {}", profile_name),
-                _ => println!("export AWS_PROFILE={}", profile_name),
+                "fish" => {
+                    print!("set -xg AWS_PROFILE {}; ", profile_name);
+                    println!("set -l __awscredx_modify_prompt {}", config.modify_shell_prompt);
+                },
+                _ => {
+                    print!("export AWS_PROFILE={}; ", profile_name);
+                    println!("__awscredx_modify_prompt={}", config.modify_shell_prompt);
+                },
             }
         }
         None => println!("export AWS_PROFILE={}", profile_name)

src/config.rs

@@ -14,6 +14,7 @@ pub struct Config {
     mfa_command: Option<String>,
     pub profiles: LinkedHashMap<ProfileName, Profile>,
     pub check_new_version_interval_days: u32,
+    pub modify_shell_prompt: bool,
 }
 
 #[derive(Deserialize, Debug)]
@@ -67,6 +68,7 @@ impl Config {
             mfa_command: Option<String>,
             profiles: LinkedHashMap<ProfileName, ProfileValue>,
             check_new_version_interval_days: Option<u32>,
+            modify_shell_prompt: Option<bool>,
         }
 
         let rc: RawConfig = toml::from_str(&content)
@@ -87,6 +89,7 @@ impl Config {
                 ProfileValue::ProfileConfig(profile) => profile,
             })).collect(),
             check_new_version_interval_days: rc.check_new_version_interval_days.unwrap_or(7),
+            modify_shell_prompt: rc.modify_shell_prompt.unwrap_or(true),
         };
         Ok(Some(config))
     }

src/init/config-template.toml

@@ -13,16 +13,21 @@ main_profile = "main"
 # MFA profile name. Default is '<main_profile>-mfa'.
 # mfa_profile = "mfa"
 
-# How often to check for new version.
-check_new_version_interval_days = 7
-
 # MFA serial number can be found in AWS web console -> IAM
 mfa_serial_number = "arn:aws:iam::MAIN_ACCOUNT_ID:mfa/USERNAME"
 
 # Uncomment and adjust the shell command if your MFA token can be provided by a script.
 # Here is an example for yubikey.
 # mfa_command = "ykman oath code | awk '/SOME_PATTERN/ {print $2}'"
 
+# How often to check for new version.
+check_new_version_interval_days = 7
+
+# This will prepend the assumed profile and the expiration time to your shell prompt.
+# You can disable this feature and directly use the shell function "__awscredx_prompt" in your prompt
+# at a more appropriate place, e.g. on a separate line, etc.
+modify_shell_prompt = true
+
 [profiles]
 
 # You can specify profiles by either providing the role ARNs

src/init/script.fish

@@ -7,13 +7,6 @@ if test $status -ne 0
   functions -c fish_prompt _original_fish_prompt
 end
 
-if test -e $HOME/.config/fish/functions/fish_prompt.fish
-  cat $HOME/.config/fish/functions/fish_prompt.fish | grep __awscredx_profile_prompt > /dev/null
-  if test $status -eq 0
-    set _original_prompt_has_aws_profile 1
-  end
-end
-
 function __awscredx_prompt
   @bin@ print-prompt
 end
@@ -24,7 +17,7 @@ function assume
   switch $s
     case 0
       eval "$output"
-      if test -z "$_original_prompt_has_aws_profile"
+      if test "$__awscredx_modify_prompt" = "true"
         function fish_prompt
           set -l old_status $status
 

src/init/script.sh

@@ -16,10 +16,12 @@ function assume {
   case $s in
     0)
       eval $out
-      if [[ $SHELL =~ zsh ]]; then
-        setopt PROMPT_SUBST
+      if [ "$__awscredx_modify_prompt" = "true" ]; then
+        if [[ $SHELL =~ zsh ]]; then
+          setopt PROMPT_SUBST
+        fi
+        PS1='$(__awscredx_prompt) '${_ORIGINAL_PS1:-}
       fi
-      PS1='$(__awscredx_prompt) '${_ORIGINAL_PS1:-}
       ;;
     50)
       "@bin@" init