CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!

1
require 'stringio'
2
require 'open3'
3
require 'English'
4
require 'tempfile'
5
require 'fileutils'
6
require 'timeout'
7
require 'shellwords'
8

9
module Acceptance
10
  class ChildProcessError < ::StandardError
11
  end
12

13
  class ChildProcessTimeoutError < ::StandardError
14
  end
15

16
  class ChildProcessRecvError < ::StandardError
17
  end
18

19
  # A wrapper around ::Open3.popen2e - allows creating a process, writing to stdin, and reading the process output
20
  # All of the data is stored for future retrieval/appending to test output
21
  class ChildProcess
22
    def initialize
23
      super
24

25
      @default_timeout = ENV['CI'] ? 120 : 40
26
      @debug = false
27
      @env ||= {}
28
      @cmd ||= []
29
      @options ||= {}
30

31
      @stdin = nil
32
      @stdout_and_stderr = nil
33
      @wait_thread = nil
34

35
      @buffer = StringIO.new
36
      @all_data = StringIO.new
37
    end
38

39
    # @return [String] All data that was read from stdout/stderr of the running process
40
    def all_data
41
      @all_data.string
42
    end
43

44
    # Runs the process
45
    # @return [nil]
46
    def run
47
      self.stdin, self.stdout_and_stderr, self.wait_thread = ::Open3.popen2e(
48
        @env,
49
        *@cmd,
50
        **@options
51
      )
52

53
      stdin.sync = true
54
      stdout_and_stderr.sync = true
55

56
      nil
57
    rescue StandardError => e
58
      warn "popen failure #{e}"
59
      raise
60
    end
61

62
    # @return [String] A line of input
63
    def recvline(timeout: @default_timeout)
64
      recvuntil($INPUT_RECORD_SEPARATOR, timeout: timeout)
65
    end
66

67
    alias readline recvline
68

69
    # @param [String|Regexp] delim
70
    def recvuntil(delim, timeout: @default_timeout, drop_delim: false)
71
      buffer = ''
72
      result = nil
73

74
      with_countdown(timeout) do |countdown|
75
        while alive? && !countdown.elapsed?
76
          data_chunk = recv(timeout: [countdown.remaining_time, 1].min)
77
          if !data_chunk
78
            next
79
          end
80

81
          buffer += data_chunk
82
          has_delimiter = delim.is_a?(Regexp) ? buffer.match?(delim) : buffer.include?(delim)
83
          next unless has_delimiter
84

85
          result, matched_delim, remaining = buffer.partition(delim)
86
          unless drop_delim
87
            result += matched_delim
88
          end
89
          unrecv(remaining)
90
          # Reset the temporary buffer to avoid the `ensure` mechanism unrecv'ing the buffer unintentionally
91
          buffer = ''
92

93
          return result
94
        end
95
      ensure
96
        unrecv(buffer)
97
      end
98

99
      result
100
    rescue ChildProcessTimeoutError
101
      raise ChildProcessRecvError, "Failed #{__method__}: Did not match #{delim.inspect}, process was alive?=#{alive?.inspect}, remaining buffer: #{self.buffer.string[self.buffer.pos..].inspect}"
102
    end
103

104
    # @return [String] Recv until additional reads would cause a block, or eof is reached, or a maximum timeout is reached
105
    def recv_available(timeout: @default_timeout)
106
      result = ''
107
      finished_reading = false
108

109
      with_countdown(timeout) do
110
        until finished_reading do
111
          data_chunk = recv(timeout: 0, wait_readable: false)
112
          if !data_chunk
113
            finished_reading = true
114
            next
115
          end
116

117
          result += data_chunk
118
        end
119
      end
120

121
      result
122
    rescue EOFError, ChildProcessTimeoutError
123
      result
124
    end
125

126
    # @param [String] data The string of bytes to put back onto the buffer; Future buffered reads will return these bytes first
127
    def unrecv(data)
128
      data.bytes.reverse.each { |b| buffer.ungetbyte(b) }
129
    end
130

131
    # @param [Integer] length Reads length bytes from the I/O stream
132
    # @param [Integer] timeout The timeout in seconds
133
    # @param [TrueClass] wait_readable True if blocking, false otherwise
134
    def recv(length = 4096, timeout: @default_timeout, wait_readable: true)
135
      buffer_result = buffer.read(length)
136
      return buffer_result if buffer_result
137

138
      retry_count = 0
139

140
      # Eagerly read, and if we fail - await a response within the given timeout period
141
      result = nil
142
      begin
143
        result = stdout_and_stderr.read_nonblock(length)
144
        unless result.nil?
145
          log("[read] #{result}")
146
          @all_data.write(result)
147
        end
148
      rescue IO::WaitReadable
149
        if wait_readable
150
          IO.select([stdout_and_stderr], nil, nil, timeout)
151
          retry_count += 1
152
          retry if retry_count == 1
153
        end
154
      end
155

156
      result
157
    end
158

159
    # @param [String] data Write the data to the tdin of the running process
160
    def write(data)
161
      log("[write] #{data}")
162
      @all_data.write(data)
163
      stdin.write(data)
164
      stdin.flush
165
    end
166

167
    # @param [String] s Send line of data to the stdin of the running process
168
    def sendline(s)
169
      write("#{s}#{$INPUT_RECORD_SEPARATOR}")
170
    end
171

172
    # @return [TrueClass, FalseClass] True if the running process is alive, false otherwise
173
    def alive?
174
      wait_thread.alive?
175
    end
176

177
    # Interact with the current process, forwarding the current stdin to the console's stdin,
178
    # and writing the console's output to stdout. Doesn't support using PTY/raw mode.
179
    def interact
180
      $stderr.puts
181
      $stderr.puts '[*] Opened interactive mode - enter "!next" to continue, or "!exit" to stop entirely. !pry for an interactive pry'
182
      $stderr.puts
183

184
      without_debugging do
185
        while alive?
186
          ready = IO.select([stdout_and_stderr, $stdin], [], [], 10)
187

188
          next unless ready
189

190
          reads, = ready
191

192
          reads.to_a.each do |read|
193
            case read
194
            when $stdin
195
              input = $stdin.gets
196
              if input.chomp == '!continue'
197
                return
198
              elsif input.chomp == '!exit'
199
                exit
200
              elsif input.chomp == '!pry'
201
                require 'pry-byebug'; binding.pry
202
              end
203

204
              write(input)
205
            when stdout_and_stderr
206
              available_bytes = recv
207
              $stdout.write(available_bytes)
208
              $stdout.flush
209
            end
210
          end
211
        end
212
      end
213
    end
214

215
    def close
216
      begin
217
        Process.kill('KILL', wait_thread.pid) if wait_thread.pid
218
      rescue StandardError => e
219
        warn "error #{e} for #{@cmd}, pid #{wait_thread.pid}"
220
      end
221
      stdin.close if stdin
222
      stdout_and_stderr.close if stdout_and_stderr
223
    end
224

225
    # @return [IO] the stdin for the child process which can be written to
226
    attr_reader :stdin
227
    # @return [IO] the stdout and stderr for the child process which can be read from
228
    attr_reader :stdout_and_stderr
229
    # @return [Process::Waiter] the waiter thread for the current process
230
    attr_reader :wait_thread
231

232
    # @return [String] The cmd that was used to execute the current process
233
    attr_reader :cmd
234

235
    private
236

237
    # @return [StringIO] the buffer for any data which was read from stdout/stderr which was read, but not consumed
238
    attr_reader :buffer
239
    # @return [IO] the stdin of the running process
240
    attr_writer :stdin
241
    # @return [IO] the stdout and stderr of the running process
242
    attr_writer :stdout_and_stderr
243
    # @return [Process::Waiter] The process wait thread which tracks if the process is alive, its pid, return value, etc.
244
    attr_writer :wait_thread
245

246
    # @param [String] s Log to stderr
247
    def log(s)
248
      return unless @debug
249

250
      $stderr.puts s
251
    end
252

253
    def without_debugging
254
      previous_debug_value = @debug
255
      @debug = false
256
      yield
257
    ensure
258
      @debug = previous_debug_value
259
    end
260

261
    # Yields a timer object that can be used to request the remaining time available
262
    def with_countdown(timeout)
263
      countdown = Acceptance::Countdown.new(timeout)
264
      # It is the caller's responsibility to honor the required countdown limits,
265
      # but let's wrap the full operation in an explicit for worse case scenario,
266
      # which may leave object state in a non-determinant state depending on the call
267
      ::Timeout.timeout(timeout * 1.5) do
268
        yield countdown
269
      end
270
      if countdown.elapsed?
271
        raise ChildProcessTimeoutError
272
      end
273
    rescue ::Timeout::Error
274
      raise ChildProcessTimeoutError
