core-client-rb/lib/ey-core/cli/ssh.rb at c859beaf4e5611cca6c69b520bb38013de42b232 · engineyard/core-client-rb · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
require 'ey-core/cli/helpers/server_sieve'
require 'ey-core/cli/subcommand'

module Ey
  module Core
    module Cli
      class Ssh < Subcommand
        title "ssh"
        summary "Open an SSH session to the environment's application master"

        option :account,
          short: "c",
          long: "account",
          description: "Name or id of account",
          argument: "account"

        option :environment,
          short: "e",
          long: "environment",
          description: "Name or id of environment",
          argument: "environment"

        option :server,
          short: 's',
          long: "server",
          description: "Specific server to ssh into. Id or amazon id (i-12345)",
          argument: "server"

        option :utilities,
          long: "utilities",
          description: "Run command on the utility servers with the given names. Specify all to run the command on all utility servers.",
          argument: "'all,resque,redis,etc'"

        option :command,
          long: "command",
          description: "Command to run",
          argument: "'command with args'"

        option :shell,
          long: "shell",
          description: "Run command in a shell other than bash",
          argument: "shell"

        option :bind_address,
          long: "bind_address",
          description: "When no command is specified, pass -L to ssh",
          argument: "bind address"

        switch :all,
          long: "all",
          description: "Run command on all servers"

        switch :app_servers,
          long: "app_servers",
          description: "Run command on all application servers"

	switch :app_master,
	  long: "app_master",
	  description: "Run command on app master"

        switch :db_servers,
          long: "db_servers",
          description: "Run command on all database servers"

        switch :db_master,
          long: "db_master",
          description: "Run command on database master"

        switch :db_slaves,
          long: "db_slaves",
          description: "Run command on database slaves"

        switch :tty,
          short: 't',
          long: "tty",
          description: "Allocated a tty for the command"

        def handle
          operator, environment = core_operator_and_environment_for(options)
          abort "Unable to find matching environment".red unless environment

          cmd      = option(:command)
          ssh_opts = []
          ssh_cmd  = ["ssh"]
          exits    = []
          user     = environment.username
          servers  = []

          if option(:command)
            if shell = option(:shell)
              cmd = Escape.shell_command([shell,'-lc',cmd])
            end

            if switch_active?(:tty)
              ssh_opts << "-t"
            elsif cmd.match(/sudo/)
              puts "sudo commands often need a tty to run correctly. Use -t option to spawn a tty.".yellow
            end

            servers += Ey::Core::Cli::Helpers::ServerSieve.filter(
              environment.servers,
              all: switch_active?(:all),
              app_servers: switch_active?(:app_servers),
              app_master: switch_active?(:app_master),
              db_servers: switch_active?(:db_servers),
              db_master: switch_active?(:db_master),
              utilities: option(:utilities)
            )
          else
            if option(:bind_address)
              ssh_opts += ["-L", option(:bind_address)]
            end

            if option(:server)
              servers += [core_server_for(server: option(:server), operator: environment)]
            else
              servers += Ey::Core::Cli::Helpers::ServerSieve.filter(
                environment.servers,
                all: switch_active?(:all),
                app_servers: switch_active?(:app_servers),
                app_master: switch_active?(:app_master),
                db_servers: switch_active?(:db_servers),
                db_master: switch_active?(:db_master),
                utilities: option(:utilities)
              )

              # Reinstate default behavior: If no command is passed and no
              # roles are passed, open a connection to the app master
              if servers.empty?
                servers += (environment.servers.all(role: "app_master") + environment.servers.all(role: "solo")).to_a
              end
            end
          end

          if servers.empty?
            abort "Unable to find any matching servers. Aborting.".red
          end

          servers.uniq!

          servers.each do |server|
            host = server.public_hostname
            name = server.name ? "#{server.role} (#{server.name})" : server.role
            puts "\nConnecting to #{name} #{host}".green
            sshcmd = Escape.shell_command((ssh_cmd + ssh_opts + ["#{user}@#{host}"] + [cmd]).compact)
            puts "Running command: #{sshcmd}".green
            system sshcmd
            exits << $?.exitstatus
          end

          exit exits.detect {|status| status != 0 } || 0
        end
      end
    end
  end
end