Skip to content

Commit

Permalink
feat: introduce UnenrichableError
Browse files Browse the repository at this point in the history
  • Loading branch information
ninoseki committed Jan 14, 2024
1 parent d19baf5 commit 6f6a9d2
Show file tree
Hide file tree
Showing 8 changed files with 186 additions and 45 deletions.
2 changes: 2 additions & 0 deletions lib/mihari/errors.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
71 changes: 39 additions & 32 deletions lib/mihari/models/artifact.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
2 changes: 2 additions & 0 deletions lib/mihari/services/enrichers.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 3 additions & 1 deletion lib/mihari/web/endpoints/artifacts.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

24 changes: 19 additions & 5 deletions spec/services/artifact_enricher_spec.rb
Original file line number Diff line number Diff line change
@@ -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
3 changes: 1 addition & 2 deletions spec/sidekiq/artifact_enrich_job_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
25 changes: 20 additions & 5 deletions spec/web/endpoints/artifacts_spec.rb
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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

Expand Down

0 comments on commit 6f6a9d2

Please sign in to comment.