Land #11126, Update sessions through the DBManager

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

@@ -20,6 +20,96 @@ module SessionDataProxy
       self.log_error(e, "Problem reporting session")
     end
   end
+
+  # Update the attributes of a session entry using opts.
+  # If opts is a Hash, the values should match the attributes to update and must contain :id.
+  # If opts is a Msf::Session object, it is converted to a Hash and used for the update.
+  # The db_record attribute of the Msf::Session object is updated using the returned Mdm::Session.
+  #
+  # @param opts [Hash|Msf::Session] Hash containing the updated values. Key should match the attribute to update.
+  #   Must contain :id of record to update. Otherwise, a Msf::Session object is used to update all attributes.
+  # @return [Mdm::Session] The updated Mdm::Session object.
+  def update_session(opts)
+    begin
+      self.data_service_operation do |data_service|
+        is_msf_session = false
+        if !opts.nil? && opts.kind_of?(Msf::Session)
+          msf_session = opts
+          is_msf_session = true
+          tmp_opts = SessionDataProxy.convert_msf_session_to_hash(msf_session)
+          # only updating session data
+          opts = tmp_opts[:session_data]
+          # add back session ID
+          opts[:id] = msf_session.db_record.id
+        end
+
+        mdm_session = data_service.update_session(opts)
+
+        # reassign returned Mdm::Session to the Msf::Session's db_record
+        msf_session.db_record = mdm_session if is_msf_session
+
+        mdm_session
+      end
+    rescue => e
+      self.log_error(e, "Problem updating session")
+    end
+  end
+
+  # TODO: handle task info
+  def self.convert_msf_session_to_hash(msf_session)
+    hash = Hash.new()
+    hash[:host_data] = parse_host_opts(msf_session)
+    hash[:session_data] = parse_session_data(msf_session)
+
+    if (msf_session.via_exploit)
+      hash[:vuln_info] = parse_vuln_info(msf_session)
+    end
+
+    return hash
+  end
+
+  #######
+  private
+  #######
+
+  def self.parse_session_data(msf_session)
+    hash = Hash.new()
+    # TODO: what to do with this shiz
+    hash[:datastore] = msf_session.exploit_datastore.to_h
+    hash[:desc] = msf_session.info
+    hash[:local_id] = msf_session.sid
+    hash[:platform] = msf_session.session_type
+    hash[:port] = msf_session.session_port
+    hash[:stype] = msf_session.type
+    hash[:via_exploit] = msf_session.via_exploit
+    hash[:via_payload] = msf_session.via_payload
+    return hash
+  end
+
+  def self.parse_host_opts(msf_session)
+    hash = Hash.new()
+    hash[:host] = msf_session.session_host
+    hash[:arch] = msf_session.arch if msf_session.respond_to?(:arch) and msf_session.arch
+    hash[:workspace] = msf_session.workspace || msf_session[:workspace]
+    return hash
+  end
+
+  def self.parse_vuln_info(msf_session)
+    hash = Hash.new()
+    if msf_session.via_exploit == "exploit/multi/handler" and msf_session.exploit_datastore['ParentModule']
+      hash[:mod_fullname] = msf_session.exploit_datastore['ParentModule']
+    else
+      hash[:mod_fullname] = msf_session.via_exploit
+    end
+
+    hash[:remote_port] = msf_session.exploit_datastore["RPORT"]
+    hash[:username] = msf_session.username
+    hash[:run_id] = msf_session.exploit.user_data.try(:[], :run_id)
+
+    hash[:mod_name] = msf_session.exploit.name
+    hash[:mod_references] = msf_session.exploit.references
+    return hash
+  end
 end
 
 

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

@@ -14,8 +14,7 @@ module RemoteSessionDataService
   def report_session(opts)
     session = opts[:session]
     if (session.kind_of? Msf::Session)
-      opts = convert_msf_session_to_hash(session)
-      opts[:session_dto] = true
+      opts = SessionDataProxy.convert_msf_session_to_hash(session)
     elsif (opts[:host])
       opts[:host] = opts[:host].address
     end
@@ -25,60 +24,14 @@ module RemoteSessionDataService
     session.db_record = sess_db
   end
 
-  #######
-  private
-  #######
-
-  # TODO: handle task info
-  def convert_msf_session_to_hash(msf_session)
-    hash = Hash.new()
-    hash[:host_data] = parse_host_opts(msf_session)
-    hash[:session_data] = parse_session_data(msf_session)
-
-    if (msf_session.via_exploit)
-      hash[:vuln_info] = parse_vuln_info(msf_session)
+  def update_session(opts)
+    path = SESSION_API_PATH
+    if opts && opts[:id]
+      id = opts.delete(:id)
+      path = "#{SESSION_API_PATH}/#{id}"
     end
 
-    return hash
+    json_to_mdm_object(self.put_data(path, opts), SESSION_MDM_CLASS, []).first
   end
 
