summaryrefslogtreecommitdiff
path: root/lib/leap_cli/ssh/backend.rb
blob: 80203b61329d34916c50a1c81a68fd4f191aca2e (plain)
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
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
#
# A custome SSHKit backend, derived from the default netssh backend.
# Our custom backend modifies the logging behavior and gracefully captures
# common exceptions.
#

require 'stringio'
require 'timeout'
require 'sshkit'
require 'leap_cli/ssh/formatter'
require 'leap_cli/ssh/scripts'

module SSHKit
  class Command
    #
    # override exit_status in order to be less verbose
    #
    def exit_status=(new_exit_status)
      @finished_at = Time.now
      @exit_status = new_exit_status
      if options[:raise_on_non_zero_exit] && exit_status > 0
        message = ""
        message += "exit status: " + exit_status.to_s + "\n"
        message += "stdout: " + (full_stdout.strip.empty? ? "Nothing written" : full_stdout.strip) + "\n"
        message += "stderr: " + (full_stderr.strip.empty? ? 'Nothing written' : full_stderr.strip) + "\n"
        raise Failed, message
      end
    end
  end
end

module LeapCli
  module SSH
    class Backend < SSHKit::Backend::Netssh

      # since the @pool is a class instance variable, we need to copy
      # the code from the superclass that initializes it. boo
      @pool = SSHKit::Backend::ConnectionPool.new

      # modify to pass itself to the block, instead of relying on instance_exec.
      def run
        Thread.current["sshkit_backend"] = self
        # was: instance_exec(@host, &@block)
        @block.call(self, @host)
      ensure
        Thread.current["sshkit_backend"] = nil
      end

      #
      # like default capture, but gracefully logs failures for us
      # last argument can be an options hash.
      #
      # available options:
      #
      #   :fail_msg    - [nil] if set, log this instead of the default
      #                  fail message.
      #
      #   :raise_error - [nil] if true, then reraise failed command exception.
      #
      #   :log_cmd     - [false] if true, log what the command is that gets run.
      #
      #   :log_output  - [true] if true, log each output from the command as
      #                  it is received.
      #
      #   :log_finish  - [false] if true, log the exit status and time
      #                  to completion
      #
      #   :log_wrap    - [nil] passed to log method as :wrap option.
      #
      def capture(*args)
        extract_options(args)
        initialize_logger(:log_output => false)
        rescue_ssh_errors(*args) do
          return super(*args)
        end
      end

      #
      # like default execute, but log the results as they come in.
      #
      # see capture() for available options
      #
      def stream(*args)
        extract_options(args)
        initialize_logger
        rescue_ssh_errors(*args) do
          execute(*args)
        end
      end

      def log(*args, &block)
        @logger ||= LeapCli.new_logger
        @logger.log(*args, &block)
      end

      # some prewritten servers-side scripts
      def scripts
        @scripts ||= LeapCli::SSH::Scripts.new(self, @host.hostname)
      end

      #
      # sshkit just passes upload! and download! to Net::SCP, but Net::SCP
      # make it impossible to set the file permissions. Here is how the mode
      # is determined, from upload.rb:
      #
      #    mode = channel[:stat] ? channel[:stat].mode & 07777 : channel[:options][:mode]
      #
      # The stat info from the file always overrides the mode you pass in options.
      # However, the channel[:options][:mode] will be applied for pure in-memory
      # uploads. So, if the mode is set, we convert the upload to be a memory
      # upload instead of a file upload.
      #
      # Stupid, but blame Net::SCP.
      #
      def upload!(src, dest, options={})
        if options[:mode]
          super(StringIO.new(File.read(src)), dest, options)
        else
          super(src, dest, options)
        end
      end

      private

      #
      # creates a new logger instance for this specific ssh command.
      # by doing this, each ssh session has its own logger and its own
      # indentation.
      #
      # potentially modifies 'args' array argument.
      #
      def initialize_logger(default_options={})
        @logger ||= LeapCli.new_logger
        @output = LeapCli::SSH::Formatter.new(@logger, @host, default_options.merge(@options))
      end

      def extract_options(args)
        if args.last.is_a? Hash
          @options = args.pop
        else
          @options = {}
        end
      end

      #
      # capture common exceptions
      #
      def rescue_ssh_errors(*args, &block)
        yield
      rescue StandardError => exc
        if exc.is_a?(SSHKit::Command::Failed) || exc.is_a?(SSHKit::Runner::ExecuteError)
          if @options[:raise_error]
            raise LeapCli::SSH::ExecuteError, exc.to_s
          elsif @options[:fail_msg]
            @logger.log(@options[:fail_msg], host: @host.hostname, :color => :red)
          else
            @logger.log(:failed, args.join(' '), host: @host.hostname) do
              @logger.log(exc.to_s.strip, wrap: true)
            end
          end
        elsif exc.is_a?(Timeout::Error) || exc.is_a?(Net::SSH::ConnectionTimeout)
          @logger.log(:failed, args.join(' '), host: @host.hostname) do
            @logger.log("Connection timed out")
          end
          if @options[:raise_error]
            raise LeapCli::SSH::TimeoutError, exc.to_s
          end
        else
          raise
        end
        return nil
      end

      def output
        @output ||= LeapCli::SSH::Formatter.new(@logger, @host)
      end

    end
  end
end