Merge pull request #6590 from mal/knife-ssh-prefix

Knife SSH prefix option

3 files changed, 168 insertions, 80 deletions

diff --git a/lib/chef/knife/ssh.rb b/lib/chef/knife/ssh.rb
index f4a025dba3..b6abd67719 100644
--- a/lib/chef/knife/ssh.rb
+++ b/lib/chef/knife/ssh.rb
@@ -46,11 +46,10 @@ class Chef
         :default => nil,
         :proc => lambda { |o| o.to_i }
 
-      option :attribute,
+      option :ssh_attribute,
         :short => "-a ATTR",
         :long => "--attribute ATTR",
-        :description => "The attribute to use for opening the connection - default depends on the context",
-        :proc => Proc.new { |key| Chef::Config[:knife][:ssh_attribute] = key.strip }
+        :description => "The attribute to use for opening the connection - default depends on the context"
 
       option :manual,
         :short => "-m",
@@ -59,6 +58,10 @@ class Chef
         :description => "QUERY is a space separated list of servers",
         :default => false
 
+      option :prefix_attribute,
+        :long => "--prefix-attribute ATTR",
+        :description => "The attribute to use for prefixing the ouput - default depends on the context"
+
       option :ssh_user,
         :short => "-x USERNAME",
         :long => "--ssh-user USERNAME",
@@ -181,27 +184,34 @@ class Chef
         session_from_list(list)
       end
 
-      def get_ssh_attribute(node)
+      def get_prefix_attribute(item)
+        # Order of precedence for prefix
+        # 1) config value (cli or knife config)
+        # 2) nil
+        msg = "Using node attribute '%s' as the prefix: %s"
+        if item["prefix"]
+          Chef::Log.debug(sprintf(msg, config[:prefix_attribute], item["prefix"]))
+          item["prefix"]
+        else
+          nil
+        end
+      end
+
+      def get_ssh_attribute(item)
         # Order of precedence for ssh target
-        # 1) command line attribute
-        # 2) configuration file
-        # 3) cloud attribute
-        # 4) fqdn
-        if node["config"]
-          Chef::Log.debug("Using node attribute '#{config[:attribute]}' as the ssh target: #{node["config"]}")
-          node["config"]
-        elsif Chef::Config[:knife][:ssh_attribute]
-          Chef::Log.debug("Using node attribute #{Chef::Config[:knife][:ssh_attribute]}: #{node["knife_config"]}")
-          node["knife_config"]
-        elsif node["cloud"] &&
-            node["cloud"]["public_hostname"] &&
-            !node["cloud"]["public_hostname"].empty?
-          Chef::Log.debug("Using node attribute 'cloud[:public_hostname]' automatically as the ssh target: #{node["cloud"]["public_hostname"]}")
-          node["cloud"]["public_hostname"]
+        # 1) config value (cli or knife config)
+        # 2) cloud attribute
+        # 3) fqdn
+        msg = "Using node attribute '%s' as the ssh target: %s"
+        if item["target"]
+          Chef::Log.debug(sprintf(msg, config[:ssh_attribute], item["target"]))
+          item["target"]
+        elsif !item.dig("cloud", "public_hostname").to_s.empty?
+          Chef::Log.debug(sprintf(msg, "cloud.public_hostname", item["cloud"]["public_hostname"]))
+          item["cloud"]["public_hostname"]
         else
-          # falling back to default of fqdn
-          Chef::Log.debug("Using node attribute 'fqdn' as the ssh target: #{node["fqdn"]}")
-          node["fqdn"]
+          Chef::Log.debug(sprintf(msg, "fqdn", item["fqdn"]))
+          item["fqdn"]
         end
       end
 
@@ -212,14 +222,12 @@ class Chef
 
         separator = ui.presenter.attribute_field_separator
 
-        # if we've set an attribute to use on the command line
-        if config[:attribute]
-          required_attributes[:config] = config[:attribute].split(separator)
+        if config[:prefix_attribute]
+          required_attributes[:prefix] = config[:prefix_attribute].split(separator)
         end
 