-  def parse_session_data(msf_session)
-    hash = Hash.new()
-    # TODO: what to do with this shiz
-    hash[:datastore] = msf_session.exploit_datastore.to_h
-    hash[:desc] = msf_session.info
-    hash[:local_id] = msf_session.sid
-    hash[:platform] = msf_session.session_type
-    hash[:port] = msf_session.session_port
-    hash[:stype] = msf_session.type
-    hash[:via_exploit] = msf_session.via_exploit
-    hash[:via_payload] = msf_session.via_payload
-    return hash
-  end
-
-  def parse_host_opts(msf_session)
-    hash = Hash.new()
-    hash[:host] = msf_session.session_host
-    hash[:arch] = msf_session.arch if msf_session.respond_to?(:arch) and msf_session.arch
-    hash[:workspace] = msf_session.workspace || msf_session[:workspace]
-    return hash
-  end
-
-  def parse_vuln_info(msf_session)
-    hash = Hash.new()
-    if msf_session.via_exploit == "exploit/multi/handler" and msf_session.exploit_datastore['ParentModule']
-      hash[:mod_fullname] = msf_session.exploit_datastore['ParentModule']
-    else
-      hash[:mod_fullname] = msf_session.via_exploit
-    end
-
-    hash[:remote_port] = msf_session.exploit_datastore["RPORT"]
-    hash[:username] = msf_session.username
-    hash[:run_id] = msf_session.exploit.user_data.try(:[], :run_id)
-
-    hash[:mod_name] = msf_session.exploit.name
-    hash[:mod_references] = msf_session.exploit.references
-    return hash
-  end
 end
-

lib/metasploit/framework/data_service/stubs/session_data_service.rb

@@ -6,4 +6,8 @@ module SessionDataService
   def report_session(opts)
     raise 'SessionDataService#report_session is not implemented'
   end
+
+  def update_session(opts)
+    raise 'SessionDataService#update_session is not implemented'
+  end
 end

lib/msf/base/sessions/meterpreter.rb

@@ -522,8 +522,7 @@ class Meterpreter < Rex::Post::Meterpreter::Client
           })
 
           if self.db_record
-            self.db_record.desc = safe_info
-            self.db_record.save!
+            framework.db.update_session(self)
           end
 
           # XXX: This is obsolete given the Mdm::Host.normalize_os() support for host.os.session_fingerprint

lib/msf/core/db_manager/host.rb

@@ -258,9 +258,9 @@ module Msf::DBManager::Host
         msf_import_timestamps(opts, host)
         host.save!
       end
-    rescue ActiveRecord::RecordNotUnique
-      # two concurrent report requests for a new host could result in a RecordNotUnique exception
-      # simply retry the report once more as an optimistic approach
+    rescue ActiveRecord::RecordNotUnique, ActiveRecord::RecordInvalid
+      # two concurrent report requests for a new host could result in a RecordNotUnique or
+      # RecordInvalid exception, simply retry the report once more as an optimistic approach
       retry if (retry_attempts+=1) <= 1
       raise
     end

lib/msf/core/db_manager/session.rb

@@ -178,6 +178,20 @@ module Msf::DBManager::Session
     }
   end
 
+  # Update the attributes of a session entry with the values in opts.
+  # The values in opts should match the attributes to update.
+  #
+  # @param opts [Hash] Hash containing the updated values. Key should match the attribute to update. Must contain :id of record to update.
+  # @return [Mdm::Session] The updated Mdm::Session object.
+  def update_session(opts)
+    return if not active
+
+    ::ActiveRecord::Base.connection_pool.with_connection {
+      id = opts.delete(:id)
+      ::Mdm::Session.update(id, opts)
+    }
+  end
+
   # Clean out any stale sessions that have been orphaned by a dead framework instance.
   # @param last_seen_interval [Integer] interval, in seconds, open sessions are marked as alive
   def remove_stale_sessions(last_seen_interval)

lib/msf/core/session.rb

@@ -285,9 +285,7 @@ module Session
   def cleanup
     if db_record and framework.db.active
       ::ActiveRecord::Base.connection_pool.with_connection {
-        db_record.closed_at = Time.now.utc
-        # ignore exceptions
-        db_record.save
+        framework.db.update_session(id: db_record.id, closed_at: Time.now.utc, close_reason: db_record.close_reason)
         db_record = nil
       }
     end

lib/msf/core/session_manager.rb

@@ -119,8 +119,7 @@ class SessionManager < Hash
               # as recently seen.  This notifies other framework instances that this
               # session is being maintained.
               if s.db_record
-                s.db_record.last_seen = Time.now.utc
-                s.db_record.save
+                s.db_record = framework.db.update_session(id: s.db_record.id, last_seen: Time.now.utc)
               end
             end
           end

lib/msf/core/web_services/servlet/session_servlet.rb

@@ -11,6 +11,7 @@ module SessionServlet
   def self.registered(app)
     app.get SessionServlet.api_path_with_id, &get_session
     app.post SessionServlet.api_path, &report_session
+    app.put SessionServlet.api_path_with_id, &update_session
   end
 
   #######
@@ -50,4 +51,19 @@ module SessionServlet
     }
   end
 
+  def self.update_session
+    lambda {
+      warden.authenticate!
+      begin
+        opts = parse_json_request(request, false)
+        tmp_params = sanitize_params(params)
+        opts[:id] = tmp_params[:id] if tmp_params[:id]
+        data = get_db.update_session(opts)
+        set_json_data_response(response: data)
+      rescue => e
+        print_error_and_create_response(error: e, message: 'There was an error updating the session:', code: 500)
+      end
+    }
+  end
+
 end

lib/msf/core/web_services/servlet/user_servlet.rb

@@ -51,11 +51,9 @@ module UserServlet
         tmp_params = sanitize_params(params)
         opts[:id] = tmp_params[:id] if tmp_params[:id]
         data = get_db.update_user(opts)
-        # Only return the single object if the id parameter is present
-        data = data.first if !sanitized_params[:id].nil? && data.count == 1
         set_json_data_response(response: data)
       rescue => e
-        print_error_and_create_response(error: e, message: 'There was an error creating the user:', code: 500)
+        print_error_and_create_response(error: e, message: 'There was an error updating the user:', code: 500)
       end
     }
   end