From 67a084f7da3ec867e16729148890fa205933989a Mon Sep 17 00:00:00 2001 From: adfoster-r7 Date: Fri, 16 Feb 2024 17:41:58 +0000 Subject: [PATCH] Allow opting out of removing whitespace --- lib/rex/text/table.rb | 16 +--- lib/rex/text/wrapped_table.rb | 11 ++- spec/rex/text/table_spec.rb | 6 +- spec/rex/text/wrapped_table_spec.rb | 121 +++++++++++++++++++++++++++- 4 files changed, 136 insertions(+), 18 deletions(-) diff --git a/lib/rex/text/table.rb b/lib/rex/text/table.rb index ccd5a00..661b5de 100644 --- a/lib/rex/text/table.rb +++ b/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 diff --git a/lib/rex/text/wrapped_table.rb b/lib/rex/text/wrapped_table.rb index d4e5e3e..8474aa8 100644 --- a/lib/rex/text/wrapped_table.rb +++ b/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] 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) diff --git a/spec/rex/text/table_spec.rb b/spec/rex/text/table_spec.rb index e0afbfd..1304725 100644 --- a/spec/rex/text/table_spec.rb +++ b/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 diff --git a/spec/rex/text/wrapped_table_spec.rb b/spec/rex/text/wrapped_table_spec.rb index 6d44af5..eb97840 100644 --- a/spec/rex/text/wrapped_table_spec.rb +++ b/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,