-        # if we've configured an attribute in our config
-        if Chef::Config[:knife][:ssh_attribute]
-          required_attributes[:knife_config] = Chef::Config[:knife][:ssh_attribute].split(separator)
+        if config[:ssh_attribute]
+          required_attributes[:target] = config[:ssh_attribute].split(separator)
         end
 
         @search_count = 0
@@ -232,8 +240,9 @@ class Chef
           # returned node object
           host = get_ssh_attribute(item)
           next if host.nil?
-          ssh_port = item[:cloud].nil? ? nil : item[:cloud][:public_ssh_port]
-          srv = [host, ssh_port]
+          prefix = get_prefix_attribute(item)
+          ssh_port = item.dig("cloud", "public_ssh_port")
+          srv = [host, ssh_port, prefix]
           list.push(srv)
         end
 
@@ -282,7 +291,8 @@ class Chef
 
       def session_from_list(list)
         list.each do |item|
-          host, ssh_port = item
+          host, ssh_port, prefix = item
+          prefix = host unless prefix
           Chef::Log.debug("Adding #{host}")
           session_opts = session_options(host, ssh_port)
           # Handle port overrides for the main connection.
@@ -291,12 +301,14 @@ class Chef
           # Handle connection timeout
           session_opts[:timeout] = Chef::Config[:knife][:ssh_timeout] if Chef::Config[:knife][:ssh_timeout]
           session_opts[:timeout] = config[:ssh_timeout] if config[:ssh_timeout]
+          # Handle session prefix
+          session_opts[:properties] = { prefix: prefix }
           # Create the hostspec.
           hostspec = session_opts[:user] ? "#{session_opts.delete(:user)}@#{host}" : host
           # Connect a new session on the multi.
           session.use(hostspec, session_opts)
 
-          @longest = host.length if host.length > @longest
+          @longest = prefix.length if prefix.length > @longest
         end
 
         session
@@ -342,9 +354,9 @@ class Chef
             chan.exec command do |ch, success|
               raise ArgumentError, "Cannot execute #{command}" unless success
               ch.on_data do |ichannel, data|
-                print_data(ichannel[:host], data)
+                print_data(ichannel.connection[:prefix], data)
                 if data =~ /^knife sudo password: /
-                  print_data(ichannel[:host], "\n")
+                  print_data(ichannel.connection[:prefix], "\n")
                   ichannel.send_data("#{get_password}\n")
                 end
               end
diff --git a/spec/functional/knife/ssh_spec.rb b/spec/functional/knife/ssh_spec.rb
index 9d6fd3ae10..3defbe781f 100644
--- a/spec/functional/knife/ssh_spec.rb
+++ b/spec/functional/knife/ssh_spec.rb
@@ -181,11 +181,11 @@ describe Chef::Knife::Ssh do
 
       it "uses the ssh_attribute" do
         @knife.run
-        expect(@knife.get_ssh_attribute({ "knife_config" => "ec2.public_hostname" })).to eq("ec2.public_hostname")
+        expect(@knife.get_ssh_attribute({ "target" => "ec2.public_hostname" })).to eq("ec2.public_hostname")
       end
     end
 
-    context "when knife[:ssh_attribute] is not provided]" do
+    context "when knife[:ssh_attribute] is not provided" do
       before do
         setup_knife(["*:*", "uptime"])
         Chef::Config[:knife][:ssh_attribute] = nil
@@ -199,22 +199,69 @@ describe Chef::Knife::Ssh do
 
     context "when -a ec2.public_public_hostname is provided" do
       before do
-        setup_knife(["-a ec2.public_hostname", "*:*", "uptime"])
+        setup_knife(["-a", "ec2.public_hostname", "*:*", "uptime"])
         Chef::Config[:knife][:ssh_attribute] = nil
       end
 
       it "should use the value on the command line" do
         @knife.run