275
    end
276
  end
277

278
  # Internally generates a temporary file with Dir::Tmpname instead of a ::Tempfile instance, otherwise windows won't allow the file to be executed
279
  # at the same time as the current Ruby process having an open handle to the temporary file
280
  class TempChildProcessFile
281
    def initialize(basename, extension)
282
      @file_path = Dir::Tmpname.create([basename, extension]) do |_path, _n, _opts, _origdir|
283
        # noop
284
      end
285

286
      ObjectSpace.define_finalizer(self, self.class.finalizer_proc_for(@file_path))
287
    end
288

289
    def path
290
      @file_path
291
    end
292

293
    def to_s
294
      path
295
    end
296

297
    def inspect
298
      "#<#{self.class} #{self.path}>"
299
    end
300

301
    def self.finalizer_proc_for(path)
302
      proc { File.delete(path) if File.exist?(path) }
303
    end
304
  end
305

306
  ###
307
  # Stores the data for a payload, including the options used to generate the payload,
308
  ###
309
  class Payload
310
    attr_reader :name, :execute_cmd, :generate_options, :datastore
311

312
    def initialize(options)
313
      @name = options.fetch(:name)
314
      @execute_cmd = options.fetch(:execute_cmd)
315
      @generate_options = options.fetch(:generate_options)
316
      @datastore = options.fetch(:datastore)
317
      @executable = options.fetch(:executable, false)
318

319
      basename = "#{File.basename(__FILE__)}_#{name}".gsub(/[^a-zA-Z]/, '-')
320
      extension = options.fetch(:extension, '')
321

322
      @file_path = TempChildProcessFile.new(basename, extension)
323
    end
324

325
    # @return [TrueClass, FalseClass] True if the payload needs marked as executable before being executed
326
    def executable?
327
      @executable
328
    end
329

330
    # @return [String] The path to the payload on disk
331
    def path
332
      @file_path.path
333
    end
334

335
    # @return [Integer] The size of the payload on disk. May be 0 when the payload doesn't exist,
336
    #   or a smaller size than expected if the payload is not fully generated by msfconsole yet.
337
    def size
338
      File.size(path)
339
    rescue StandardError => _e
340
      0
341
    end
342

343
    def [](k)
344
      options[k]
345
    end
346

347
    # @return [Array<String>] The command which can be used to execute this payload. For instance ["python3", "/tmp/path.py"]
348
    def execute_command
349
      @execute_cmd.map do |val|
350
        val.gsub('${payload_path}', path)
351
      end
352
    end
353

354
    # @param [Hash] default_global_datastore
355
    # @return [String] The setg commands for setting the global datastore
356
    def setg_commands(default_global_datastore: {})
357
      commands = []
358
      # Ensure the global framework datastore is always clear
359
      commands << "irb -e '(self.respond_to?(:framework) ? framework : self).datastore.user_defined.clear'"
360
      # Call setg
361
      global_datastore = default_global_datastore.merge(@datastore[:global])
362
      global_datastore.each do |key, value|
363
        commands << "setg #{key} #{value}"
364
      end
365
      commands.join("\n")
366
    end
367

368
    # @param [Hash] default_module_datastore
369
    # @return [String] The command which can be used on msfconsole to generate the payload
370
    def generate_command(default_module_datastore: {})
371
      generate_options = @generate_options.map do |key, value|
372
        "#{key} #{value}"
373
      end
374
      "generate -o #{path} #{generate_options.join(' ')} #{datastore_options(default_module_datastore: default_module_datastore)}"
375
    end
376

377
    # @param [Hash] default_module_datastore
378
    # @return [String] The command which can be used on msfconsole to create the listener
379
    def handler_command(default_module_datastore: {})
380
      "to_handler #{datastore_options(default_module_datastore: default_module_datastore)}"
381
    end
382

383
    # @param [Hash] default_module_datastore
384
    # @return [String] The datastore options string
385
    def datastore_options(default_module_datastore: {})
386
      module_datastore = default_module_datastore.merge(@datastore[:module])
387
      module_options = module_datastore.map do |key, value|
388
        "#{key}=#{value}"
389
      end
390

391
      module_options.join(' ')
392
    end
393

394
    # @param [Hash] default_global_datastore
395
    # @param [Hash] default_module_datastore
396
    # @return [String] A human readable representation of the payload configuration object
397
    def as_readable_text(default_global_datastore: {}, default_module_datastore: {})
