Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place.

1
# -*- coding: binary -*-
2
require 'rex/text/color'
3

4
module Rex
5
module Ui
6
module Text
7

8
###
9
#
10
# The shell class provides a command-prompt style interface in a
11
# generic fashion.
12
#
13
###
14
module Shell
15

16
  include Rex::Text::Color
17

18
  ###
19
  #
20
  # This module is meant to be mixed into an input medium class instance as a
21
  # means of extending it to display a prompt before each call to gets.
22
  #
23
  ###
24
  module InputShell
25
    attr_accessor :prompt, :output
26

27
    def pgets
28

29
      output.print(prompt)
30
      output.flush
31

32
      output.prompting
33
      buf = gets
34
      output.prompting(false)
35

36
      buf
37
    end
38
  end
39

40
  #
41
  # Initializes a shell that has a prompt and can be interacted with.
42
  #
43
  def initialize(prompt, prompt_char = '>', histfile = nil, framework = nil, name = nil)
44
    # Set the stop flag to false
45
    self.stop_flag      = false
46
    self.disable_output = false
47
    self.stop_count     = 0
48
    self.name = name
49

50
    # Initialize the prompt
51
    self.cont_prompt = ' > '
52
    self.cont_flag = false
53
    self.prompt = prompt
54
    self.prompt_char = prompt_char
55

56
    self.histfile = histfile
57
    self.hist_last_saved = 0
58

59
    # Static prompt variables
60
    self.local_hostname = ENV['HOSTNAME'] || try_exec('hostname')&.split('.')&.first&.rstrip || ENV['COMPUTERNAME']
61
    self.local_username = ENV['USER'] || try_exec('whoami')&.rstrip || ENV['USERNAME']
62

63
    self.framework = framework
64
  end
65

66
  def init_tab_complete
67
    if (self.input and self.input.supports_readline)
68
      # Unless cont_flag because there's no tab complete for continuation lines
69
      self.input = Input::Readline.new(lambda { |str| tab_complete(str) unless cont_flag })
70
      self.input.output = self.output
71
    end
72
  end
73

74
  #
75
  # Initializes the user interface input/output classes.
76
  #
77
  def init_ui(in_input = nil, in_output = nil)
78
    # Initialize the input and output methods
79
    self.input  = in_input
80
    self.output = in_output
81

82
    if (self.input)
83
      # Extend the input medium as an input shell if the input medium
84
      # isn't intrinsicly a shell.
85
      if (self.input.intrinsic_shell? == false)
86
        self.input.extend(InputShell)
87
      end
88

89
      self.input.output = self.output
90
    end
91
  end
92

93
  #
94
  # Resets the user interface handles.
95
  #
96
  def reset_ui
97
    init_ui
98
  end
99

100
  #
101
  # Sets the log source that should be used for logging input and output.
102
  #
103
  def set_log_source(log_source)
104
    self.log_source = log_source
105
  end
106

107
  #
108
  # Unsets the log source so that logging becomes disabled.
109
  #
110
  def unset_log_source
111
    set_log_source(nil)
112
  end
113

114
  #
115
  # Performs tab completion on the supplied string.
116
  #
117
  def tab_complete(str)
118
    return tab_complete_proc(str) if (tab_complete_proc)
119
  end
120

121
  #
122
  # Run the command processing loop.
123
  #
124
  def run(&block)
125
    begin
126
      require 'pry'
127
      # pry history will not be loaded by default when pry is used as a breakpoint like `binding.pry`
128
      Pry.config.history_load = false
129
    rescue LoadError
130
      # Pry is a development dependency, if not available suppressing history_load can be safely ignored.
131
    end
132

133
    with_history_manager_context do
134
      begin
135
        while true
136
          # If the stop flag was set or we've hit EOF, break out
137
          break if self.stop_flag || self.stop_count > 1
138

139
          init_tab_complete
140
          update_prompt
141

142
          line = get_input_line
143

144
          # If you have sessions active, this will give you a shot to exit
145
          # gracefully. If you really are ambitious, 2 eofs will kick this out
146
          if input.eof? || line == nil
147
            self.stop_count += 1
148
            next if self.stop_count > 1
149

150
            if block
151
              block.call('quit')
152
            elsif respond_to?(:run_single)
153
              # PseudoShell does not provide run_single
154
              run_single('quit')
155
            end
156

157
            # If a block was passed in, pass the line to it.  If it returns true,
158
            # break out of the shell loop.
159
          elsif block
160
            break if block.call(line)
161

162
            # Otherwise, call what should be an overridden instance method to
163
            # process the line.
164
          else
165
            run_single(line)
166
            self.stop_count = 0
167
          end
168
        end
169
        # Prevent accidental console quits
170
      rescue ::Interrupt
171
        output.print("Interrupt: use the 'exit' command to quit\n")
172
        retry
173
      end
174
    end
175
  end
176

177
  #
178
  # Stop processing user input.
179
  #
180
  def stop
181
    self.stop_flag = true
182
  end
183

184
  #
