rex 1.0.2 → 2.0.2

data/lib/rex/ui/text/shell.rb

@@ -0,0 +1,403 @@
+# -*- coding: binary -*-
+require 'rex/ui'
+
+module Rex
+module Ui
+module Text
+
+###
+#
+# The shell class provides a command-prompt style interface in a
+# generic fashion.
+#
+###
+module Shell
+
+  ###
+  #
+  # This module is meant to be mixed into an input medium class instance as a
+  # means of extending it to display a prompt before each call to gets.
+  #
+  ###
+  module InputShell
+    attr_accessor :prompt, :output
+
+    def pgets()
+
+      output.print(prompt)
+      output.flush
+
+      output.prompting
+      buf = gets
+      output.prompting(false)
+
+      buf
+    end
+  end
+
+  #
+  # Initializes a shell that has a prompt and can be interacted with.
+  #
+  def initialize(prompt, prompt_char = '>', histfile = nil, framework = nil)
+    # Set the stop flag to false
+    self.stop_flag      = false
+    self.disable_output = false
+    self.stop_count	    = 0
+
+    # Initialize the prompt
+    self.init_prompt = prompt
+    self.prompt_char = prompt_char
+
+    self.histfile = histfile
+    self.hist_last_saved = 0
+
+    self.framework = framework
+  end
+
+  def init_tab_complete
+    if (self.input and self.input.supports_readline)
+      self.input = Input::Readline.new(lambda { |str| tab_complete(str) })
+      if Readline::HISTORY.length == 0 and histfile and File.exists?(histfile)
+        File.readlines(histfile).each { |e|
+          Readline::HISTORY << e.chomp
+        }
+        self.hist_last_saved = Readline::HISTORY.length
+      end
+      self.input.output = self.output
+      update_prompt(input.prompt)
+    end
+  end
+
+  #
+  # Initializes the user interface input/output classes.
+  #
+  def init_ui(in_input = nil, in_output = nil)
+    # Initialize the input and output methods
+    self.input  = in_input
+    self.output = in_output
+
+    if (self.input)
+      # Extend the input medium as an input shell if the input medium
+      # isn't intrinsicly a shell.
+      if (self.input.intrinsic_shell? == false)
+        self.input.extend(InputShell)
+      end
+
+      self.input.output = self.output
+    end
+    update_prompt('')
+  end
+
+  #
+  # Resets the user interface handles.
+  #
+  def reset_ui
+    init_ui
+  end
+
+  #
+  # Sets the log source that should be used for logging input and output.
+  #
+  def set_log_source(log_source)
+    self.log_source = log_source
+  end
+
+  #
+  # Unsets the log source so that logging becomes disabled.
+  #
+  def unset_log_source
+    set_log_source(nil)
+  end
+
+  #
+  # Performs tab completion on the supplied string.
+  #
+  def tab_complete(str)
+    return tab_complete_proc(str) if (tab_complete_proc)
+  end
+
+  #
+  # Run the command processing loop.
+  #
+  def run(&block)
+
+    begin
+
+      while true
+        # If the stop flag was set or we've hit EOF, break out
+        break if (self.stop_flag or self.stop_count > 1)
+
+        init_tab_complete
+
+        if framework
+          if input.prompt.include?("%T")
+            t = Time.now
+            if framework.datastore['PromptTimeFormat']
+              t = t.strftime(framework.datastore['PromptTimeFormat'])
+            end
+            input.prompt.gsub!(/%T/, t.to_s)
+          end
+
+          if input.prompt.include?("%H")
+            hostname = ENV['HOSTNAME']
+            if hostname.nil?
+              hostname = `hostname`.split('.')[0]
+            end
+
+            # check if hostname is still nil
+            if hostname.nil?
+              hostname = ENV['COMPUTERNAME']
+            end
+
+            if hostname.nil?
+              hostname = 'unknown'
+            end
+
+            input.prompt.gsub!(/%H/, hostname.chomp)
+          end
+
+          if input.prompt.include?("%U")
+            user = ENV['USER']
+            if user.nil?
+              user = `whoami`
+            end
+
+            # check if username is still nil
+            if user.nil?
+              user = ENV['USERNAME']
+            end
+
+            if user.nil?
+              user = 'unknown'
+            end
+
+            input.prompt.gsub!(/%U/, user.chomp)
+          end
+
+          input.prompt.gsub!(/%S/, framework.sessions.length.to_s)
+          input.prompt.gsub!(/%J/, framework.jobs.length.to_s)
+          input.prompt.gsub!(/%L/, Rex::Socket.source_address("50.50.50.50"))
+          input.prompt.gsub!(/%D/, ::Dir.getwd)
+          self.init_prompt = input.prompt
+        end
+
+        line = input.pgets()
+        log_output(input.prompt)
+
+        # If a block was passed in, pass the line to it.  If it returns true,
+        # break out of the shell loop.
+        if (block)
+          break if (line == nil or block.call(line))
+        elsif(input.eof? or line == nil)
+        # If you have sessions active, this will give you a shot to exit gravefully
+        # If you really are ambitious, 2 eofs will kick this out
+          self.stop_count += 1
+          next if(self.stop_count > 1)
+          run_single("quit")
+        else
+        # Otherwise, call what should be an overriden instance method to
+        # process the line.
+          ret = run_single(line)
+          # don't bother saving lines that couldn't be found as a
+          # command, create the file if it doesn't exist
+          if ret and self.histfile
+            File.open(self.histfile, "a+") { |f|
+              f.puts(line)
+            }
+          end
+          self.stop_count = 0
+        end
+
+      end
+    # Prevent accidental console quits
+    rescue ::Interrupt
+      output.print("Interrupt: use the 'exit' command to quit\n")
+      retry
+    end
+  end
+
+  #
+  # Stop processing user input.
+  #
+  def stop
+    self.stop_flag = true
+  end
+
+  #
+  # Checks to see if the shell has stopped.
+  #
+  def stopped?
+    self.stop_flag
+  end
+
+  #
+  # Change the input prompt.
+  #
+  # prompt - the actual prompt
+  # new_prompt_char the char to append to the prompt
+  # mode - append or not to append - false = append true = make a new prompt
+  def update_prompt(prompt = nil, new_prompt_char = nil, mode = false)
+    if (self.input)
+      if prompt
+        new_prompt = self.init_prompt + ' ' + prompt + prompt_char + ' '
+      else
+        new_prompt = self.prompt || ''
+      end
+
+      if mode
+        new_prompt = prompt + (new_prompt_char || prompt_char) + ' '
+      end
+
+      # Save the prompt before any substitutions
+      self.prompt = new_prompt
+
+      # Set the actual prompt to the saved prompt with any substitutions
+      # or updates from our output driver, be they color or whatever
+      self.input.prompt = self.output.update_prompt(new_prompt)
+      self.prompt_char  = new_prompt_char if (new_prompt_char)
+    end
+  end
+
+  #
+  # Output shortcuts
+  #
+
+  #
+  # Prints an error message to the output handle.
+  #
+  def print_error(msg='')
+    return if (output.nil?)
+
+    self.on_print_proc.call(msg) if self.on_print_proc
+    # Errors are not subject to disabled output
+    log_output(output.print_error(msg))
+  end
+
+  #
+  # Prints a status message to the output handle.
+  #
+  def print_status(msg='')
+    return if (disable_output == true)
+
+    self.on_print_proc.call(msg) if self.on_print_proc
+    log_output(output.print_status(msg))
+  end
+
+  #
+  # Prints a good message to the output handle.
+  #
+  def print_good(msg='')
+    return if (disable_output == true)
+
+    self.on_print_proc.call(msg) if self.on_print_proc
+    log_output(output.print_good(msg))
+  end
+
+  #
+  # Prints a line of text to the output handle.
+  #
+  def print_line(msg='')
+    return if (disable_output == true)
+
+    self.on_print_proc.call(msg) if self.on_print_proc
+    log_output(output.print_line(msg))
+  end
+
+  #
+  # Prints a warning message to the output handle.
+  #
+  def print_warning(msg='')
+    return if (disable_output == true)
+
+    self.on_print_proc.call(msg) if self.on_print_proc
+    log_output(output.print_warning(msg))
+  end
+
+  #
+  # Prints a raw message to the output handle.
+  #
+  def print(msg='')
+    return if (disable_output == true)
+    self.on_print_proc.call(msg) if self.on_print_proc
+    log_output(output.print(msg))
+  end
+
+  #
+  # Whether or not output has been disabled.
+  #
+  attr_accessor :disable_output
+  #
+  # The input handle to read user input from.
+  #
+  attr_reader   :input
+  #
+  # The output handle to write output to.
+  #
+  attr_reader   :output
+
+  attr_accessor :on_command_proc
+  attr_accessor :on_print_proc
+  attr_accessor :framework
+
+protected
+
+  #
+  # Parse a line into an array of arguments.
+  #
+  def parse_line(line)
+    log_input(line)
+
+    line.gsub!(/(\r|\n)/, '')
+
+    begin
+      return args = Rex::Parser::Arguments.from_s(line)
+    rescue ::ArgumentError
+      print_error("Parse error: #{$!}")
+    end
+
+    return []
+  end
+
+  #
+  # Print the prompt, but do not log it.
+  #
+  def _print_prompt(prompt)
+    output.print(prompt)
+  end
+
+  #
+  # Writes the supplied input to the log source if one has been registered.
+  #
+  def log_input(buf)
+    rlog(buf, log_source) if (log_source)
+  end
+
+  #
+  # Writes the supplied output to the log source if one has been registered.
+  #
+  def log_output(buf)
+    rlog(buf, log_source) if (log_source)
+  end
+
+  attr_writer   :input, :output # :nodoc:
+  attr_accessor :stop_flag, :init_prompt # :nodoc:
+  attr_accessor :prompt # :nodoc:
+  attr_accessor :prompt_char, :tab_complete_proc # :nodoc:
+  attr_accessor :histfile # :nodoc:
+  attr_accessor :hist_last_saved # the number of history lines when last saved/loaded
+  attr_accessor :log_source, :stop_count # :nodoc:
+
+end
+
+###
+#
+# Pseudo-shell interface that simply includes the Shell mixin.
+#
+###
+class PseudoShell
+  include Shell
+end
+
+
+end end end
+

