Create Mdm::Session relationships

MSP-12172

app/models/mdm/session.rb

@@ -34,12 +34,29 @@ class Mdm::Session < ActiveRecord::Base
   # @!attribute [rw] routes
   #   Routes tunneled throug this session.
   #
-  #   @return [Array<Mdm::Route>]
+  #   @return [ActiveRecord::Relation<Mdm::Route>]
   has_many :routes,
            class_name: 'Mdm::Route',
            dependent: :delete_all,
            inverse_of: :session
 
+  # @!attribute [rw] originating_module_run
+  #   Records the Metasploit modules run that created this session
+  #
+  #   @return [MetasploitDataModels::ModuleRun]
+  belongs_to :originating_module_run,
+           class_name: 'MetasploitDataModels::ModuleRun',
+           foreign_key: :module_run_id,
+           inverse_of: :spawned_session
+
+  # @!attribute [rw] target_module_runs
+  #   Records the Metasploit modules run on this session
+  #
+  #   @return [ActiveRecord::Relation<MetasploitDataModels::ModuleRun>]
+  has_many :target_module_runs,
+           class_name: 'MetasploitDataModels::ModuleRun',
+           inverse_of: :target_session
+
   # @!attribute vuln_attempt
   #   Vulnerability attempt that created this session.
   #

app/models/metasploit_data_models/module_run.rb

@@ -65,16 +65,49 @@ class MetasploitDataModels::ModuleRun < ActiveRecord::Base
   #
 
 
-  # A reference to the Metasploit content module in the DB cache
-  # @return [Mdm::Module::Detail]
+  # @!attribute [rw] module_detail
+  #  A reference to the Metasploit content module in the DB cache
+  #
+  #  @return [Mdm::Module::Detail]
   belongs_to :module_detail,
              class_name: 'Mdm::Module::Detail',
              inverse_of: :module_runs
 
+
+  # @!attribute [rw] spawned_session
+  #
+  #  The session created by running this module.
+  #  Note that this is NOT the session that modules are run on.
+  #
+  #  @return [Mdm::Session]
+  has_one :spawned_session,
+             class_name: 'Mdm::Session',
+             inverse_of: :originating_module_run
+
+
+  # @!attribute [rw] target_session
+  #
+  #  The session this module was run on, if any.
+  #  Note that this is NOT a session created by this module run
+  #  of exploit modules.
+  #
+  #  @return [Mdm::Session]
+  belongs_to :target_session,
+             class_name: 'Mdm::Session',
+             foreign_key: :session_id,
+             inverse_of: :target_module_runs
+
+
+
+  # Declares this model to implement a polymorphic relationship with other models.
   belongs_to :trackable, polymorphic: true
 
-  # The user that launched this module
-  # @return [Mdm::User]
+
+  # @!attribute [rw] user
+  #
+  #  The user that launched this module
+  #
+  #  @return [Mdm::User]
   belongs_to :user,
              class_name:  'Mdm::User',
              foreign_key: 'user_id',

db/migrate/20150219215039_add_module_run_to_session.rb

@@ -0,0 +1,8 @@
+class AddModuleRunToSession < ActiveRecord::Migration
+  def change
+    change_table :sessions do |t|
+      t.integer :module_run_id
+    end
+    add_index :sessions, :module_run_id
+  end
+end

spec/app/models/mdm/session_spec.rb

@@ -25,31 +25,34 @@
   context 'database' do
 
     context 'timestamps'do
-      it { should have_db_column(:opened_at).of_type(:datetime).with_options(:null => false) }
-      it { should have_db_column(:closed_at).of_type(:datetime) }
-      it { should have_db_column(:last_seen).of_type(:datetime) }
+      it { is_expected.to have_db_column(:closed_at).of_type(:datetime) }
+      it { is_expected.to have_db_column(:last_seen).of_type(:datetime) }
+      it { is_expected.to have_db_column(:opened_at).of_type(:datetime).with_options(:null => false) }
     end
 
     context 'columns' do
