Merge pull request #1027 from ninoseki/uneinrichable-error

feat: introduce UnenrichableError

lib/mihari/errors.rb

@@ -7,6 +7,8 @@ class Error < StandardError; end
 
   class ValueError < Error; end
 
+  class UnenrichableError < Error; end
+
   class ConfigurationError < Error
     # @return [Array<String>, nil]
     attr_reader :detail

lib/mihari/models/artifact.rb

@@ -221,6 +221,45 @@ def enrich_by_enricher(enricher)
         methods.each { |method| send(method, enricher) if respond_to?(method) }
       end
 
+      def can_enrich_whois?
+        %w[domain url].include?(data_type) && whois_record.nil?
+      end
+
+      def can_enrich_dns?
+        %w[domain url].include?(data_type) && dns_records.empty?
+      end
+
+      def can_enrich_reverse_dns?
+        data_type == "ip" && reverse_dns_names.empty?
+      end
+
+      def can_enrich_geolocation?
+        data_type == "ip" && geolocation.nil?
+      end
+
+      def can_enrich_autonomous_system?
+        data_type == "ip" && autonomous_system.nil?
+      end
+
+      def can_enrich_ports?
+        data_type == "ip" && ports.empty?
+      end
+
+      def can_enrich_cpes?
+        data_type == "ip" && cpes.empty?
+      end
+
+      def can_enrich_vulnerabilities?
+        data_type == "ip" && vulnerabilities.empty?
+      end
+
+      def enrichable?
+        enrich_methods = methods.map(&:to_s).select { |method| method.start_with?("can_enrich_") }
+        enrich_methods.map(&:to_sym).any? do |method|
+          send(method) if respond_to?(method)
+        end
+      end
+
       class << self
         # @!method search_by_filter(filter)
         #   @param [Mihari::Structs::Filters::Search] filter
@@ -260,38 +299,6 @@ def domain
           Addressable::URI.parse(data).host
         end
       end
-
-      def can_enrich_whois?
-        %w[domain url].include?(data_type) && whois_record.nil?
-      end
-
-      def can_enrich_dns?
-        %w[domain url].include?(data_type) && dns_records.empty?
-      end
-
-      def can_enrich_reverse_dns?
-        data_type == "ip" && reverse_dns_names.empty?
-      end
-
-      def can_enrich_geolocation?
-        data_type == "ip" && geolocation.nil?
-      end
-
-      def can_enrich_autonomous_system?
-        data_type == "ip" && autonomous_system.nil?
-      end
-
-      def can_enrich_ports?
-        data_type == "ip" && ports.empty?
-      end
-
-      def can_enrich_cpes?
-        data_type == "ip" && cpes.empty?
-      end
-
-      def can_enrich_vulnerabilities?
-        data_type == "ip" && vulnerabilities.empty?
-      end
     end
   end
 end

lib/mihari/services/enrichers.rb

@@ -17,6 +17,8 @@ def call(id)
           :ports
         ).find(id)
 
+        raise UnenrichableError.new, "#{artifact.id} is already enriched or unenrichable" unless artifact.enrichable?
+
         artifact.enrich_all
         artifact.save
       end

lib/mihari/web/endpoints/artifacts.rb

@@ -53,7 +53,7 @@ class Artifacts < Grape::API
 
           desc "Enrich an artifact", {
             success: { code: 201, model: Entities::Message },
-            failure: [{ code: 404, model: Entities::ErrorMessage }],
+            failure: [{ code: 400, model: Entities::ErrorMessage }, { code: 404, model: Entities::ErrorMessage }],
             summary: "Enrich an artifact"
           }
           params do
@@ -78,6 +78,8 @@ class Artifacts < Grape::API
             return present({ message: message, queued: queued }, with: Entities::QueueMessage) if result.success?
 
             case result.failure
+            when UnenrichableError
+              error!({ message: result.failure.message }, 400)
             when ActiveRecord::RecordNotFound
               error!({ message: "ID:#{id} not found" }, 404)
             end

spec/services/artifact_enricher_spec.rb

@@ -1,13 +1,27 @@
 # frozen_string_literal: true
 
-RSpec.describe Mihari::Services::ArtifactEnricher do
+RSpec.describe Mihari::Services::ArtifactEnricher, vcr: "Mihari_Services_ArtifactEnricher/ip:1.1.1.1" do
   describe ".call" do
-    let!(:artifact) do
-      FactoryBot.create(:artifact, :unenrichable)
+    context "with enrichable artifact" do
+      let(:artifact) do
+        artifact = FactoryBot.build(:artifact, :ip).tap { |tapped| tapped.data = "1.1.1.1" }
+        artifact.save
+        artifact
+      end
+
+      it do
+        expect(described_class.call(artifact.id)).to eq(true)
+      end
     end
 
-    it do
-      expect(described_class.call(artifact.id)).to eq(true)
+    context "with unenrichable artifact" do
+      let(:artifact) do
+        FactoryBot.create(:artifact, :unenrichable)
+      end
+
+      it do
+        expect { described_class.call(artifact.id) }.to raise_error(Mihari::UnenrichableError)
+      end
     end
   end
 end

spec/sidekiq/artifact_enrich_job_spec.rb

@@ -7,8 +7,7 @@
 
   it do
     Sidekiq::Testing.inline! do
-      res = described_class.perform_async(artifact.id)
-      expect(res).not_to be nil
+      expect { described_class.perform_async(artifact.id) }.to raise_error(Mihari::UnenrichableError)
     end
   end
 end

spec/web/endpoints/artifacts_spec.rb

@@ -1,9 +1,13 @@
 # frozen_string_literal: true
 
-RSpec.describe Mihari::Web::Endpoints::Artifacts do
+RSpec.describe Mihari::Web::Endpoints::Artifacts, vcr: "Mihari_Services_ArtifactEnricher/ip:1.1.1.1" do
   include Rack::Test::Methods
 
-  let_it_be(:artifact) { FactoryBot.create(:artifact, :unenrichable) }
+  let_it_be(:artifact) do
+    artifact = FactoryBot.build(:artifact, :ip).tap { |tapped| tapped.data = "1.1.1.1" }
+    artifact.save
+    artifact
+  end
   let_it_be(:artifact_to_delete) { FactoryBot.create(:artifact) }
 
   def app
@@ -42,9 +46,20 @@ def app
   end
 
   describe "get /api/artifacts/:id/enrich" do
-    it "returns 201" do
-      post "/api/artifacts/#{artifact.id}/enrich"
-      expect(last_response.status).to eq(201)
+    context "with enrichable artifact" do
+      it "returns 201" do
+        post "/api/artifacts/#{artifact.id}/enrich"
+        expect(last_response.status).to eq(201)
+      end
+    end
+
+    context "with unenrichable artifact" do
+      let(:artifact) { FactoryBot.create(:artifact, :unenrichable) }
+
+      it "returns 400" do
+        post "/api/artifacts/#{artifact.id}/enrich"
+        expect(last_response.status).to eq(400)
+      end
     end
   end