data/lib/rex/ui/text/table.rb

@@ -0,0 +1,346 @@
+# -*- coding: binary -*-
+require 'rex/ui'
+require 'rex/socket'
+
+module Rex
+module Ui
+module Text
+
+###
+#
+# Prints text in a tablized format.  Pretty lame at the moment, but
+# whatever.
+#
+###
+class Table
+
+  #
+  # Initializes a text table instance using the supplied properties.  The
+  # Table class supports the following hash attributes:
+  #
+  # Header
+  #
+  #	The string to display as a heading above the table.  If none is
+  #	specified, no header will be displayed.
+  #
+  # HeaderIndent
+  #
+  # 	The amount of space to indent the header.  The default is zero.
+  #
+  # Columns
+  #
+  # 	The array of columns that will exist within the table.
+  #
+  # Rows
+  #
+  # 	The array of rows that will exist.
+  #
+  # Width
+  #
+  # 	The maximum width of the table in characters.
+  #
+  # Indent
+  #
+  # 	The number of characters to indent the table.
+  #
+  # CellPad
+  #
+  # 	The number of characters to put between each horizontal cell.
+  #
+  # Prefix
+  #
+  # 	The text to prefix before the table.
+  #
+  # Postfix
+  #
+  # 	The text to affix to the end of the table.
+  #
+  # Sortindex
+  #
+  #	The column to sort the table on, -1 disables sorting.
+  #
+  def initialize(opts = {})
+    self.header   = opts['Header']
+    self.headeri  = opts['HeaderIndent'] || 0
+    self.columns  = opts['Columns'] || []
+    # updated below if we got a "Rows" option
+    self.rows     = []
+
+    self.width    = opts['Width']   || 80
+    self.indent   = opts['Indent']  || 0
+    self.cellpad  = opts['CellPad'] || 2
+    self.prefix   = opts['Prefix']  || ''
+    self.postfix  = opts['Postfix'] || ''
+    self.colprops = []
+
+    self.sort_index  = opts['SortIndex'] || 0
+
+    # Default column properties
+    self.columns.length.times { |idx|
+      self.colprops[idx] = {}
+      self.colprops[idx]['MaxWidth'] = self.columns[idx].length
+    }
+
+    # ensure all our internal state gets updated with the given rows by
+    # using add_row instead of just adding them to self.rows.  See #3825.
+    opts['Rows'].each { |row| add_row(row) } if opts['Rows']
+
+    # Merge in options
+    if (opts['ColProps'])
+      opts['ColProps'].each_key { |col|
+        idx = self.columns.index(col)
+
+        if (idx)
+          self.colprops[idx].merge!(opts['ColProps'][col])
+        end
+      }
+    end
+
+  end
+
+  #
+  # Converts table contents to a string.
+  #
+  def to_s
+    str  = prefix.dup
+    str << header_to_s || ''
+    str << columns_to_s || ''
+    str << hr_to_s || ''
+
+    sort_rows
+    rows.each { |row|
+      if (is_hr(row))
+        str << hr_to_s
+      else
+        str << row_to_s(row)
+      end
+    }
+
+    str << postfix
+
+    return str
+  end
+
+  #
+  # Converts table contents to a csv
+  #
+  def to_csv
+    str = ''
+    str << ( columns.join(",") + "\n" )
+    rows.each { |row|
+      next if is_hr(row)
+      str << ( row.map{|x|
+        x = x.to_s
+
+        x.gsub(/[\r\n]/, ' ').gsub(/\s+/, ' ').gsub('"', '""')
+      }.map{|x| "\"#{x}\"" }.join(",") + "\n" )
+    }
+    str
+  end
+
+  #
+  #
+  # Returns the header string.
+  #
+  def header_to_s # :nodoc:
+    if (header)
+      pad = " " * headeri
+
+      return pad + header + "\n" + pad + "=" * header.length + "\n\n"
+    end
+
+    return ''
+  end
+
+  #
+  # Prints the contents of the table.
+  #
+  def print
+    puts to_s
+  end
+
+  #
+  # Adds a row using the supplied fields.
+  #
+  def <<(fields)
+    add_row(fields)
+  end
+
+  #
+  # Adds a row with the supplied fields.
+  #
+  def add_row(fields = [])
+    if fields.length != self.columns.length
+      raise RuntimeError, 'Invalid number of columns!'
+    end
+    fields.each_with_index { |field, idx|
+      if (colprops[idx]['MaxWidth'] < field.to_s.length)
+        colprops[idx]['MaxWidth'] = field.to_s.length
+      end
+    }
+
+    rows << fields
+  end
+
+  #
+  # Sorts the rows based on the supplied index of sub-arrays
+  # If the supplied index is an IPv4 address, handle it differently, but
+  # avoid actually resolving domain names.
+  #
+  def sort_rows(index=sort_index)
+    return if index == -1
+    return unless rows
+    rows.sort! do |a,b|
+      if a[index].nil?
+        -1
+      elsif b[index].nil?
+        1
+      elsif Rex::Socket.dotted_ip?(a[index]) and Rex::Socket.dotted_ip?(b[index])
+        Rex::Socket::addr_atoi(a[index]) <=> Rex::Socket::addr_atoi(b[index])
+      elsif a[index] =~ /^[0-9]+$/ and b[index] =~ /^[0-9]+$/
+        a[index].to_i <=> b[index].to_i
+      else
+        a[index] <=> b[index] # assumes otherwise comparable.
+      end
+    end
+  end
+
+  #
+  # Adds a horizontal line.
+  #
+  def add_hr
+    rows << '__hr__'
+  end
+
+  #
+  # Returns new sub-table with headers and rows maching column names submitted
+  #
+  def [](*col_names)
+    tbl = self.class.new('Indent' => self.indent,
+                         'Header' => self.header,
+                         'Columns' => col_names)
+    indexes = []
+
+    col_names.each do |col_name|
+      index = self.columns.index(col_name)
+      raise RuntimeError, "Invalid column name #{col_name}" if index.nil?
+      indexes << index
+    end
+
+    self.rows.each do |old_row|
+      new_row = []
+      indexes.map {|i| new_row << old_row[i]}
+      tbl << new_row
+    end
+
+    return tbl
+  end
+
+
+  alias p print
+
+  attr_accessor :header, :headeri # :nodoc:
+  attr_accessor :columns, :rows, :colprops # :nodoc:
+  attr_accessor :width, :indent, :cellpad # :nodoc:
+  attr_accessor :prefix, :postfix # :nodoc:
+  attr_accessor :sort_index # :nodoc:
+
+protected
+
+  #
+  # Defaults cell widths and alignments.
+  #
+  def defaults # :nodoc:
+    self.columns.length.times { |idx|
+    }
+  end
+
+  #
+  # Checks to see if the row is an hr.
+  #
+  def is_hr(row) # :nodoc:
+    return ((row.kind_of?(String)) && (row == '__hr__'))
+  end
+
+  #
+  # Converts the columns to a string.
+  #
+  def columns_to_s # :nodoc:
+    nameline = ' ' * indent
+    barline  = nameline.dup
+    last_col = nil
+    last_idx = nil
+    columns.each_with_index { |col,idx|
+      if (last_col)
+        nameline << pad(' ', last_col, last_idx)
+
+        remainder = colprops[last_idx]['MaxWidth'] - last_col.length
+        if (remainder < 0)
+          remainder = 0
+        end
+        barline << (' ' * (cellpad + remainder))
+      end
+      nameline << col
+      barline << ('-' * col.length)
+
+      last_col = col
+      last_idx = idx
+    }
+
+    return "#{nameline}\n#{barline}"
+  end
+
+  #
+  # Converts an hr to a string.
+  #
+  def hr_to_s # :nodoc:
+    return "\n"
+  end
+
+  #
+  # Converts a row to a string.
+  #
+  def row_to_s(row) # :nodoc:
+    line = ' ' * indent
+    last_cell = nil
+    last_idx = nil
+    row.each_with_index { |cell, idx|
+      if (idx != 0)
+        line << pad(' ', last_cell.to_s, last_idx)
+      end
+      # line << pad(' ', cell.to_s, idx)
+      # Limit wide cells
+      if colprops[idx]['MaxChar']
+        last_cell = cell.to_s[0..colprops[idx]['MaxChar'].to_i]
+        line << last_cell
+      else
+        line << cell.to_s
+        last_cell = cell
+      end
+      last_idx = idx
+    }
+
+    return line + "\n"
+  end
+
+  #
+  # Pads out with the supplied character for the remainder of the space given
+  # some text and a column index.
+  #
+  def pad(chr, buf, colidx, use_cell_pad = true) # :nodoc:
+    remainder = colprops[colidx]['MaxWidth'] - buf.length
+    val       = chr * remainder;
+
+    if (use_cell_pad)
+      val << ' ' * cellpad
+    end
+
+    return val
+  end
+
+
+end
+
+end
+end
+end