Fix search command for single characters

rapid7 · Mar 22, 2021 · 33ea973 · 33ea973
1 parent 2126caf
commit 33ea973
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 12 deletions.
diff --git a/lib/msf/ui/console/table_print/highlight_substring_styler.rb b/lib/msf/ui/console/table_print/highlight_substring_styler.rb
@@ -5,25 +5,18 @@ module Ui
     module Console
       module TablePrint
         class HighlightSubstringStyler
-          COLOR = '%bgmag'
+          HIGHLIGHT_COLOR = '%bgmag'
+          RESET_COLOR = '%clr'
 
           def initialize(substrings)
             @substrings = substrings
           end
 
           def style(value)
-            value_cp = value.clone
+            search_terms = @substrings.map { |substring| Regexp.escape(substring) }
+            search_pattern = /#{search_terms.join('|')}/i
 
-            @substrings.each do |s|
-              # Regex used to pull out matches and preserve case sensitivity
-              matches = value_cp.scan(%r{#{Regexp.escape(s)}}i)
-
-              matches.each do |m|
-                value_cp.gsub!(m, COLOR + m + '%clr')
-              end
-            end
-
-            value_cp
+            value.gsub(search_pattern) { |match| "#{HIGHLIGHT_COLOR}#{match}#{RESET_COLOR}" }
           end
         end
       end

diff --git a/spec/lib/msf/ui/console/table_print/highlight_substring_styler_spec.rb b/spec/lib/msf/ui/console/table_print/highlight_substring_styler_spec.rb
@@ -23,5 +23,12 @@
 
       expect(styler.style(str)).to eql "AAAAA%bgmagBBB%clrAAAAA%bgmagCCC%clr"
     end
+
+    it 'should work with single characters' do
+      str = 'ABCABC'
+      styler = described_class.new(%w(a b c))
+
+      expect(styler.style(str)).to eql "%bgmagA%clr%bgmagB%clr%bgmagC%clr%bgmagA%clr%bgmagB%clr%bgmagC%clr"
+    end
   end
 end