artifact.rb in mihari-7.0.0

- old
+ new

@@ -20,78 +20,68 @@
       belongs_to :alert
 
       has_one :autonomous_system, dependent: :destroy
       has_one :geolocation, dependent: :destroy
       has_one :whois_record, dependent: :destroy
+      has_one :rule, through: :alert
 
       has_many :cpes, dependent: :destroy
       has_many :dns_records, dependent: :destroy
       has_many :ports, dependent: :destroy
       has_many :reverse_dns_names, dependent: :destroy
+      has_many :tags, through: :alert
 
       include ActiveModel::Validations
+      include SearchCop
+      include Concerns::Searchable
 
+      search_scope :search do
+        attributes :id, :data, :data_type, :source, :query, :created_at, "alert.id", "rule.id", "rule.title",
+          "rule.description"
+        attributes tag: "tags.name"
+        attributes asn: "autonomous_system.asn"
+        attributes country_code: "geolocation.country_code"
+        attributes "dns_record.value": "dns_records.value"
+        attributes "dns_record.resource": "dns_records.resource"
+        attributes reverse_dns_name: "reverse_dns_names.name"
+        attributes cpe: "cpes.name"
+        attributes port: "ports.port"
+      end
+
       validates_with ArtifactValidator
 
-      # @return [Array<Mihari::Tag>] Tags
-      attr_accessor :tags
+      after_initialize :set_data_type, :set_rule_id, if: :new_record?
 
-      # @return [String, nil] Rule ID
+      # @return [String, nil]
       attr_accessor :rule_id
 
-      def initialize(*args, **kwargs)
-        attrs = args.first || kwargs
-        data_ = attrs[:data]
-
-        raise TypeError if data_.is_a?(Array) || data_.is_a?(Hash)
-
-        super(*args, **kwargs)
-
-        self.data_type = DataType.type(data)
-
-        @tags = []
-        @rule_id = ""
-      end
-
       #
-      # Check uniqueness of artifact
+      # Check uniqueness
       #
       # @param [Time, nil] base_time Base time to check decaying
-      # @param [Integer, nil] artifact_lifetime Artifact lifetime (TTL) in seconds
+      # @param [Integer, nil] artifact_ttl Artifact TTL in seconds
       #
       # @return [Boolean] true if it is unique. Otherwise false.
       #
-      def unique?(base_time: nil, artifact_lifetime: nil)
+      def unique?(base_time: nil, artifact_ttl: nil)
         artifact = self.class.joins(:alert).where(
           data: data,
           alert: { rule_id: rule_id }
         ).order(created_at: :desc).first
         return true if artifact.nil?
 
         # check whether the artifact is decayed or not
-        return false if artifact_lifetime.nil?
+        return false if artifact_ttl.nil?
 
         # use the current UTC time if base_time is not given (for testing)
         base_time ||= Time.now.utc
 
-        decayed_at = base_time - (artifact_lifetime || -1).seconds
+        decayed_at = base_time - (artifact_ttl || -1).seconds
         artifact.created_at < decayed_at
       end
 
       #
-      # Count artifacts
-      #
-      # @param [Mihari::Structs::Filters::Artifact::SearchFilter] filter
-      #
-      # @return [Integer]
-      #
-      def count(filter)
-        relation = build_relation(filter)
-        relation.distinct("artifact.id").count
-      end
-
-      #
       # Enrich whois record
       #
       # @param [Mihari::Enrichers::Whois] enricher
       #
       def enrich_whois(enricher = Enrichers::Whois.new)
@@ -125,22 +115,22 @@
       #
       # Enrich geolocation
       #
       # @param [Mihari::Enrichers::IPInfo] enricher
       #
-      def enrich_geolocation(enricher = Enrichers::IPInfo.new)
+      def enrich_geolocation(enricher = Enrichers::MMDB.new)
         return unless can_enrich_geolocation?
 
         self.geolocation = Geolocation.build_by_ip(data, enricher: enricher)
       end
 
       #
       # Enrich AS
       #
       # @param [Mihari::Enrichers::IPInfo] enricher
       #
-      def enrich_autonomous_system(enricher = Enrichers::IPInfo.new)
+      def enrich_autonomous_system(enricher = Enrichers::MMDB.new)
         return unless can_enrich_autonomous_system?
 
         self.autonomous_system = AutonomousSystem.build_by_ip(data, enricher: enricher)
       end
 
@@ -168,24 +158,24 @@
 
       #
       # Enrich all the enrichable relationships of the artifact
       #
       def enrich_all
-        enrich_autonomous_system ipinfo
+        enrich_autonomous_system mmdb
         enrich_dns
-        enrich_geolocation ipinfo
+        enrich_geolocation mmdb
         enrich_reverse_dns shodan
         enrich_whois
         enrich_ports shodan
         enrich_cpes shodan
       end
 
       ENRICH_METHODS_BY_ENRICHER = {
         Enrichers::Whois => %i[
           enrich_whois
         ],
-        Enrichers::IPInfo => %i[
+        Enrichers::MMDB => %i[
           enrich_autonomous_system
           enrich_geolocation
         ],
         Enrichers::Shodan => %i[
           enrich_ports
@@ -206,62 +196,30 @@
         methods = ENRICH_METHODS_BY_ENRICHER[enricher.class] || []
         methods.each { |method| send(method, enricher) if respond_to?(method) }
       end
 
       class << self
-        #
-        # Search artifacts
-        #
-        # @param [Mihari::Structs::Filters::Artifact::SearchFilterWithPagination] filter
-        #
-        # @return [Array<Artifact>]
-        #
-        def search(filter)
-          limit = filter.limit.to_i
-          raise ArgumentError, "limit should be bigger than zero" unless limit.positive?
+        # @!method search_by_filter(filter)
+        #   @param [Mihari::Structs::Filters::Search] filter
+        #   @return [Array<Mihari::Models::Alert>]
 
-          page = filter.page.to_i
-          raise ArgumentError, "page should be bigger than zero" unless page.positive?
+        # @!method count_by_filter(filter)
+        #   @param [Mihari::Structs::Filters::Search] filter
+        #   @return [Integer]
+      end
 
-          offset = (page - 1) * limit
+      private
 
-          relation = build_relation(filter.without_pagination)
-          relation.limit(limit).offset(offset).order(id: :desc)
-        end
+      def set_data_type
+        self.data_type = DataType.type(data)
+      end
 
-        #
-        # Count artifacts
-        #
-        # @param [Mihari::Structs::Filters::Artifact::SearchFilter] filter
-        #
-        # @return [Integer]
-        #
-        def count(filter)
-          relation = build_relation(filter)
-          relation.distinct("artifacts.id").count
-        end
-
-        #
-        # @param [Mihari::Structs::Filters::Artifact::SearchFilter] filter
-        #
-        # @return [Mihari::Models::Artifact]
-        #
-        def build_relation(filter)
-          relation = eager_load(alert: :tags)
-
-          relation = relation.where(alert: { rule_id: filter.rule_id }) if filter.rule_id
-          relation = relation.where(alert: { tags: { name: filter.tag } }) if filter.tag
-          relation = relation.where("artifacts.created_at >= ?", filter.from_at) if filter.from_at
-          relation = relation.where("artifacts.created_at <= ?", filter.to_at) if filter.to_at
-
-          relation
-        end
+      def set_rule_id
+        @set_rule_id ||= nil
       end
 
-      private
-
-      def ipinfo
-        @ipinfo ||= Enrichers::IPInfo.new
+      def mmdb
+        @mmdb ||= Enrichers::MMDB.new
       end
 
       def shodan
         @shodan ||= Enrichers::Shodan.new
       end