gistfile1.rb

#
# Author:: Adam Jacob (<adam@opscode.com>)
# Copyright:: Copyright (c) 2009 Opscode, Inc.
# License:: Apache License, Version 2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

require 'chef/knife'

class Chef
  class Knife
    class Mosh < Knife

      deps do
        require 'chef/search/query'
      end

      attr_writer :password

      banner "knife mosh QUERY COMMAND (options)"

      option :concurrency,
        :short => "-C NUM",
        :long => "--concurrency NUM",
        :description => "The number of concurrent connections",
        :default => nil,
        :proc => lambda { |o| o.to_i }

      option :attribute,
        :short => "-a ATTR",
        :long => "--attribute ATTR",
        :description => "The attribute to use for opening the connection - default is fqdn",
        :default => "fqdn"

      option :manual,
        :short => "-m",
        :long => "--manual-list",
        :boolean => true,
        :description => "QUERY is a space separated list of servers",
        :default => false

      option :ssh_user,
        :short => "-x USERNAME",
        :long => "--ssh-user USERNAME",
        :description => "The ssh username"

      option :ssh_password,
        :short => "-P PASSWORD",
        :long => "--ssh-password PASSWORD",
        :description => "The ssh password"

      option :ssh_port,
        :short => "-p PORT",
        :long => "--ssh-port PORT",
        :description => "The ssh port",
        :default => "22",
        :proc => Proc.new { |key| Chef::Config[:knife][:ssh_port] = key }

      option :identity_file,
        :short => "-i IDENTITY_FILE",
        :long => "--identity-file IDENTITY_FILE",
        :description => "The SSH identity file used for authentication"

      option :no_host_key_verify,
        :long => "--no-host-key-verify",
        :description => "Disable host key verification",
        :boolean => true,
        :default => false


      def configure_session
        @node_list = case config[:manual]
               when true
                 @name_args[0].split(" ")
               when false
                 r = Array.new
                 q = Chef::Search::Query.new
                 @action_nodes = q.search(:node, @name_args[0])[0]
                 @action_nodes.each do |item|
                   i = format_for_display(item)[config[:attribute]]
                   r.push(i) unless i.nil?
                 end
                 r
               end
        (ui.fatal("No nodes returned from search!"); exit 10) if @node_list.length == 0
      end


      def tmux
        ssh_dest = lambda do |server|
          "'mosh #{server}'"
        end

        new_window_cmds = lambda do
          if @node_list.length > 1
            [""] + @node_list[1..-1].map do |server|
              "new-window -a -n '#{server}' #{ssh_dest.call(server)}"
            end
          else
            []
          end.join(" \\; ")
        end

        tmux_name = "'knife mosh #{@name_args[0].gsub(/:/, '=')}'"
        begin
          server = @node_list.first
          cmd = ["tmux new-session -d -s #{tmux_name}",
                 "-n '#{server}'", ssh_dest.call(server),
                 new_window_cmds.call].join(" ")
          Chef::Mixin::Command.run_command(:command => cmd)
          exec("tmux attach-session -t #{tmux_name}")
        rescue Chef::Exceptions::Exec
        end
      end

      def run

        @longest = 0

        configure_session

        case @name_args[1]
        when "tmux"
          tmux
        end

        session.close
      end

    end
  end
end