From d63ff8a559de970b8de755f86ec8fbc805c1ddde Mon Sep 17 00:00:00 2001
From: Alan Foster <alandavid_foster@rapid7.com>
Date: Mon, 22 Mar 2021 18:52:57 +0000
Subject: [PATCH] Fix search command for single characters

---
 .../table_print/highlight_substring_styler.rb   | 17 +++++------------
 .../highlight_substring_styler_spec.rb          |  7 +++++++
 2 files changed, 12 insertions(+), 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
index 7cf0ffdeeb..5ca540615e 100644
--- 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 Msf
     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
index ef7bb4440a..12cb15ecdb 100644
--- 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 @@ RSpec.describe Msf::Ui::Console::TablePrint::HighlightSubstringStyler do
 
       expect(styler.style(str)).to eql "AAAAA%bgmagBBB%clrAAAAA%bgmagCCC%clr"
     end
+
+    it 'should highlight 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