185
  # Checks to see if the shell has stopped.
186
  #
187
  def stopped?
188
    self.stop_flag
189
  end
190

191
  #
192
  # Change the input prompt.
193
  #
194
  # prompt - the actual prompt
195
  # new_prompt_char the char to append to the prompt
196
  def update_prompt(new_prompt = self.prompt, new_prompt_char = self.prompt_char)
197
    if (self.input)
198
      p = substitute_colors(new_prompt + ' ' + new_prompt_char + ' ', true)
199

200
      # Save the prompt before any substitutions
201
      self.prompt = new_prompt
202
      self.prompt_char  = new_prompt_char
203

204
      # Set the actual prompt to the saved prompt with any substitutions
205
      # or updates from our output driver, be they color or whatever
206
      self.input.prompt = self.output.update_prompt(format_prompt(p))
207
    end
208
  end
209

210
  #
211
  # Output shortcuts
212
  #
213

214
  #
215
  # Prints an error message to the output handle.
216
  #
217
  def print_error(msg='')
218
    return if (output.nil?)
219
    return if (msg.nil?)
220

221
    self.on_print_proc.call(msg) if self.on_print_proc
222
    # Errors are not subject to disabled output
223
    log_output(output.print_error(msg))
224
  end
225

226
  alias_method :print_bad, :print_error
227

228
  #
229
  # Prints a status message to the output handle.
230
  #
231
  def print_status(msg='')
232
    return if (disable_output == true)
233

234
    self.on_print_proc.call(msg) if self.on_print_proc
235
    log_output(output.print_status(msg))
236
  end
237

238
  #
239
  # Prints a good message to the output handle.
240
  #
241
  def print_good(msg='')
242
    return if (disable_output == true)
243

244
    self.on_print_proc.call(msg) if self.on_print_proc
245
    log_output(output.print_good(msg))
246
  end
247

248
  #
249
  # Prints a line of text to the output handle.
250
  #
251
  def print_line(msg='')
252
    return if (disable_output == true)
253

254
    self.on_print_proc.call(msg) if self.on_print_proc
255
    log_output(output.print_line(msg))
256
  end
257

258
  #
259
  # Prints a warning message to the output handle.
260
  #
261
  def print_warning(msg='')
262
    return if (disable_output == true)
263

264
    self.on_print_proc.call(msg) if self.on_print_proc
265
    log_output(output.print_warning(msg))
266
  end
267

268
  #
269
  # Prints a raw message to the output handle.
270
  #
271
  def print(msg='')
272
    return if (disable_output == true)
273
    self.on_print_proc.call(msg) if self.on_print_proc
274
    log_output(output.print(msg))
275
  end
276

277
  #
278
  # Whether or not output has been disabled.
279
  #
280
  attr_accessor :disable_output
281
  #
282
  # The input handle to read user input from.
283
  #
284
  attr_reader   :input
285
  #
286
  # The output handle to write output to.
287
  #
288
  attr_reader   :output
289

290
  attr_reader   :prompt, :prompt_char
291
  attr_accessor :on_command_proc
292
  attr_accessor :on_print_proc
293
  attr_accessor :framework
294
  attr_accessor :history_manager
295
  attr_accessor :hist_last_saved # the number of history lines when last saved/loaded
296

297
  protected
298

299
  # Executes the yielded block under the context of a new HistoryManager context. The shell's history will be flushed
300
  # to disk when no longer interacting with the shell. If no history manager is available, the history will not be persisted.
301
  def with_history_manager_context
302
    history_manager = self.history_manager || framework&.history_manager
303
    return yield unless history_manager
304

305
    begin
306
      history_manager.with_context(history_file: histfile, name: name) do
307
        self.hist_last_saved = Readline::HISTORY.length
308

309
        yield
310
      end
311
    ensure
312
      history_manager.flush
313
      self.hist_last_saved = Readline::HISTORY.length
314
    end
315
  end
316

317
  def supports_color?
318
    true
319
  end
320

321
  #
322
  # Get a single line of input, following continuation directives as necessary.
323
  #
324
  def get_input_line
325
    line = "\\\n"
326
    prompt_needs_reset = false
327

328
    self.cont_flag = false
329
    while line =~ /(^|[^\\])\\\s*$/
330
      # Strip \ and all the trailing whitespace
331
      line.sub!(/\\\s*/, '')
332

333
      if line.length > 0
334
        # Using update_prompt will overwrite the primary prompt
335
        input.prompt = output.update_prompt(self.cont_prompt)
336
        self.cont_flag = true
337
        prompt_needs_reset = true
338
      end
339

340
      output.input = input
341
      str = input.pgets
342
      if str
343
        line << str
344
      else
345
        line = nil
346
      end
347

348
      output.input = nil
349
      log_output(input.prompt)
350
    end
351
    self.cont_flag = false
352

353
    if prompt_needs_reset
354
      # The continuation prompt was used so reset the prompt
355
      update_prompt
356
    end
357

358
    line
359
  end
360

361
  #
