Merge pull request #15 from clee-r7/MS-2941-remote-host-update

MS-2941: Remote Host Update

Author: jbarnett-r7

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

@@ -6,7 +6,7 @@ def hosts(wspace = workspace, non_dead = false, addresses = nil, search_term = n
       opts = {}
       opts[:wspace] = wspace
       opts[:non_dead] = non_dead
-      opts[:addresses] = addresses
+      opts[:address] = addresses
       opts[:search_term] = search_term
       data_service.hosts(opts)
     rescue Exception => e
@@ -40,6 +40,15 @@ def report_hosts(hosts)
     end
   end
 
+  def update_host(opts)
+    begin
+      data_service = self.get_data_service()
+      data_service.update_host(opts)
+    rescue Exception => e
+      elog "Problem updating host: #{e.message}"
+    end
+  end
+
   def delete_host(opts)
     begin
       data_service = self.get_data_service()

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

@@ -23,6 +23,15 @@ def report_hosts(hosts)
     self.post_data(HOST_API_PATH, hosts)
   end
 
+  def update_host(opts)
+    path = HOST_API_PATH
+    if opts && opts[:id]
+      id = opts.delete(:id)
+      path = "#{HOST_API_PATH}/#{id}"
+    end
+    json_to_mdm_object(self.put_data(path, opts), HOST_MDM_CLASS, [])
+  end
+
   def delete_host(opts)
     json_to_mdm_object(self.delete_data(HOST_API_PATH, opts), HOST_MDM_CLASS, [])
   end

lib/msf/core/db_manager/host.rb

@@ -152,7 +152,8 @@ def hosts(opts)
 
       conditions = {}
       conditions[:state] = [Msf::HostState::Alive, Msf::HostState::Unknown] if opts[:non_dead]
-      conditions[:address] = opts[:addresses] if opts[:addresses] && !opts[:addresses].empty?
+      conditions[:address] = opts[:address] if opts[:address] && !opts[:address].empty?
+      conditions[:id] = opts[:id] if opts[:id] && !opts[:id].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])
@@ -268,6 +269,20 @@ def report_host(opts)
   }
   end
 
+  def update_host(opts)
+    # process workspace string for update if included in opts
+    wspace = opts.delete(:workspace)
+    if wspace.kind_of? String
+      wspace = find_workspace(wspace)
+      opts[:workspace] = wspace
+    end
+
+    ::ActiveRecord::Base.connection_pool.with_connection {
+      id = opts.delete(:id)
+      Mdm::Host.update(id, opts)
+    }
+  end
+
   def split_windows_os_name(os_name)
     return [] if os_name.nil?
     flavor_match = os_name.match(/Windows\s+(.*)/)

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

@@ -4,9 +4,14 @@ def self.api_path
     '/api/v1/hosts'
   end
 
+  def self.api_path_with_id
+    "#{HostServlet.api_path}/?:id?"
+  end
+
   def self.registered(app)
-    app.get HostServlet.api_path, &get_host
+    app.get HostServlet.api_path_with_id, &get_host
     app.post HostServlet.api_path, &report_host
+    app.put HostServlet.api_path_with_id, &update_host
     app.delete HostServlet.api_path, &delete_host
   end
 
@@ -18,7 +23,7 @@ def self.get_host
     lambda {
       begin
         opts = parse_json_request(request, false)
-        data = get_db().hosts(params)
+        data = get_db().hosts(params.symbolize_keys)
         includes = [:loots]
         set_json_response(data, includes)
       rescue Exception => e
@@ -40,6 +45,20 @@ def self.report_host
     }
   end
 
+  def self.update_host
+    lambda {
+      begin
+        opts = parse_json_request(request, false)
+        tmp_params = params.symbolize_keys
+        opts[:id] = tmp_params[:id] if tmp_params[:id]
+        data = get_db().update_host(opts)
+        set_json_response(data)
+      rescue Exception => e
+        set_error_on_response(e)
+      end
+    }
+  end
+
   def self.delete_host
     lambda {
       begin

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

@@ -231,47 +231,28 @@ def cmd_hosts_help
             cmd_hosts("-h")
           end
 
-          def change_host_info(rws, data)
-            if rws == [nil]
-              print_error("In order to change the host info, you must provide a range of hosts")
+          # Changes the specified host data
+          #
+          # @param host_ranges - range of hosts to process
+          # @param host_data - hash of host data to be updated
+          def change_host_data(host_ranges, host_data)
+            if !host_data || host_data.length != 1
+              print_error("A single key-value data hash is required to change the host data")
               return
             end
+            attribute = host_data.keys[0]
 
-            rws.each do |rw|
-              rw.each do |ip|
-                id = framework.db.get_host(:address => ip).id
-                framework.db.hosts.update(id, :info => data)
-                framework.db.report_note(:host => ip, :type => 'host.info', :data => data)
-              end
-            end
-          end
-
-          def change_host_name(rws, data)
-            if rws == [nil]
-              print_error("In order to change the host name, you must provide a range of hosts")
+            if host_ranges == [nil]
+              print_error("In order to change the host #{attribute}, you must provide a range of hosts")
               return
             end
 
-            rws.each do |rw|
-              rw.each do |ip|
-                id = framework.db.get_host(:address => ip).id
-                framework.db.hosts.update(id, :name => data)
-                framework.db.report_note(:host => ip, :type => 'host.name', :data => data)
-              end
-            end
-          end
-
-          def change_host_comment(rws, data)
-            if rws == [nil]
-              print_error("In order to change the comment, you must provide a range of hosts")
-              return
-            end
+            each_host_range_chunk(host_ranges) do |host_search|
+              break if !host_search.nil? && host_search.empty?
 
-            rws.each do |rw|
-              rw.each do |ip|
-                id = framework.db.get_host(:address => ip).id
-                framework.db.hosts.update(id, :comments => data)
-                framework.db.report_note(:host => ip, :type => 'host.comments', :data => data)
+              framework.db.hosts(framework.db.workspace, false, host_search).each do |host|
+                framework.db.update_host(host_data.merge(id: host.id))
+                framework.db.report_note(host: host.address, type: "host.#{attribute}", data: host_data[attribute])
               end
             end
           end
@@ -492,13 +473,13 @@ def cmd_hosts(*args)
 
             case
               when mode == [:new_info]
-                change_host_info(host_ranges, info_data)
+                change_host_data(host_ranges, info: info_data)
                 return
               when mode == [:new_name]
-                change_host_name(host_ranges, name_data)
+                change_host_data(host_ranges, name: name_data)
                 return
               when mode == [:new_comment]
-                change_host_comment(host_ranges, comment_data)
+                change_host_data(host_ranges, comments: comment_data)
                 return
               when mode == [:tag]
                 begin