Land #68, Allow opting out of whitespace removal

lib/rex/text/table.rb

@@ -21,21 +21,13 @@ class Table
   # To enforce all tables to be wrapped to the terminal's current width, call `Table.wrap_tables!`
   # before invoking `Table.new` as normal.
   def self.new(*args, &block)
-    if wrap_table?(args)
-      table_options = args.first
-      return ::Rex::Text::WrappedTable.new(table_options)
-    end
-    return super(*args, &block)
-  end
-
-  def self.wrap_table?(args)
-    return false unless wrapped_tables?
+    return super(*args, &block) unless wrap_table?(args)
 
     table_options = args.first
-    if table_options&.key?('WordWrap')
-      return table_options['WordWrap']
-    end
+    ::Rex::Text::WrappedTable.new(table_options)
+  end
 
+  def self.wrap_table?(args)
     wrapped_tables?
   end
 

lib/rex/text/wrapped_table.rb

@@ -72,6 +72,7 @@ def initialize(opts = {})
     self.rows     = []
 
     self.width    = opts['Width']   || ::IO.console&.winsize&.[](1) || ::BigDecimal::INFINITY
+    self.word_wrap = opts.fetch('WordWrap', true)
     self.indent   = opts['Indent']  || 0
     self.cellpad  = opts['CellPad'] || 2
     self.prefix   = opts['Prefix']  || ''
@@ -87,6 +88,7 @@ def initialize(opts = {})
       self.colprops[idx] = {}
       self.colprops[idx]['Width'] = nil
       self.colprops[idx]['WordWrap'] = true
+      self.colprops[idx]['Strip'] = true
       self.colprops[idx]['Stylers'] = []
       self.colprops[idx]['Formatters'] = []
       self.colprops[idx]['ColumnStylers'] = []
@@ -328,7 +330,7 @@ def [](*col_names)
 
   attr_accessor :header, :headeri # :nodoc:
   attr_accessor :columns, :rows, :colprops # :nodoc:
-  attr_accessor :width, :indent, :cellpad # :nodoc:
+  attr_accessor :width, :word_wrap, :indent, :cellpad # :nodoc:
   attr_accessor :prefix, :postfix # :nodoc:
   attr_accessor :sort_index, :sort_order, :scterm # :nodoc:
 
@@ -478,6 +480,8 @@ def find_color_type_and_value(string)
   # @param [Array<String>] values
   # @param [Integer] optimal_widths
   def chunk_values(values, optimal_widths)
+    return values.map { |value| [value] } unless word_wrap
+
     # First split long strings into an array of chunks, where each chunk size is the calculated column width
     values_as_chunks = values.each_with_index.map do |value, idx|
       color_state = {}
@@ -586,6 +590,8 @@ def calculate_optimal_widths(styled_columns, styled_rows)
       end
     end
 
+    return display_width_metadata.map { |metadata| metadata[:max_display_width] } unless word_wrap
+
     # Calculate the sizes set by the user
     user_influenced_column_widths = colprops.map do |colprop|
       if colprop['Width']
@@ -647,7 +653,8 @@ def calculate_optimal_widths(styled_columns, styled_rows)
   end
 
   def format_table_field(str, idx)
-    str_cp = str.to_s.encode('UTF-8', invalid: :replace, undef: :replace).strip
+    str_cp = str.to_s.encode('UTF-8', invalid: :replace, undef: :replace)
+    str_cp = str_cp.strip if colprops[idx]['Strip']
 
     colprops[idx]['Formatters'].each do |f|
       str_cp = f.format(str_cp)

spec/rex/text/table_spec.rb

@@ -39,12 +39,12 @@ def style(str)
         expect(described_class.wrap_table?(default_table_options)).to be true
       end
 
-      it "returns the user preference when set to true" do
+      it "ignores the user preference when set to true" do
         expect(described_class.wrap_table?(enable_wrapped_table_options)).to be true
       end
 
-      it "returns the user preference when set to false" do
-        expect(described_class.wrap_table?(disable_wrapped_table_options)).to be false
+      it "ignores the user preference when set to false" do
+        expect(described_class.wrap_table?(disable_wrapped_table_options)).to be true
       end
     end
 

