Merge pull request #13 from clee-r7/externalize-host-data-search

Externalize host data search

Author: jbarnett-r7

lib/metasploit/framework/data_service/proxy/host_data_proxy.rb

@@ -1,12 +1,13 @@
 module HostDataProxy
 
-  def hosts(wspace = workspace, non_dead = false, addresses = nil)
+  def hosts(wspace = workspace, non_dead = false, addresses = nil, search_term = nil)
     begin
       data_service = self.get_data_service()
       opts = {}
       opts[:wspace] = wspace
       opts[:non_dead] = non_dead
       opts[:addresses] = addresses
+      opts[:search_term] = search_term
       data_service.hosts(opts)
     rescue Exception => e
       puts "Call to #{data_service.class}#hosts threw exception: #{e.message}"

lib/metasploit/framework/data_service/remote/http/core.rb

@@ -108,7 +108,8 @@ def put_data(path, data_hash, query = nil)
   #
   def make_request(request_type, path, data_hash = nil, query = nil)
     begin
-      query_str = (!query.nil? && !query.empty?) ? URI.encode_www_form(append_workspace(query)) : nil
+      # simplify query by removing nil values
+      query_str = (!query.nil? && !query.empty?) ? append_workspace(query).compact.to_query : nil
       uri = URI::HTTP::build({path: path, query: query_str})
       puts "#{Time.now} - HTTP #{request_type} request to #{uri.request_uri} with #{data_hash ? data_hash : "nil"}"
 

lib/metasploit/framework/data_service/remote/http/remote_host_data_service.rb

@@ -8,7 +8,7 @@ module RemoteHostDataService
   HOST_MDM_CLASS = 'Mdm::Host'
 
   def hosts(opts)
-    json_to_mdm_object(self.get_data(HOST_API_PATH, opts), HOST_MDM_CLASS, [])
+    json_to_mdm_object(self.get_data(HOST_API_PATH, nil, opts), HOST_MDM_CLASS, [])
   end
 
   def report_host(opts)

lib/msf/core/db_manager/host.rb

@@ -81,19 +81,6 @@ def add_host_tag(opts)
     end
   end
 
-  def find_hosts_with_tag(opts)
-    workspace_id = opts[:workspace_id]
-    host_address = opts[:host_address]
-    tag_name = opts[:tag_name]
-    Mdm::Tag.joins(:hosts).where("hosts.workspace_id = ? and hosts.address = ? and tags.name = ?", workspace_id, host_address, tag_name).references(:hosts).order("tags.id DESC")
-  end
-
-  def find_host_tags(opts)
-    workspace_id = opts[:workspace_id]
-    host_address = opts[:host_address]
-    Mdm::Tag.joins(:hosts).where("hosts.workspace_id = ? and hosts.address = ?", workspace.id, host_address).order("tags.id DESC")
-  end
-
   def delete_host_tag(opts)
     workspace = opts[:workspace]
     if workspace.kind_of? String
@@ -161,12 +148,21 @@ def hosts(opts)
       wspace = find_workspace(wspace)
     end
 
-  ::ActiveRecord::Base.connection_pool.with_connection {
-    conditions = {}
-    conditions[:state] = [Msf::HostState::Alive, Msf::HostState::Unknown] if opts[:non_dead]
-    conditions[:address] = opts[:addresses] if opts[:addresses]
-    wspace.hosts.where(conditions).order(:address)
-  }
+    ::ActiveRecord::Base.connection_pool.with_connection {
+
+      conditions = {}
+      conditions[:state] = [Msf::HostState::Alive, Msf::HostState::Unknown] if opts[:non_dead]
+      conditions[:address] = opts[:addresses] if opts[:addresses] && !opts[:addresses].empty?
+
+      if opts[:search_term] && !opts[:search_term].empty?
+        column_search_conditions = Msf::Util::DBManager.create_all_column_search_conditions(Mdm::Host, opts[:search_term])
+        tag_conditions = Arel::Nodes::Regexp.new(Mdm::Tag.arel_table[:name], Arel::Nodes.build_quoted("(?mi)#{opts[:search_term]}"))
+        search_conditions = column_search_conditions.or(tag_conditions)
+        wspace.hosts.where(conditions).where(search_conditions).includes(:tags).references(:tags).order(:address)
+      else
+        wspace.hosts.where(conditions).order(:address)
+      end
+    }
   end
 
   #

lib/msf/core/db_manager/http/servlet/host_servlet.rb

@@ -18,7 +18,7 @@ def self.get_host
     lambda {
       begin
         opts = parse_json_request(request, false)
-        data = get_db().hosts(opts)
+        data = get_db().hosts(params)
         includes = [:loots]
         set_json_response(data, includes)
       rescue Exception => e

lib/msf/ui/console/command_dispatcher/db.rb

@@ -542,7 +542,7 @@ def cmd_hosts(*args)
                 when '-R', '--rhosts'
                   set_rhosts = true
                 when '-S', '--search'
-                  search_term = /#{args.shift}/nmi
+                  search_term = args.shift
                 when '-i', '--info'
                   mode << :new_info
                   info_data = args.shift
@@ -648,12 +648,12 @@ def cmd_hosts(*args)
             each_host_range_chunk(host_ranges) do |host_search|
               break if !host_search.nil? && host_search.empty?
 
-              framework.db.hosts(framework.db.workspace, onlyup, host_search).each do |host|
-                if search_term
-                  next unless (
-                  host.attribute_names.any? { |a| host[a.intern].to_s.match(search_term) } || !find_hosts_with_tag(framework.db.workspace.id, host.address, search_term.source).empty?
-                  )
-                end
+              framework.db.hosts(framework.db.workspace, onlyup, host_search, search_term = search_term).each do |host|
+                # if search_term
+                #   next unless (
+                #   host.attribute_names.any? { |a| host[a.intern].to_s.match(search_term) } || !find_hosts_with_tag(framework.db.workspace.id, host.address, search_term.source).empty?
+                #   )
+                # end
 
                 columns = col_names.map do |n|
                   # Deal with the special cases

lib/msf/util.rb

@@ -25,3 +25,6 @@ module Util
 
 # Host helpers
 require 'msf/util/host'
+
+# DBManager helpers
+require 'msf/util/db_manager'

lib/msf/util/db_manager.rb

@@ -0,0 +1,19 @@
+module Msf
+  module Util
+    module DBManager
+      # Creates search conditions to match the specified search string against all of the model's columns.
+      #
+      # @param model - An ActiveRecord model object
+      # @param search - A string regex search
+      # @return Arel::Nodes::Or object that represents a search of all of the model's columns
+      def self.create_all_column_search_conditions(model, search)
+        search = "(?mi)#{search}"
+        condition_set = model.columns.map do |column|
+          Arel::Nodes::Regexp.new(Arel::Nodes::NamedFunction.new("CAST", [model.arel_table[column.name].as("TEXT")]),
+                                  Arel::Nodes.build_quoted(search))
+        end
+        condition_set.reduce { |conditions, condition| conditions.or(condition).expr }
+      end
+    end
+  end
+end