-      it { should have_db_column(:host_id).of_type(:integer) }
-      it { should have_db_column(:stype).of_type(:string) }
-      it { should have_db_column(:via_exploit).of_type(:string) }
-      it { should have_db_column(:via_payload).of_type(:string) }
-      it { should have_db_column(:desc).of_type(:string) }
-      it { should have_db_column(:port).of_type(:integer) }
-      it { should have_db_column(:platform).of_type(:string) }
-      it { should have_db_column(:datastore).of_type(:text) }
-      it { should have_db_column(:local_id).of_type(:integer) }
+      it { is_expected.to have_db_column(:datastore).of_type(:text) }
+      it { is_expected.to have_db_column(:desc).of_type(:string) }
+      it { is_expected.to have_db_column(:host_id).of_type(:integer) }
+      it { is_expected.to have_db_column(:local_id).of_type(:integer) }
+      it { is_expected.to have_db_column(:module_run_id).of_type(:integer) }
+      it { is_expected.to have_db_column(:platform).of_type(:string) }
+      it { is_expected.to have_db_column(:port).of_type(:integer) }
+      it { is_expected.to have_db_column(:stype).of_type(:string) }
+      it { is_expected.to have_db_column(:via_exploit).of_type(:string) }
+      it { is_expected.to have_db_column(:via_payload).of_type(:string) }
     end
   end
 
   context 'associations' do
-    it { should belong_to(:host).class_name('Mdm::Host') }
-    it { should have_many(:events).class_name('Mdm::SessionEvent').dependent(:delete_all) }
-    it { should have_many(:routes).class_name('Mdm::Route').dependent(:delete_all) }
-    it { should have_one(:workspace).class_name('Mdm::Workspace').through(:host) }
-    it { should have_many(:task_sessions).class_name('Mdm::TaskSession').dependent(:destroy)}
-    it { should have_many(:tasks).class_name('Mdm::Task').through(:task_sessions)}
+    it { is_expected.to have_many(:events).class_name('Mdm::SessionEvent').dependent(:delete_all) }
+    it { is_expected.to belong_to(:host).class_name('Mdm::Host') }
+    it { is_expected.to belong_to(:originating_module_run).class_name('MetasploitDataModels::ModuleRun') }
+    it { is_expected.to have_many(:routes).class_name('Mdm::Route').dependent(:delete_all) }
+    it { is_expected.to have_many(:target_module_runs).class_name('MetasploitDataModels::ModuleRun') }
+    it { is_expected.to have_many(:tasks).class_name('Mdm::Task').through(:task_sessions)}
+    it { is_expected.to have_many(:task_sessions).class_name('Mdm::TaskSession').dependent(:destroy) }
+    it { is_expected.to have_one(:workspace).class_name('Mdm::Workspace').through(:host) }
   end
 
   context 'scopes' do

spec/app/models/metasploit_data_models/module_run_spec.rb

@@ -22,10 +22,34 @@
 
   context "associations" do
     it { is_expected.to belong_to(:user).class_name('Mdm::User') }
+    it { is_expected.to belong_to(:target_session).class_name('Mdm::Session') }
     it { is_expected.to belong_to(:trackable) }
+    it { is_expected.to have_one(:spawned_session).class_name('Mdm::Session') }
   end
 
   context "validations" do
+    describe "when a session is set on the module run" do
+      before(:each) do
+        module_run.target_session = FactoryGirl.build(:mdm_session)
+      end
+
+      context "when module_name is present" do
+        context "when the module is an exploit" do
+          before(:each){ module_run.module_name = 'exploit/windows/mah-crazy-exploit' }
+
+          it { is_expected.to_not be_valid }
+        end
+      end
+
+      context "when module_detail is present" do
+        before(:each) do
+          module_run.module_detail = FactoryGirl.create(:mdm_module_detail, fullname: 'exploit/windows/some-evil')
+        end
+
+        it { is_expected.to_not be_valid }
+      end
+    end
+
     describe "attempted_at" do
       before(:each){ module_run.attempted_at = nil }
 

spec/dummy/db/structure.sql

@@ -1188,7 +1188,8 @@ CREATE TABLE sessions (
     closed_at timestamp without time zone,
     close_reason character varying(255),
     local_id integer,
-    last_seen timestamp without time zone
+    last_seen timestamp without time zone,
+    module_run_id integer
 );
 
 
@@ -2845,6 +2846,13 @@ CREATE INDEX index_services_on_proto ON services USING btree (proto);
 CREATE INDEX index_services_on_state ON services USING btree (state);
 
 
+--
+-- Name: index_sessions_on_module_run_id; Type: INDEX; Schema: public; Owner: -; Tablespace: 
+--
+
+CREATE INDEX index_sessions_on_module_run_id ON sessions USING btree (module_run_id);
+
+
 --
 -- Name: index_vulns_on_name; Type: INDEX; Schema: public; Owner: -; Tablespace: 
 --
@@ -3134,6 +3142,8 @@ INSERT INTO schema_migrations (version) VALUES ('20150212214222');
 
 INSERT INTO schema_migrations (version) VALUES ('20150219173821');
 
+INSERT INTO schema_migrations (version) VALUES ('20150219215039');
+
 INSERT INTO schema_migrations (version) VALUES ('21');
 
 INSERT INTO schema_migrations (version) VALUES ('22');