-        expect(@knife.config[:attribute]).to eq("ec2.public_hostname")
+        expect(@knife.config[:ssh_attribute]).to eq("ec2.public_hostname")
       end
 
       it "should override what is set in knife.rb" do
         # This is the setting imported from knife.rb
         Chef::Config[:knife][:ssh_attribute] = "fqdn"
         # Then we run knife with the -a flag, which sets the above variable
-        setup_knife(["-a ec2.public_hostname", "*:*", "uptime"])
+        setup_knife(["-a", "ec2.public_hostname", "*:*", "uptime"])
         @knife.run
-        expect(@knife.config[:attribute]).to eq("ec2.public_hostname")
+        expect(@knife.config[:ssh_attribute]).to eq("ec2.public_hostname")
+      end
+    end
+  end
+
+  describe "prefix" do
+    context "when knife[:prefix_attribute] is set" do
+      before do
+        setup_knife(["*:*", "uptime"])
+        Chef::Config[:knife][:prefix_attribute] = "name"
+      end
+
+      it "uses the prefix_attribute" do
+        @knife.run
+        expect(@knife.get_prefix_attribute({ "prefix" => "name" })).to eq("name")
+      end
+    end
+
+    context "when knife[:prefix_attribute] is not provided" do
+      before do
+        setup_knife(["*:*", "uptime"])
+        Chef::Config[:knife][:prefix_attribute] = nil
+      end
+
+      it "falls back to nil" do
+        @knife.run
+        expect(@knife.get_prefix_attribute({})).to eq(nil)
+      end
+    end
+
+    context "when --prefix-attribute ec2.public_public_hostname is provided" do
+      before do
+        setup_knife(["--prefix-attribute", "ec2.public_hostname", "*:*", "uptime"])
+        Chef::Config[:knife][:prefix_attribute] = nil
+      end
+
+      it "should use the value on the command line" do
+        @knife.run
+        expect(@knife.config[:prefix_attribute]).to eq("ec2.public_hostname")
+      end
+
+      it "should override what is set in knife.rb" do
+        # This is the setting imported from knife.rb
+        Chef::Config[:knife][:prefix_attribute] = "fqdn"
+        # Then we run knife with the -b flag, which sets the above variable
+        setup_knife(["--prefix-attribute", "ec2.public_hostname", "*:*", "uptime"])
+        @knife.run
+        expect(@knife.config[:prefix_attribute]).to eq("ec2.public_hostname")
       end
     end
   end
@@ -305,7 +352,7 @@ describe Chef::Knife::Ssh do
     Chef::Config[:chef_server_url] = "http://localhost:9000"
 
     @api.post("/search/node?q=*:*&start=0&rows=1000", 200) do
-      %({"total":1, "start":0, "rows":[{"data": {"fqdn":"the.fqdn", "config": "the_public_hostname", "knife_config": "the_public_hostname" }}]})
+      %({"total":1, "start":0, "rows":[{"data": {"fqdn":"the.fqdn", "target": "the_public_hostname"}}]})
     end
   end
 
diff --git a/spec/unit/knife/ssh_spec.rb b/spec/unit/knife/ssh_spec.rb
index 815995b687..af54115ac7 100644
--- a/spec/unit/knife/ssh_spec.rb
+++ b/spec/unit/knife/ssh_spec.rb
@@ -49,28 +49,36 @@ describe Chef::Knife::Ssh do
 
       def self.should_return_specified_attributes
         it "returns an array of the attributes specified on the command line OR config file, if only one is set" do
-          @node_bar["config"] = "10.0.0.2"
-          @node_foo["config"] = "10.0.0.1"
-          @knife.config[:attribute] = "ipaddress"
+          @node_bar["target"] = "10.0.0.2"
+          @node_foo["target"] = "10.0.0.1"
+          @node_bar["prefix"] = "bar"
+          @node_foo["prefix"] = "foo"
+          @knife.config[:ssh_attribute] = "ipaddress"
+          @knife.config[:prefix_attribute] = "name"
           Chef::Config[:knife][:ssh_attribute] = "ipaddress" # this value will be in the config file
