Rescue Train::UserError and re-try with su_password

Signed-off-by: Vivek Singh <vivek.singh@msystechnologies.com>

2 files changed, 19 insertions, 21 deletions

diff --git a/lib/chef/knife/bootstrap.rb b/lib/chef/knife/bootstrap.rb
index 13a62b4de0..340ffaecfd 100644
--- a/lib/chef/knife/bootstrap.rb
+++ b/lib/chef/knife/bootstrap.rb
@@ -606,25 +606,26 @@ class Chef
 
       # Actual bootstrap command to be run on the node.
       # Handles recursive calls if su USER failed to authenticate.
-      def bootstrap_run_command(cmd, limit = 3)
-        r = connection.run_command(cmd) do |data, ch|
+      def bootstrap_run_command(cmd)
+        r = connection.run_command(cmd) do |data, channel|
           ui.msg("#{ui.color(" [#{connection.hostname}]", :cyan)} #{data}")
-          ch.send_data("#{config[:su_password] || config[:connection_password]}\n") if data.match?("Password:")
+          channel.send_data("#{config[:su_password] || config[:connection_password]}\n") if data.match?("Password:")
         end
+
         if r.exit_status != 0
-          stderr = (r.stderr + r.stdout).strip
-
-          if stderr.match?("su: Authentication failure")
-            limit -= 1
-            ui.warn("Failed to authenticate su - #{config[:su_user]} to #{server_name}")
-            password = ui.ask("Enter password for su - #{config[:su_user]}@#{server_name}:", echo: false)
-            config[:su_password] = password
-            bootstrap_run_command(cmd, limit) if limit > 0
-          else
-            ui.error("The following error occurred on #{server_name}:")
-            ui.error(stderr)
-            exit(r.exit_status)
-          end
+          ui.error("The following error occurred on #{server_name}:")
+          ui.error("#{r.stdout} #{r.stderr}".strip)
+          exit(r.exit_status)
+        end
+      rescue Train::UserError => e
+        limit ||= 0
+        if e.reason == :bad_su_user_password && limit < 3
+          limit += 1
+          ui.warn("Failed to authenticate su - #{config[:su_user]} to #{server_name}")
+          config[:su_password] = ui.ask("Enter password for su - #{config[:su_user]}@#{server_name}:", echo: false)
+          retry
+        else
+          raise
         end
       end
 
diff --git a/spec/unit/knife/bootstrap_spec.rb b/spec/unit/knife/bootstrap_spec.rb
index ede8d029c6..294942c229 100644
--- a/spec/unit/knife/bootstrap_spec.rb
+++ b/spec/unit/knife/bootstrap_spec.rb
@@ -1763,7 +1763,6 @@ describe Chef::Knife::Bootstrap do
       let(:exit_status) { 1 }
       let(:stdout) { "su: Authentication failure" }
       let(:connection_obj) { double("connection", transport_options: {}) }
-      let(:result_mock2) { double("result", exit_status: 1, stderr: "A message", stdout: "") }
       it "shows an error and exits" do
         allow(connection).to receive(:connection).and_return(connection_obj)
         expect(knife.ui).to receive(:info).with(/Bootstrapping.*/)
@@ -1774,10 +1773,8 @@ describe Chef::Knife::Bootstrap do
           .to receive(:run_command)
           .with("su - USER -c 'sh /path.sh'")
           .and_yield("output here", nil)
-          .and_return result_mock
-        expect(knife.ui).to receive(:ask).and_return("password").twice
-        expect(connection).to receive(:run_command).with("su - USER -c 'sh /path.sh'").and_return(result_mock, result_mock2)
-        expect { knife.perform_bootstrap("/path.sh") }.to raise_error(SystemExit)
+          .and_raise(Train::UserError)
+        expect { knife.perform_bootstrap("/path.sh") }.to raise_error(Train::UserError)
       end
     end
   end