Update report_vuln to support tracking origin

lib/msf/core/db_manager/service.rb

@@ -48,6 +48,7 @@ def find_or_create_service(opts)
   # +:info+:: Detailed information about the service such as name and version information
   # +:state+:: The current listening state of the service (one of: open, closed, filtered, unknown)
   #
+  # @return [Mdm::Service,nil]
   def report_service(opts)
     return if !active
   ::ApplicationRecord.connection_pool.with_connection { |conn|
@@ -81,8 +82,6 @@ def report_service(opts)
       return nil
     end
 
-    ret  = {}
-
     proto = opts[:proto] || Msf::DBManager::DEFAULT_SERVICE_PROTO
 
     service = host.services.where(port: opts[:port].to_i, proto: proto).first_or_initialize
@@ -116,13 +115,13 @@ def report_service(opts)
     end
 
     if opts[:task]
-      Mdm::TaskService.create(
+      Mdm::TaskService.where(
           :task => opts[:task],
           :service => service
-      )
+      ).first_or_create
     end
 
-    ret[:service] = service
+    service
   }
   end
 

lib/msf/core/db_manager/vuln.rb

@@ -223,6 +223,8 @@ def report_vuln(opts)
     # Set the exploited_at value if provided
     vuln.exploited_at = exploited_at if exploited_at
 
+    vuln.origin = opts[:origin] if opts[:origin]
+
     # Merge the references
     if rids
       vuln.refs << (rids - vuln.refs)

spec/support/shared/examples/msf/db_manager/service.rb

@@ -8,4 +8,61 @@
   it { is_expected.to respond_to :find_or_create_service }
   it { is_expected.to respond_to :services }
   it { is_expected.to respond_to :report_service }
+
+  describe '#report_service', if: !ENV['REMOVE_DB'] do
+    let(:workspace) do
+      subject.default_workspace
+    end
+
+    let(:task) do
+      subject.report_task(workspace: workspace, user: 'test_user', info: 'info', path: 'mock/path')
+    end
+
+    context 'without a task' do
+      it 'creates a service' do
+        service = subject.report_service(
+          host: '192.0.2.1',
+          port: '5000',
+          name: 'test_service',
+          proto: 'tcp',
+          info: 'banner',
+          workspace: workspace
+        )
+        expect(subject.services({ workspace: workspace }).count).to eq 1
+        expect(service.name).to eq 'test_service'
+        expect(service.port).to eq 5000
+        expect(service.proto).to eq 'tcp'
+        expect(service.info).to eq 'banner'
+        expect(service.host.address.to_s).to eq '192.0.2.1'
+        expect(service.host.workspace).to eq workspace
+        expect(service.task_services).to be_empty
+        expect(task.task_services).to be_empty
+      end
+    end
+
+    context 'with a task and calling multiple times' do
+      it 'creates a service' do
+        service = 3.times.map do |count|
+          subject.report_service(
+            host: '192.0.2.1',
+            port: '5000',
+            name: 'test_service',
+            proto: 'tcp',
+            info: "banner #{count}",
+            workspace: workspace,
+            task: task
+          )
+        end.last
+        expect(subject.services({ workspace: workspace }).count).to eq 1
+        expect(service.name).to eq 'test_service'
+        expect(service.port).to eq 5000
+        expect(service.proto).to eq 'tcp'
+        expect(service.info).to eq 'banner 2'
+        expect(service.host.address.to_s).to eq '192.0.2.1'
+        expect(service.host.workspace).to eq workspace
+        expect(service.task_services.length).to eq 1
+        expect(task.task_services.length).to eq 1
+      end
+    end
+  end
 end

spec/support/shared/examples/msf/db_manager/vuln.rb

@@ -9,6 +9,74 @@
     it { is_expected.to respond_to :find_vuln_by_details }
   end
 
-  it { is_expected.to respond_to :report_vuln }
+  describe '#report_vuln', if: !ENV['REMOTE_DB'] do
+    let(:workspace) do
+      subject.default_workspace
+    end
+
+    let(:task) do
+      subject.report_task(workspace: workspace, user: 'test_user', info: 'info', path: 'mock/path')
+    end
+
+    let(:service) do
+      subject.report_service(
+        host: '192.0.2.1',
+        port: '5000',
+        name: 'test_service',
+        proto: 'tcp',
+        info: 'banner',
+        workspace: workspace
+      )
+    end
+
+    context 'without an origin' do
+      it 'creates a vuln' do
+        vuln = subject.report_vuln(
+          host: '192.0.2.1',
+          sname: 'AD CS',
+          name: "vuln name",
+          info: 'vuln info',
+          refs: ['https://example.com'],
+          workspace: workspace,
+          service: service,
+          origin: task
+        )
+        expect(subject.vulns({ workspace: workspace }).count).to eq 1
+        expect(vuln.name).to eq 'vuln name'
+        expect(vuln.service.name).to eq 'test_service'
+        expect(vuln.service.port).to eq 5000
+        expect(vuln.info).to eq 'vuln info'
+        expect(vuln.host.address.to_s).to eq '192.0.2.1'
+        expect(vuln.host.workspace).to eq workspace
+        expect(service.task_services).to be_empty
+      end
+    end
+
+    context 'with a task origin and calling multiple times' do
+      it 'creates a service' do
+        vuln = 3.times.map do |count|
+          subject.report_vuln(
+            host: '192.0.2.1',
+            sname: 'AD CS',
+            name: "vuln name",
+            info: 'vuln info',
+            refs: ['https://example.com'],
+            workspace: workspace,
+            service: service,
+            origin: task
+          )
+        end.last
+        expect(subject.services({ workspace: workspace }).count).to eq 1
+        expect(vuln.name).to eq 'vuln name'
+        expect(vuln.service.name).to eq 'test_service'
+        expect(vuln.service.port).to eq 5000
+        expect(vuln.info).to eq 'vuln info'
+        expect(vuln.host.address.to_s).to eq '192.0.2.1'
+        expect(vuln.host.workspace).to eq workspace
+        expect(service.task_services).to be_empty
+      end
+    end
+  end
+
   it { is_expected.to respond_to :vulns }
-end
+end