-          expect(@knife).to receive(:session_from_list).with([["10.0.0.1", nil], ["10.0.0.2", nil]])
+          Chef::Config[:knife][:prefix_attribute] = "name" # this value will be in the config file
+          expect(@knife).to receive(:session_from_list).with([["10.0.0.1", nil, "foo"], ["10.0.0.2", nil, "bar"]])
           @knife.configure_session
         end
 
         it "returns an array of the attributes specified on the command line even when a config value is set" do
-          @node_bar["config"] = "10.0.0.2"
-          @node_foo["config"] = "10.0.0.1"
+          @node_bar["target"] = "10.0.0.2"
+          @node_foo["target"] = "10.0.0.1"
+          @node_bar["prefix"] = "bar"
+          @node_foo["prefix"] = "foo"
           Chef::Config[:knife][:ssh_attribute] = "config_file" # this value will be in the config file
-          @knife.config[:attribute] = "ipaddress" # this is the value of the command line via #configure_attribute
-          expect(@knife).to receive(:session_from_list).with([["10.0.0.1", nil], ["10.0.0.2", nil]])
+          Chef::Config[:knife][:prefix_attribute] = "config_file" # this value will be in the config file
+          @knife.config[:ssh_attribute] = "ipaddress" # this is the value of the command line via #configure_attribute
+          @knife.config[:prefix_attribute] = "name" # this is the value of the command line via #configure_attribute
+          expect(@knife).to receive(:session_from_list).with([["10.0.0.1", nil, "foo"], ["10.0.0.2", nil, "bar"]])
           @knife.configure_session
         end
       end
 
-      it "searchs for and returns an array of fqdns" do
+      it "searches for and returns an array of fqdns" do
         expect(@knife).to receive(:session_from_list).with([
-          ["foo.example.org", nil],
-          ["bar.example.org", nil],
+          ["foo.example.org", nil, nil],
+          ["bar.example.org", nil, nil],
         ])
         @knife.configure_session
       end
@@ -84,8 +92,8 @@ describe Chef::Knife::Ssh do
         end
         it "returns an array of cloud public hostnames" do
           expect(@knife).to receive(:session_from_list).with([
-            ["ec2-10-0-0-1.compute-1.amazonaws.com", nil],
-            ["ec2-10-0-0-2.compute-1.amazonaws.com", nil],
+            ["ec2-10-0-0-1.compute-1.amazonaws.com", nil, nil],
+            ["ec2-10-0-0-2.compute-1.amazonaws.com", nil, nil],
           ])
           @knife.configure_session
         end
@@ -101,8 +109,8 @@ describe Chef::Knife::Ssh do
 
         it "returns an array of fqdns" do
           expect(@knife).to receive(:session_from_list).with([
-            ["foo.example.org", nil],
-            ["bar.example.org", nil],
+            ["foo.example.org", nil, nil],
+            ["bar.example.org", nil, nil],
           ])
           @knife.configure_session
         end
@@ -144,15 +152,35 @@ describe Chef::Knife::Ssh do
     end
   end
 
+  describe "#get_prefix_attribute" do
+    # Order of precedence for prefix
+    # 1) config value (cli or knife config)
+    # 2) nil
+    before do
+      Chef::Config[:knife][:prefix_attribute] = nil
+      @knife.config[:prefix_attribute] = nil
+      @node_foo["cloud"]["public_hostname"] = "ec2-10-0-0-1.compute-1.amazonaws.com"
+      @node_bar["cloud"]["public_hostname"] = ""
+    end
+
+    it "should return nil by default" do
+      expect(@knife.get_prefix_attribute({})).to eq(nil)
+    end
+
+    it "should favor config over nil" do
+      @node_foo["prefix"] = "config"
+      expect( @knife.get_prefix_attribute(@node_foo)).to eq("config")
+    end
+  end
+
   describe "#get_ssh_attribute" do
     # Order of precedence for ssh target
-    # 1) command line attribute
-    # 2) configuration file
-    # 3) cloud attribute
-    # 4) fqdn
+    # 1) config value (cli or knife config)
+    # 2) cloud attribute
+    # 3) fqdn
     before do
       Chef::Config[:knife][:ssh_attribute] = nil
