From b1401e2e4e518e70a54fcba8294c71c2e39438ca Mon Sep 17 00:00:00 2001
From: Adam Cammack <adam_cammack@rapid7.com>
Date: Fri, 17 Aug 2018 14:10:28 -0500
Subject: [PATCH] Update the prompt every shell tick

This was inadvertently done as part of the tab completion
initialization but is what we want. Also move the prompt formatting to
the prompt update and make the code more readable.
---
 lib/rex/ui/text/shell.rb | 99 ++++++++++++++++------------------------
 1 file changed, 40 insertions(+), 59 deletions(-)

diff --git a/lib/rex/ui/text/shell.rb b/lib/rex/ui/text/shell.rb
index b7256f6854..103b092a50 100644
--- a/lib/rex/ui/text/shell.rb
+++ b/lib/rex/ui/text/shell.rb
@@ -67,7 +67,6 @@ module Shell
         self.hist_last_saved = Readline::HISTORY.length
       end
       self.input.output = self.output
-      update_prompt(input.prompt)
     end
   end
 
@@ -88,7 +87,6 @@ module Shell
 
       self.input.output = self.output
     end
-    update_prompt('')
   end
 
   #
@@ -131,62 +129,7 @@ module Shell
         break if self.stop_flag || self.stop_count > 1
 
         init_tab_complete
-
-        if framework
-          if input.prompt.include?("%T")
-            t = Time.now
-            # This %T is the strftime shorthand for %H:%M:%S
-            format = framework.datastore['PromptTimeFormat'] || "%T"
-            t = t.strftime(format)
-            # This %T is the marker in the prompt where we need to place the time
-            input.prompt.gsub!(/%T/, t.to_s)
-          end
-
-          if input.prompt.include?("%H")
-            hostname = ENV['HOSTNAME']
-            if hostname.nil?
-              hostname = `hostname`.split('.')[0]
-            end
-
-            # check if hostname is still nil
-            if hostname.nil?
-              hostname = ENV['COMPUTERNAME']
-            end
-
-            if hostname.nil?
-              hostname = 'unknown'
-            end
-
-            input.prompt.gsub!(/%H/, hostname.chomp)
-          end
-
-          if input.prompt.include?("%U")
-            user = ENV['USER']
-            if user.nil?
-              user = `whoami`
-            end
-
-            # check if username is still nil
-            if user.nil?
-              user = ENV['USERNAME']
-            end
-
-            if user.nil?
-              user = 'unknown'
-            end
-
-            input.prompt.gsub!(/%U/, user.chomp)
-          end
-
-          input.prompt.gsub!(/%S/, framework.sessions.length.to_s)
-          input.prompt.gsub!(/%J/, framework.jobs.length.to_s)
-          input.prompt.gsub!(/%L/, Rex::Socket.source_address("50.50.50.50"))
-          input.prompt.gsub!(/%D/, ::Dir.getwd)
-          if framework.db.active
-            input.prompt.gsub!(/%W/, framework.db.workspace.name)
-          end
-          self.init_prompt = input.prompt
-        end
+        update_prompt
 
         line = get_input_line
 
@@ -260,7 +203,7 @@ module Shell
 
       # Set the actual prompt to the saved prompt with any substitutions
       # or updates from our output driver, be they color or whatever
-      self.input.prompt = self.output.update_prompt(new_prompt)
+      self.input.prompt = self.output.update_prompt(format_prompt(new_prompt))
       self.prompt_char  = new_prompt_char if (new_prompt_char)
     end
   end
@@ -441,6 +384,44 @@ protected
     update_prompt *old_p, true
   end
 
+  #
+  # Handle prompt substitutions
+  #
+  def format_prompt(str)
+    if framework
+      if str.include?("%T")
+        t = Time.now
+        # This %T is the strftime shorthand for %H:%M:%S
+        format = framework.datastore['PromptTimeFormat'] || "%T"
+        t = t.strftime(format)
+        # This %T is the marker in the prompt where we need to place the time
+        str.gsub!(/%T/, t.to_s)
+      end
+
+      if str.include?("%H")
+        hostname = ENV['HOSTNAME'] || `hostname`.split('.')[0] ||
+          ENV['COMPUTERNAME'] || 'unknown'
+
+        str.gsub!(/%H/, hostname.chomp)
+      end
+
+      if str.include?("%U")
+        user = ENV['USER'] || `whoami` || ENV['USERNAME'] || 'unknown'
+        str.gsub!(/%U/, user.chomp)
+      end
+
+      str.gsub!(/%S/, framework.sessions.length.to_s)
+      str.gsub!(/%J/, framework.jobs.length.to_s)
+      str.gsub!(/%L/, Rex::Socket.source_address("50.50.50.50"))
+      str.gsub!(/%D/, ::Dir.getwd)
+      if framework.db.active
+        str.gsub!(/%W/, framework.db.workspace.name)
+      end
+    end
+
+    str
+  end
+
   attr_writer   :input, :output # :nodoc:
   attr_accessor :stop_flag, :init_prompt, :cont_prompt # :nodoc:
   attr_accessor :prompt # :nodoc: