Skip to content

Commit

Permalink
refactor: renew enrichers
Browse files Browse the repository at this point in the history
  • Loading branch information
ninoseki committed Jan 20, 2024
1 parent a20f696 commit 2c6966d
Show file tree
Hide file tree
Showing 27 changed files with 328 additions and 673 deletions.
66 changes: 54 additions & 12 deletions lib/mihari/enrichers/base.rb
Original file line number Diff line number Diff line change
Expand Up @@ -6,46 +6,88 @@ module Enrichers
# Base class for enrichers
#
class Base < Actor
prepend MemoWise

#
# @param [Hash, nil] options
#
def initialize(options: nil)
super(options: options)
end

#
# @param [String] value
# Enrich an artifact
#
# @param [Mihari::Models::Artifact] artifact
#
def call(value)
# @return [Mihari::Models::Artifact]
#
def call(artifact)
raise NotImplementedError, "You must implement #{self.class}##{__method__}"
end

#
# @param [Mihari::Models::Artifact] value
# @param [Mihari::Models::Artifact] artifact
#
# @return [Dry::Monads::Result::Success<Object>, Dry::Monads::Result::Failure]
#
def result(value)
def result(artifact)
return unless callable?(artifact)

result = Try[StandardError] do
retry_on_error(
times: retry_times,
interval: retry_interval,
exponential_backoff: retry_exponential_backoff
) { call value }
retry_on_error(times: retry_times, interval: retry_interval,
exponential_backoff: retry_exponential_backoff) do
call artifact
end
end.to_result

if result.failure?
Mihari.logger.warn("Enricher:#{self.class.key} for #{value.truncate(32)} failed: #{result.failure}")
Mihari.logger.warn("Enricher:#{self.class.key} for #{artifact.data.truncate(32)} failed: #{result.failure}")
end

result
end

#
# @param [Mihari::Models::Artifact] artifact
#
# @return [Boolean]
#
def callable?(artifact)
callable_data_type?(artifact) && callable_relationships?(artifact)
end

class << self
def inherited(child)
super
Mihari.enrichers << child
end
end

private

#
# @param [Mihari::Models::Artifact] artifact
#
# @return [Boolean]
#
def callable_data_type?(artifact)
supported_data_types.include? artifact.data_type
end

#
# @param [Mihari::Models::Artifact] artifact
#
# @return [Boolean]
#
def callable_relationships?(artifact)
raise NotImplementedError, "You must implement #{self.class}##{__method__}"
end

#
# @return [Array<String>]
#
def supported_data_types
[]
end
end
end
end
33 changes: 26 additions & 7 deletions lib/mihari/enrichers/google_public_dns.rb
Original file line number Diff line number Diff line change
Expand Up @@ -7,14 +7,20 @@ module Enrichers
#
class GooglePublicDNS < Base
#
# Query Google Public DNS
# @param [Mihari::Models::Artifact] artifact
#
# @param [String] name
# @return [Mihari::Models::Artifact]
#
# @return [Mihari::Structs::GooglePublicDNS::Response]
#
def call(name)
client.query_all name
def call(artifact)
res = client.query_all(artifact.domain)

artifact.tap do |tapped|
if tapped.dns_records.empty?
tapped.dns_records = res.answers.map do |answer|
Models::DnsRecord.new(resource: answer.resource_type, value: answer.data)
end
end
end
end

class << self
Expand All @@ -28,8 +34,21 @@ def key

private

#
# @param [Mihari::Models::Artifact] artifact
#
# @return [Boolean]
#
def callable_relationships?(artifact)
artifact.dns_records.empty?
end

def supported_data_types
%w[url domain]
end

def client
Clients::GooglePublicDNS.new(timeout: timeout)
@client ||= Clients::GooglePublicDNS.new(timeout: timeout)
end
end
end
Expand Down
32 changes: 25 additions & 7 deletions lib/mihari/enrichers/mmdb.rb
Original file line number Diff line number Diff line change
Expand Up @@ -7,18 +7,36 @@ module Enrichers
#
class MMDB < Base
#
# Query MMDB
# @param [Mihari::Models::Artifact] artifact
#
# @param [String] ip
def call(artifact)
res = client.query(artifact.data)

artifact.tap do |tapped|
tapped.autonomous_system ||= Models::AutonomousSystem.new(number: res.asn) if res.asn
if res.country_code
tapped.geolocation ||= Models::Geolocation.new(
country: NormalizeCountry(res.country_code, to: :short),
country_code: res.country_code
)
end
end
end

private

#
# @param [Mihari::Models::Artifact] artifact
#
# @return [Mihari::Structs::MMDB::Response]
# @return [Boolean]
#
def call(ip)
client.query ip
def callable_relationships?(artifact)
artifact.geolocation.nil? || artifact.autonomous_system.nil?
end
memo_wise :call

private
def supported_data_types
%w[ip]
end

def client
@client ||= Clients::MMDB.new(timeout: timeout)
Expand Down
39 changes: 35 additions & 4 deletions lib/mihari/enrichers/shodan.rb
Original file line number Diff line number Diff line change
Expand Up @@ -9,17 +9,48 @@ class Shodan < Base
#
# Query Shodan Internet DB
#
# @param [String] ip
# @param [Mihari::Models::Artifact] artifact
#
# @return [Mihari::Structs::Shodan::InternetDBResponse, nil]
#
def call(ip)
client.query ip
def call(artifact)
res = client.query(artifact.data)

artifact.tap do |tapped|
tapped.cpes = (res&.cpes || []).map { |cpe| Models::CPE.new(name: cpe) } if tapped.cpes.empty?
tapped.ports = (res&.ports || []).map { |port| Models::Port.new(number: port) } if tapped.ports.empty?
if tapped.reverse_dns_names.empty?
tapped.reverse_dns_names = (res&.hostnames || []).map do |name|
Models::ReverseDnsName.new(name: name)
end
end
end
end

#
# @param [Mihari::Models::Artifact] artifact
#
# @return [Boolean]
#
def callable?(artifact)
false unless supported_data_types.include?(artifact.data_type)
end
memo_wise :call

private

#
# @param [Mihari::Models::Artifact] artifact
#
# @return [Boolean]
#
def callable_relationships?(artifact)
artifact.cpes.empty? || artifact.ports.empty? || artifact.reverse_dns_names.empty?
end

def supported_data_types
%w[ip]
end

def client
@client ||= Clients::ShodanInternetDB.new(timeout: timeout)
end
Expand Down
26 changes: 17 additions & 9 deletions lib/mihari/enrichers/whois.rb
Original file line number Diff line number Diff line change
Expand Up @@ -8,26 +8,34 @@ module Enrichers
# Whois enricher
#
class Whois < Base
#
# @param [Hash, nil] options
#
def initialize(options: nil)
super(options: options)
end
prepend MemoWise

#
# Query IAIA Whois API
#
# @param [String] domain
# @param [Mihari::Models::Artifact] artifact
#
# @return [Mihari::Models::WhoisRecord, nil]
#
def call(domain)
memoized_call PublicSuffix.domain(domain)
def call(artifact)
artifact.whois_record ||= memoized_call(PublicSuffix.domain(artifact.domain))
end

private

#
# @param [Mihari::Models::Artifact] artifact
#
# @return [Boolean]
#
def callable_relationships?(artifact)
artifact.whois_record.nil?
end

def supported_data_types
%w[url domain]
end

#
# @param [String] domain
#
Expand Down
12 changes: 12 additions & 0 deletions lib/mihari/models/alert.rb
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,18 @@ module Models
# Alert model
#
class Alert < ActiveRecord::Base
# @!attribute [r] id
# @return [Integer, nil]

# @!attribute [rw] created_at
# @return [DateTime]

# @!attribute [r] rule
# @return [Mihari::Models::Rule]

# @!attribute [r] artifacts
# @return [Array<Mihari::Models::Artifact>]

belongs_to :rule

has_many :artifacts, dependent: :destroy
Expand Down
Loading

0 comments on commit 2c6966d

Please sign in to comment.