-      @knife.config[:attribute] = nil
+      @knife.config[:ssh_attribute] = nil
       @node_foo["cloud"]["public_hostname"] = "ec2-10-0-0-1.compute-1.amazonaws.com"
       @node_bar["cloud"]["public_hostname"] = ""
     end
@@ -165,18 +193,9 @@ describe Chef::Knife::Ssh do
       expect(@knife.get_ssh_attribute(@node_foo)).to eq("ec2-10-0-0-1.compute-1.amazonaws.com")
     end
 
-    it "should favor to attribute_from_cli over config file and cloud" do
-      @knife.config[:attribute] = "command_line"
-      Chef::Config[:knife][:ssh_attribute] = "config_file"
-      @node_foo["config"] = "command_line"
-      @node_foo["knife_config"] = "config_file"
-      expect( @knife.get_ssh_attribute(@node_foo)).to eq("command_line")
-    end
-
-    it "should favor config file over cloud and default" do
-      Chef::Config[:knife][:ssh_attribute] = "config_file"
-      @node_foo["knife_config"] = "config_file"
-      expect( @knife.get_ssh_attribute(@node_foo)).to eq("config_file")
+    it "should favor config over cloud and default" do
+      @node_foo["target"] = "config"
+      expect( @knife.get_ssh_attribute(@node_foo)).to eq("config")
     end
 
     it "should return fqdn if cloud.hostname is empty" do
@@ -192,40 +211,50 @@ describe Chef::Knife::Ssh do
     end
 
     it "uses the port from an ssh config file" do
-      @knife.session_from_list([["the.b.org", nil]])
+      @knife.session_from_list([["the.b.org", nil, nil]])
       expect(@knife.session.servers[0].port).to eq(23)
     end
 
     it "uses the port from a cloud attr" do
-      @knife.session_from_list([["the.b.org", 123]])
+      @knife.session_from_list([["the.b.org", 123, nil]])
       expect(@knife.session.servers[0].port).to eq(123)
     end
 
+    it "uses the prefix from list" do
+      @knife.session_from_list([["the.b.org", nil, "b-team"]])
+      expect(@knife.session.servers[0][:prefix]).to eq("b-team")
+    end
+
+    it "defaults to a prefix of host" do
+      @knife.session_from_list([["the.b.org", nil, nil]])
+      expect(@knife.session.servers[0][:prefix]).to eq("the.b.org")
+    end
+
     it "defaults to a timeout of 120 seconds" do
-      @knife.session_from_list([["the.b.org", nil]])
+      @knife.session_from_list([["the.b.org", nil, nil]])
       expect(@knife.session.servers[0].options[:timeout]).to eq(120)
     end
 
     it "uses the timeout from Chef Config" do
       Chef::Config[:knife][:ssh_timeout] = 5
       @knife.config[:ssh_timeout] = nil
-      @knife.session_from_list([["the.b.org", nil]])
+      @knife.session_from_list([["the.b.org", nil, nil]])
       expect(@knife.session.servers[0].options[:timeout]).to eq(5)
     end
 
     it "uses the timeout from knife config" do
       @knife.config[:ssh_timeout] = 6
-      @knife.session_from_list([["the.b.org", nil]])
+      @knife.session_from_list([["the.b.org", nil, nil]])
       expect(@knife.session.servers[0].options[:timeout]).to eq(6)
     end
 
     it "uses the user from an ssh config file" do
-      @knife.session_from_list([["the.b.org", 123]])
+      @knife.session_from_list([["the.b.org", 123, nil]])
       expect(@knife.session.servers[0].user).to eq("locutus")
     end
 
     it "uses keepalive settings from an ssh config file" do
-      @knife.session_from_list([["the.b.org", 123]])
+      @knife.session_from_list([["the.b.org", 123, nil]])
       expect(@knife.session.servers[0].options[:keepalive]).to be true
       expect(@knife.session.servers[0].options[:keepalive_interval]).to eq 60
     end