398
      <<~EOF
399
        ## Payload
400
        use #{name}
401

402
        ## Set global datastore
403
        #{setg_commands(default_global_datastore: default_global_datastore)}
404

405
        ## Generate command
406
        #{generate_command(default_module_datastore: default_module_datastore)}
407

408
        ## Create listener
409
        #{handler_command(default_module_datastore: default_module_datastore)}
410

411
        ## Execute command
412
        #{Shellwords.join(execute_command)}
413
      EOF
414
    end
415
  end
416

417
  class PayloadProcess
418
    # @return [Process::Waiter] the waiter thread for the current process
419
    attr_reader :wait_thread
420

421
    # @return [String] the executed command
422
    attr_reader :cmd
423

424
    # @return [String] the payload path on disk
425
    attr_reader :payload_path
426

427
    # @param [Array<String>] cmd The command which can be used to execute this payload. For instance ["python3", "/tmp/path.py"]
428
    # @param [path] payload_path The payload path on disk
429
    # @param [Hash] opts the opts to pass to the Process#spawn call
430
    def initialize(cmd, payload_path, opts = {})
431
      super()
432

433
      @payload_path = payload_path
434
      @debug = false
435
      @env = {}
436
      @cmd = cmd
437
      @options = opts
438
    end
439

440
    # @return [Process::Waiter] the waiter thread for the payload process
441
    def run
442
      pid = Process.spawn(
443
        @env,
444
        *@cmd,
445
        **@options
446
      )
447
      @wait_thread = Process.detach(pid)
448
      @wait_thread
449
    end
450

451
    def alive?
452
      @wait_thread.alive?
453
    end
454

455
    def close
456
      begin
457
        Process.kill('KILL', wait_thread.pid) if wait_thread.pid
458
      rescue StandardError => e
459
        warn "error #{e} for #{@cmd}, pid #{wait_thread.pid}"
460
      end
461
      [:in, :out, :err].each do |name|
462
        @options[name].close if @options[name]
463
      end
464
      @wait_thread.join
465
    end
466
  end
467

468
  class ConsoleDriver
469
    def initialize
470
      @console = nil
471
      @payload_processes = []
472
      ObjectSpace.define_finalizer(self, self.class.finalizer_proc_for(self))
473
    end
474

475
    # @param [Acceptance::Payload] payload
476
    # @param [Hash] opts
477
    def run_payload(payload, opts)
478
      if payload.executable? && !File.executable?(payload.path)
479
        FileUtils.chmod('+x', payload.path)
480
      end
481

482
      payload_process = PayloadProcess.new(payload.execute_command, payload.path, opts)
483
      payload_process.run
484
      @payload_processes << payload_process
485
      payload_process
486
    end
487

488
    # @return [Acceptance::Console]
489
    def open_console
490
      @console = Console.new
491
      @console.run
492
      @console.recvuntil(Console.prompt, timeout: 120)
493

494
      @console
495
    end
496

497
    def close_payloads
498
      close_processes(@payload_processes)
499
    end
500

501
    def close
502
      close_processes(@payload_processes + [@console])
503
    end
504

505
    def self.finalizer_proc_for(instance)
506
      proc { instance.close }
507
    end
508

509
    private
510

511
    def close_processes(processes)
512
      while (process = processes.pop)
513
        begin
514
          process.close
515
        rescue StandardError => e
516
          $stderr.puts e.to_s
517
        end
518
      end
519
    end
520
  end
521

522
  class Console < ChildProcess
523
    def initialize
524
      super
525

526
      framework_root = Dir.pwd
527
      @debug = true
528
      @env = {
529
        'BUNDLE_GEMFILE' => File.join(framework_root, 'Gemfile'),
530
        'PATH' => "#{framework_root.shellescape}:#{ENV['PATH']}"
531
      }
532
      @cmd = [
533
        'bundle', 'exec', 'ruby', 'msfconsole',
534
        '--no-readline',
535
        # '--logger', 'Stdout',
536
        '--quiet'
537
      ]
538
      @options = {
539
        chdir: framework_root
540
      }
541
    end
542

543
    def self.prompt
544
      /msf6.*>\s+/
545
    end
546

547
    def reset
548
      sendline('sessions -K')
549
      recvuntil(Console.prompt)
550

551
      sendline('jobs -K')
552
      recvuntil(Console.prompt)
553
    ensure
554
      @all_data.reopen('')
555
    end
556
  end
557
end
558

559