362
  # Parse a line into an array of arguments.
363
  #
364
  def parse_line(line)
365
    log_input(line)
366

367
    line.gsub!(/(\r|\n)/, '')
368

369
    begin
370
      return args = Rex::Parser::Arguments.from_s(line)
371
    rescue ::ArgumentError
372
      print_error("Parse error: #{$!}")
373
    end
374

375
    return []
376
  end
377

378
  #
379
  # Print the prompt, but do not log it.
380
  #
381
  def _print_prompt(prompt)
382
    output.print(prompt)
383
  end
384

385
  #
386
  # Writes the supplied input to the log source if one has been registered.
387
  #
388
  def log_input(buf)
389
    rlog(buf, log_source) if (log_source)
390
  end
391

392
  #
393
  # Writes the supplied output to the log source if one has been registered.
394
  #
395
  def log_output(buf)
396
    rlog(buf, log_source) if (log_source)
397
  end
398

399
  #
400
  # Prompt the user for input if possible. Special edition for use inside commands.
401
  #
402
  def prompt_yesno(query)
403
    p = "#{query} [y/N]"
404
    old_p = [self.prompt, self.prompt_char]
405
    update_prompt p, ' '
406
    /^y/i === get_input_line
407
  ensure
408
    update_prompt *old_p
409
  end
410

411
  #
412
  # Handle prompt substitutions
413
  #
414
  def format_prompt(str)
415
    return str unless framework
416

417
    # find the active session
418
    session = framework.sessions.values.find { |session| session.interacting }
419
    default = 'unknown'
420

421
    formatted = ''
422
    skip_next = false
423
    for prefix, spec in str.split('').each_cons(2) do
424
      if skip_next
425
        skip_next = false
426
        next
427
      end
428

429
      unless prefix == '%'
430
        formatted << prefix
431
        skip_next = false
432
        next
433
      end
434

435
      skip_next = true
436
      if spec == 'T'
437
        if framework.datastore['PromptTimeFormat']
438
          strftime_format = framework.datastore['PromptTimeFormat']
439
        else
440
          strftime_format = ::Time::DATE_FORMATS[:db].to_s
441
        end
442
        formatted << ::Time.now.strftime(strftime_format).to_s
443
      elsif spec == 'W' && framework.db.active
444
        formatted << framework.db.workspace.name
445
      elsif session
446
        sysinfo = session.respond_to?(:sys) ? session.sys.config.sysinfo : nil
447

448
        case spec
449
        when 'A'
450
          formatted << (sysinfo.nil? ? default : sysinfo['Architecture'])
451
        when 'D'
452
          formatted << (session.respond_to?(:fs) ? session.fs.dir.getwd(refresh: false) : default)
453
        when 'd'
454
          formatted << ::Dir.getwd
455
        when 'H'
456
          formatted << (sysinfo.nil? ? default : sysinfo['Computer'])
457
        when 'h'
458
          formatted << (self.local_hostname || default).chomp
459
        when 'I'
460
          formatted << session.tunnel_peer
461
        when 'i'
462
          formatted << session.tunnel_local
463
        when 'M'
464
          formatted << session.session_type
465
        when 'S'
466
          formatted << session.sid.to_s
467
        when 'U'
468
          formatted << (session.respond_to?(:sys) ? session.sys.config.getuid(refresh: false) : default)
469
        when 'u'
470
          formatted << (self.local_username || default).chomp
471
        else
472
          formatted << prefix
473
          skip_next = false
474
        end
475
      else
476
        case spec
477
        when 'H'
478
          formatted << (self.local_hostname || default).chomp
479
        when 'J'
480
          formatted << framework.jobs.length.to_s
481
        when 'U'
482
          formatted << (self.local_username || default).chomp
483
        when 'S'
484
          formatted << framework.sessions.length.to_s
485
        when 'L'
486
          formatted << Rex::Socket.source_address
487
        when 'D'
488
          formatted << ::Dir.getwd
489
        else
490
          formatted << prefix
491
          skip_next = false
492
        end
493
      end
494
    end
495

496
    if str.length > 0 && !skip_next
497
      formatted << str[-1]
498
    end
499

500
    formatted
501
  end
502

503
  attr_writer   :input, :output # :nodoc:
504
  attr_writer   :prompt, :prompt_char # :nodoc:
505
  attr_accessor :stop_flag, :cont_prompt # :nodoc:
506
  attr_accessor :tab_complete_proc # :nodoc:
507
  attr_accessor :histfile # :nodoc:
508
  attr_accessor :log_source, :stop_count # :nodoc:
509
  attr_accessor :local_hostname, :local_username # :nodoc:
510
  attr_reader   :cont_flag # :nodoc:
511
  attr_accessor :name
512
private
513

514
  def try_exec(command)
515
    begin
516
      %x{ #{ command } }
517
    rescue SystemCallError
518
      nil
519
    end
520
  end
521

522
  attr_writer   :cont_flag # :nodoc:
523

524
end
525

526
end end end
527

528