spec/rex/text/wrapped_table_spec.rb

@@ -34,6 +34,14 @@ def style(str)
     allow(Rex::Text::Table).to receive(:wrap_table?).with(anything).and_return(true)
   end
 
+  before(:each) do
+    described_class.wrap_tables!
+  end
+
+  after(:each) do
+    described_class.unwrap_tables!
+  end
+
   describe "#to_csv" do
     it "handles strings in different encodings" do
       options = {
@@ -221,6 +229,82 @@ def style(str)
       TABLE
     end
 
+    context 'when values have trailing and preceding whitespae' do
+      it 'strips values by default' do
+        whitespace = "     "
+        col_1_field = whitespace + "A" * 5 + whitespace
+        col_2_field = whitespace + "B" * 50 + whitespace
+        col_3_field = whitespace + "C" * 15 + whitespace
+
+        options = {
+          'Header' => 'Header',
+          'Columns' => [
+            'Column 1',
+            'Column 2',
+            'Column 3'
+          ]
+        }
+
+        tbl = Rex::Text::Table.new(options)
+
+        tbl << [
+          col_1_field,
+          col_2_field,
+          col_3_field
+        ]
+
+        expect(tbl).to match_table <<~TABLE
+          Header
+          ======
+
+          Column 1  Column 2                                            Column 3
+          --------  --------                                            --------
+          AAAAA     BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB  CCCCCCCCCCCCCCC
+        TABLE
+      end
+
+      it 'allows conditional stripping of values' do
+        whitespace = "     "
+        col_1_field = whitespace + "A" * 5 + whitespace
+        col_2_field = whitespace + "B" * 50 + whitespace
+        col_3_field = whitespace + "C" * 15 + whitespace
+
+        options = {
+          'Header' => 'Header',
+          'Columns' => [
+            'Column 1',
+            'Column 2',
+            'Column 3'
+          ],
+          'ColProps' => {
+            'Column 1' => {
+              'Strip' => true
+            },
+            'Column 2' => {
+              'Strip' => false
+            }
+          }
+        }
+
+        tbl = Rex::Text::Table.new(options)
+
+        tbl << [
+          col_1_field,
+          col_2_field,
+          col_3_field
+        ]
+
+        expect(tbl).to match_table <<~TABLE
+          Header
+          ======
+
+          Column 1  Column 2                                                      Column 3
+          --------  --------                                                      --------
+          AAAAA          BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB       CCCCCCCCCCCCCCC
+        TABLE
+      end
+    end
+
     context 'when using arrow indicators' do
       let(:empty_column_header_styler) do
         clazz = Class.new do
@@ -442,6 +526,41 @@ def style(str)
       TABLE
     end
 
+    it 'allows wordwrapping to be disabled globally' do
+      options = {
+        'Header' => 'Header',
+        'Width' => 3,
+        'WordWrap' => false,
+        'Columns' => [
+          'Column 1',
+          'Column 2',
+          'Column 3'
+        ],
+        'ColProps' => {
+          'Column 2' => {
+            'Stylers' => [styler]
+          }
+        }
+      }
+
+      tbl = Rex::Text::Table.new(options)
+
+      tbl << [
+        "A" * 5,
+        "ABC ABCD ABC" * 1,
+        "C" * 5
+      ]
+
+      expect(tbl).to match_table <<~TABLE
+        Header
+        ======
+
+        Column 1  Column 2      Column 3
+        --------  --------      --------
+        AAAAA     %bluABC ABCD ABC%clr  CCCCC
+      TABLE
+    end
+
     it 'should apply field stylers correctly and NOT increase column length when having a low width value' do
       options = {
         'Header' => 'Header',
@@ -773,7 +892,7 @@ def style(str)
         expect(tbl.to_s.lines).to all(have_maximum_display_width(80))
       end
 
-      it "Wraps columns as well as values" do
+      it "wraps columns as well as values" do
         options = {
           'Header' => 'Header',
           'Indent' => 2,