Merge pull request #1036 from ninoseki/renew-enrichers

refactor: renew enrichers

lib/mihari/enrichers/base.rb

@@ -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

lib/mihari/enrichers/google_public_dns.rb

@@ -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
@@ -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

lib/mihari/enrichers/mmdb.rb

@@ -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)

lib/mihari/enrichers/shodan.rb

@@ -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

lib/mihari/enrichers/whois.rb

@@ -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
       #

lib/mihari/models/alert.rb

@@ -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