Allow opting out of removing whitespace

adfoster-r7 · Feb 16, 2024 · 5e70ebd · 5e70ebd
1 parent 1a7b639
commit 5e70ebd
Show file tree

Hide file tree

Showing 2 changed files with 79 additions and 1 deletion.
diff --git a/lib/rex/text/wrapped_table.rb b/lib/rex/text/wrapped_table.rb
@@ -87,6 +87,7 @@ def initialize(opts = {})
       self.colprops[idx] = {}
       self.colprops[idx]['Width'] = nil
       self.colprops[idx]['WordWrap'] = true
+      self.colprops[idx]['Strip'] = false
       self.colprops[idx]['Stylers'] = []
       self.colprops[idx]['Formatters'] = []
       self.colprops[idx]['ColumnStylers'] = []
@@ -647,7 +648,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]
 
     colprops[idx]['Formatters'].each do |f|
       str_cp = f.format(str_cp)

diff --git a/spec/rex/text/wrapped_table_spec.rb b/spec/rex/text/wrapped_table_spec.rb
@@ -221,6 +221,82 @@ def style(str)
       TABLE
     end
 
+    context 'when values have trailing and preeding 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