Add update_session method

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

@@ -21,6 +21,61 @@ module SessionDataProxy
     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
+        $stderr.puts("*** SessionDataProxy.update_session(): opts=#{opts}")  # TODO: remove
+        # session = opts[:session]
+        # $stderr.puts("*** SessionDataProxy.update_session(): session.class=#{session.class}")  # TODO: remove
+        # if !session.nil? && session.kind_of?(Msf::Session)
+        if !opts.nil? && opts.kind_of?(Msf::Session)
+          msf_session = opts
+          is_msf_session = true
+          $stderr.puts("*** SessionDataProxy.update_session(): is_msf_session=#{is_msf_session}, opts=#{opts}, opts.class=#{opts.class}")  # TODO: remove
+          # save session ID
+          # id = opts.delete(:id)
+          # tmp_opts = convert_msf_session_to_hash(session)
+          tmp_opts = SessionDataProxy.convert_msf_session_to_hash(msf_session)
+          # only updating session data
+          # opts = tmp_opts[:session_data]
+          opts = tmp_opts[:session_data]
+          # add back session ID
+          # opts[:id] = id
+          opts[:id] = msf_session.db_record.id
+          $stderr.puts("*** SessionDataProxy.update_session(): after convert_msf_session_to_hash, opts=#{opts}")  # TODO: remove
+        end
+
+        mdm_session = data_service.update_session(opts)
+
+        $stderr.puts("*** SessionDataProxy.update_session(): after update: is_msf_session=#{is_msf_session}, mdm_session=#{mdm_session}, #{mdm_session.attributes}")  # TODO: remove
+        # reassign returned Mdm::Session to the Msf::Session's db_record
+        msf_session.db_record = mdm_session if is_msf_session
+        $stderr.puts("*** SessionDataProxy.update_session(): msf_session.db_record=#{msf_session.db_record}") if is_msf_session  # TODO: remove
+
+        # TODO: remove block
+        if is_msf_session && !msf_session.db_record.nil?
+          $stderr.puts("SessionDataProxy.update_session(): msf_session.db_record=#{msf_session.db_record}, msf_session.db_record.id=#{msf_session.db_record.id}")  # TODO: remove
+        else
+          $stderr.puts("SessionDataProxy.update_session(): msf_session.db_record is nil")  # TODO: remove
+        end
+
+        mdm_session
+      end
+    rescue => e
+      $stderr.puts("SessionDataProxy.update_session(): e.backtrace=#{e.backtrace}")  # TODO: remove
+      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()

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

@@ -12,18 +12,40 @@ module RemoteSessionDataService
   end
 
   def report_session(opts)
+    $stderr.puts("RemoteSessionDataService.report_session(): opts=#{opts}")  # TODO: remove
     session = opts[:session]
     if (session.kind_of? Msf::Session)
       opts = SessionDataProxy.convert_msf_session_to_hash(session)
       opts[:session_dto] = true
     elsif (opts[:host])
+      $stderr.puts("*** RemoteSessionDataService.report_session(): executing path where session is not a kind_of Msf::Session...")  # TODO: remove
       opts[:host] = opts[:host].address
     end
 
     opts[:time_stamp] = Time.now.utc
+    $stderr.puts("RemoteSessionDataService.report_session(): opts=#{opts}")  # TODO: remove
     sess_db = json_to_mdm_object(self.post_data(SESSION_API_PATH, opts), SESSION_MDM_CLASS, []).first
+    if !sess_db.nil?
+      $stderr.puts("RemoteSessionDataService.report_session(): sess_db=#{sess_db}, sess_db.id=#{sess_db.id}")  # TODO: remove
+    else
+      $stderr.puts("RemoteSessionDataService.report_session(): sess_db is nil")  # TODO: remove
+    end
     session.db_record = sess_db
   end
 
-end
+  def update_session(opts)
+    path = SESSION_API_PATH
+    if opts && opts[:id]
+      id = opts.delete(:id)
+      path = "#{SESSION_API_PATH}/#{id}"
+    end
 
+    $stderr.puts("RemoteSessionDataService.update_session(): path=#{path}, opts=#{opts}")  # TODO: remove
+
+    sess_db = json_to_mdm_object(self.put_data(path, opts), SESSION_MDM_CLASS, []).first
+    $stderr.puts("RemoteSessionDataService.update_session(): returning... sess_db=#{sess_db}")  # TODO: remove
+
+    sess_db
+  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/core/db_manager/session.rb

@@ -178,6 +178,29 @@ 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 {
+      $stderr.puts("#{DateTime.now}  Msf::DBManager::Session.update_session(): opts=#{opts}")  # TODO: remove
+
+      id = opts.delete(:id)
+      $stderr.puts("#{DateTime.now}  Msf::DBManager::Session.update_session(): id=#{id}, opts=#{opts}")  # TODO: remove
+      $stderr.puts("#{DateTime.now}  Msf::DBManager::Session.update_session(): id=#{id}, before update: #{Mdm::Session.find(id).attributes}")  # TODO: remove
+      session_db_record = ::Mdm::Session.update(id, opts)
+      $stderr.puts("#{DateTime.now}  Msf::DBManager::Session.update_session(): session_db_record=#{session_db_record}")  # TODO: remove
+      $stderr.puts("#{DateTime.now}  Msf::DBManager::Session.update_session(): session_db_record.id=#{session_db_record.id}") unless session_db_record.nil?   # TODO: remove
+      $stderr.puts("#{DateTime.now}  Msf::DBManager::Session.update_session(): id=#{id}, after update: #{Mdm::Session.find(id).attributes}")  # TODO: remove
+
+      session_db_record
+    }
+  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/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,25 @@ module SessionServlet
     }
   end
 
+  def self.update_session
+    lambda {
+      warden.authenticate!
+      begin
+        opts = parse_json_request(request, false)
+        $stderr.puts("#{DateTime.now}  SessionServlet.update_session(): opts=#{opts}")  # TODO: remove
+        tmp_params = sanitize_params(params)
+        opts[:id] = tmp_params[:id] if tmp_params[:id]
+        $stderr.puts("#{DateTime.now}  SessionServlet.update_session(): (after mod) opts=#{opts}")  # TODO: remove
+        data = get_db.update_session(opts)
+
+        $stderr.puts("#{DateTime.now}  SessionServlet.update_session(): data=#{data}")  # TODO: remove
+        $stderr.puts("#{DateTime.now}  SessionServlet.update_session(): data.class=#{data.class}") unless data.nil?  # TODO: remove
+
+        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