remote_shell.rb 4.79 KB
Newer Older
1
require 'base64'
2
require 'json'
3
require 'socket'
4
require 'timeout'
5

anonym's avatar
anonym committed
6
module RemoteShell
anonym's avatar
anonym committed
7
  class ServerFailure < StandardError
8
9
  end

10
11
12
13
  # Used to differentiate vs Timeout::Error, which is thrown by
  # try_for() (by default) and often wraps around remote shell usage
  # -- in that case we don't want to catch that "outer" exception in
  # our handling of remote shell timeouts below.
anonym's avatar
anonym committed
14
  class Timeout < ServerFailure
15
16
  end

anonym's avatar
anonym committed
17
18
  DEFAULT_TIMEOUT = 20*60

anonym's avatar
anonym committed
19
20
  # Counter providing unique id:s for each communicate() call.
  @@request_id ||= 0
21

anonym's avatar
anonym committed
22
  def communicate(vm, *args, **opts)
23
    opts[:timeout] ||= DEFAULT_TIMEOUT
24
    socket = TCPSocket.new("127.0.0.1", vm.get_remote_shell_port)
anonym's avatar
anonym committed
25
    id = (@@request_id += 1)
26
27
28
29
30
31
32
33
    # Since we already have defined our own Timeout in the current
    # scope, we have to be more careful when referring to the Timeout
    # class from the 'timeout' module. However, note that we want it
    # to throw our own Timeout exception.
    Object::Timeout.timeout(opts[:timeout], Timeout) do
      socket.puts(JSON.dump([id] + args))
      socket.flush
      loop do
34
        line = socket.readline("\n").chomp("\n")
35
36
37
38
39
        response_id, status, *rest = JSON.load(line)
        if response_id == id
          if status != "success"
            if status == "error" and rest.class == Array and rest.size == 1
              msg = rest.first
anonym's avatar
anonym committed
40
              raise ServerFailure.new("#{msg}")
41
            else
anonym's avatar
anonym committed
42
              raise ServerFailure.new("Uncaught exception: #{status}: #{rest}")
43
            end
44
          end
45
46
47
48
          return rest
        else
          debug_log("Dropped out-of-order remote shell response: " +
                    "got id #{response_id} but expected id #{id}")
49
        end
50
      end
51
    end
52
53
  ensure
    socket.close if defined?(socket) && socket
54
55
  end

anonym's avatar
anonym committed
56
57
  module_function :communicate
  private :communicate
58

59
  class ShellCommand
anonym's avatar
anonym committed
60
61
62
63
64
    # If `:spawn` is false the server will block until it has finished
    # executing `cmd`. If it's true the server won't block, and the
    # response will always be [0, "", ""] (only used as an
    # ACK). execute() will always block until a response is received,
    # though. Spawning is useful when starting processes in the
65
66
    # background (or running scripts that does the same) or any
    # application we want to interact with.
anonym's avatar
anonym committed
67
68
    def self.execute(vm, cmd, **opts)
      opts[:user] ||= "root"
69
      opts[:spawn] = false unless opts.has_key?(:spawn)
anonym's avatar
anonym committed
70
71
      type = opts[:spawn] ? "spawn" : "call"
      debug_log("#{type}ing as #{opts[:user]}: #{cmd}")
72
      ret = RemoteShell.communicate(vm, 'sh_' + type, opts[:user], cmd, **opts)
anonym's avatar
anonym committed
73
      debug_log("#{type} returned: #{ret}") if not(opts[:spawn])
anonym's avatar
anonym committed
74
75
      return ret
    end
anonym's avatar
anonym committed
76

anonym's avatar
anonym committed
77
    attr_reader :cmd, :returncode, :stdout, :stderr
78

anonym's avatar
anonym committed
79
    def initialize(vm, cmd, **opts)
anonym's avatar
anonym committed
80
      @cmd = cmd
anonym's avatar
anonym committed
81
      @returncode, @stdout, @stderr = self.class.execute(vm, cmd, **opts)
anonym's avatar
anonym committed
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
    end

    def success?
      return @returncode == 0
    end

    def failure?
      return not(success?)
    end

    def to_s
      "Return status: #{@returncode}\n" +
        "STDOUT:\n" +
        @stdout +
        "STDERR:\n" +
        @stderr
    end
  end
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
  class PythonCommand
    def self.execute(vm, code, **opts)
      opts[:user] ||= "root"
      show_code = code.chomp
      if show_code["\n"]
        show_code = "\n" + show_code.lines.map { |l| " "*4 + l.chomp } .join("\n")
      end
      debug_log("executing Python as #{opts[:user]}: #{show_code}")
      ret = RemoteShell.communicate(
        vm, 'python_execute', opts[:user], code, **opts
      )
      debug_log("execution complete")
      return ret
    end

    attr_reader :code, :exception, :stdout, :stderr

    def initialize(vm, code, **opts)
      @code = code
      @exception, @stdout, @stderr = self.class.execute(vm, code, **opts)
    end

    def success?
      return @exception == nil
    end

    def failure?
      return not(success?)
    end

    def to_s
      "Exception: #{@exception}\n" +
        "STDOUT:\n" +
        @stdout +
        "STDERR:\n" +
        @stderr
    end
  end

140
141
142
  # An IO-like object that is more or less equivalent to a File object
  # opened in rw mode.
  class File
anonym's avatar
anonym committed
143
    def self.open(vm, mode, path, *args, **opts)
144
      debug_log("opening file #{path} in '#{mode}' mode")
145
      ret = RemoteShell.communicate(vm, 'file_' + mode, path, *args, **opts)
146
      if ret.size != 1
anonym's avatar
anonym committed
147
        raise ServerFailure.new("expected 1 value but got #{ret.size}")
148
149
150
151
152
153
154
155
156
157
158
159
      end
      debug_log("#{mode} complete")
      return ret.first
    end

    attr_reader :vm, :path

    def initialize(vm, path)
      @vm, @path = vm, path
    end

    def read()
160
      Base64.decode64(self.class.open(@vm, 'read', @path))
161
162
163
    end

    def write(data)
164
      self.class.open(@vm, 'write', @path, Base64.encode64(data))
165
166
167
    end

    def append(data)
168
      self.class.open(@vm, 'append', @path, Base64.encode64(data))
169
170
    end
  end
171
end