Skip to content

Commit

Permalink
feat: improve enrichment parallelism (#1108)
Browse files Browse the repository at this point in the history
  • Loading branch information
ninoseki authored Aug 10, 2024
1 parent 30cbef2 commit 136c9c5
Show file tree
Hide file tree
Showing 3 changed files with 46 additions and 11 deletions.
4 changes: 3 additions & 1 deletion lib/mihari/enrichers/whois.rb
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,9 @@ class Whois < Base
def call(artifact)
return if artifact.domain.nil?

artifact.whois_record ||= memoized_lookup(PublicSuffix.domain(artifact.domain))
artifact.tap do |tapped|
tapped.whois_record ||= memoized_lookup(PublicSuffix.domain(artifact.domain))
end
end

private
Expand Down
48 changes: 42 additions & 6 deletions lib/mihari/models/artifact.rb
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
# frozen_string_literal: true

require "ostruct"

module Mihari
module Models
#
Expand Down Expand Up @@ -158,10 +160,7 @@ class Artifact < ActiveRecord::Base
# @return [Boolean] true if it is unique. Otherwise false.
#
def unique?(base_time: nil, artifact_ttl: nil)
artifact = self.class.joins(:alert).where(
data:,
alert: {rule_id:}
).order(created_at: :desc).first
artifact = self.class.joins(:alert).where(data:, alert: {rule_id:}).order(created_at: :desc).first
return true if artifact.nil?

# check whether the artifact is decayed or not
Expand All @@ -179,7 +178,32 @@ def enrichable?
end

def enrich
callable_enrichers.each { |enricher| enricher.result self }
enrich_by_enrichers callable_enrichers
end

#
# @param [Array<Mihari::Enrichers::Base>] enrichers
# @param [Boolean] parallel
#
# @return [Mihari::Models::Artifact]
#
def enrich_by_enrichers(enrichers)
# NOTE: doing parallel with ActiveRecord objects is troublesome (e.g. connection issue, etc.)
# so converting the object to an OpenStruct object
s = struct
results = Parallel.map(enrichers) { |enricher| enricher.result s }
enriched = results.compact.map { |result| result.value_or(nil) }.compact

self.dns_records = enriched.map(&:dns_records).flatten.compact
self.cpes = enriched.map(&:cpes).flatten.compact
self.ports = enriched.map(&:ports).flatten.compact
self.vulnerabilities = enriched.map(&:vulnerabilities).flatten.compact

self.autonomous_system = enriched.map(&:autonomous_system).compact.first
self.geolocation = enriched.map(&:geolocation).compact.first
self.whois_record = enriched.map(&:whois_record).compact.first

self
end

#
Expand All @@ -195,6 +219,18 @@ def domain
end
end

def struct
OpenStruct.new(attributes).tap do |s|
s.domain = domain
s.cpes ||= []
s.dns_records ||= []
s.ports ||= []
s.reverse_dns_names ||= []
s.vulnerabilities ||= []
s.tags ||= []
end
end

class << self
# @!method search_by_filter(filter)
# @param [Mihari::Structs::Filters::Search] filter
Expand All @@ -212,7 +248,7 @@ class << self
#
def callable_enrichers
@callable_enrichers ||= Mihari.enrichers.map(&:new).select do |enricher|
enricher.callable?(self)
enricher.callable? self
end
end

Expand Down
5 changes: 1 addition & 4 deletions lib/mihari/rule.rb
Original file line number Diff line number Diff line change
Expand Up @@ -178,10 +178,7 @@ def unique_artifacts
#
def enriched_artifacts
@enriched_artifacts ||= Parallel.map(unique_artifacts) do |artifact|
artifact.tap do |tapped|
# NOTE: To apply changes correctly, enrichers should be applied to an artifact serially
enrichers.each { |enricher| enricher.result(tapped) }
end
artifact.enrich_by_enrichers enrichers
end
end

Expand Down

0 comments on commit 136c9c5

Please sign in to comment.