miteru 1.2.2 → 2.0.1

data/lib/miteru/database.rb

@@ -1,11 +1,9 @@
 # frozen_string_literal: true
 
-require "active_record"
-
-class InitialSchema < ActiveRecord::Migration[7.0]
+class V2Schema < ActiveRecord::Migration[7.0]
   def change
     create_table :records, if_not_exists: true do |t|
-      t.string :hash, null: false, index: { unique: true }
+      t.string :sha256, null: false, index: {unique: true}
       t.string :hostname, null: false
       t.json :headers, null: false
       t.text :filename, null: false
@@ -13,61 +11,75 @@ class InitialSchema < ActiveRecord::Migration[7.0]
       t.integer :filesize, null: false
       t.string :mime_type, null: false
       t.text :url, null: false
+      t.string :source, null: false
 
       t.timestamps
     end
   end
 end
 
-class V11Schema < ActiveRecord::Migration[7.0]
-  def change
-    add_column :records, :source, :string, if_not_exists: true
-  end
-end
-
-def adapter
-  return "postgresql" if Miteru.configuration.database.start_with?("postgresql://", "postgres://")
-  return "mysql2" if Miteru.configuration.database.start_with?("mysql2://")
-
-  "sqlite3"
+#
+# @return [Array<ActiveRecord::Migration>] schemas
+#
+def schemas
+  [V2Schema]
 end
 
 module Miteru
   class Database
     class << self
+      #
+      # DB migration
+      #
+      # @param [Symbol] direction
+      #
+      def migrate(direction)
+        schemas.each { |schema| schema.migrate direction }
+      end
+
+      #
+      # Establish DB connection
+      #
       def connect
-        case adapter
-        when "postgresql", "mysql2"
-          ActiveRecord::Base.establish_connection(Miteru.configuration.database)
-        else
-          ActiveRecord::Base.establish_connection(
-            adapter: adapter,
-            database: Miteru.configuration.database
-          )
-        end
+        return if connected?
 
-        # ActiveRecord::Base.logger = Logger.new STDOUT
-        ActiveRecord::Migration.verbose = false
+        ActiveRecord::Base.establish_connection Miteru.config.database_url.to_s
+        ActiveRecord::Base.logger = Logger.new($stdout) if Miteru.development?
+      end
 
-        InitialSchema.migrate(:up)
-        V11Schema.migrate(:up)
-      rescue StandardError => _e
-        # Do nothing
+      #
+      # @return [Boolean]
+      #
+      def connected?
+        ActiveRecord::Base.connected?
       end
 
+      #
+      # Close DB connection(s)
+      #
       def close
-        ActiveRecord::Base.clear_active_connections!
-        ActiveRecord::Base.connection.close
+        return unless connected?
+
+        ActiveRecord::Base.connection_handler.clear_active_connections!
       end
 
-      def destroy!
-        return unless ActiveRecord::Base.connected?
+      def with_db_connection
+        Miteru::Database.connect unless connected?
+        yield
+      rescue ActiveRecord::StatementInvalid
+        Miteru.logger.error("DB migration is not yet complete. Please run 'miteru db migrate'.")
+      ensure
+        Miteru::Database.close
+      end
 
-        InitialSchema.migrate(:down)
-        V11Schema.migrate(:down)
+      private
+
+      def adapter
+        return "postgresql" if %w[postgresql postgres].include?(Miteru.config.database_url.scheme)
+        return "mysql2" if Miteru.config.database_url.scheme == "mysql2"
+
+        "sqlite3"
       end
     end
   end
 end
-
-Miteru::Database.connect

data/lib/miteru/downloader.rb

@@ -2,65 +2,76 @@
 
 require "digest"
 require "fileutils"
-require "uri"
+
+require "down/http"
 
 module Miteru
-  class Downloader
-    attr_reader :base_dir, :memo
+  class Downloader < Service
+    prepend MemoWise
 
-    def initialize(base_dir = "/tmp")
-      @base_dir = base_dir
-      @memo = {}
-      raise ArgumentError, "#{base_dir} doesn't exist." unless Dir.exist?(base_dir)
-    end
+    # @return [String]
+    attr_reader :base_dir
 
-    def download_kits(kits)
-      kits.each { |kit| download_kit kit }
-    end
+    # @return [Miteru::Kit]
+    attr_reader :kit
 
-    private
+    #
+    # <Description>
+    #
+    # @param [Miteru::Kit] kit
+    # @param [String] base_dir
+    #
+    def initialize(kit, base_dir: Miteru.config.download_to)
+      super()
+      @kit = kit
+      @base_dir = base_dir
+    end
 
-    def download_kit(kit)
+    #
+    # @return [String]
+    #
+    def call
       destination = kit.filepath_to_download
 
-      begin
-        downloaded_as = HTTPClient.download(kit.url, destination)
-      rescue Down::Error => e
-        Miteru.logger.error "Failed to download: #{kit.url} (#{e})"
-        return
-      end
+      # downloader.download(kit.url, destination:, max_size:)
+      downloader.download(kit.url, destination:, max_size:)
 
-      # check filesize
-      size = File.size downloaded_as
-      if size > Miteru.configuration.file_maxsize
-        Miteru.logger.info "#{kit.url}'s filesize exceeds the limit: #{size}"
-        FileUtils.rm downloaded_as
-        return
+      unless Record.unique_sha256?(sha256(destination))
+        FileUtils.rm destination
+        raise UniquenessError, "Kit:#{sha256(destination)} is registered already."
       end
 
-      hash = sha256(downloaded_as)
+      # Record a kit in DB
+      Record.create_by_kit_and_hash(kit, sha256: sha256(destination))
+      Miteru.logger.info "Download #{kit.url} as #{destination}"
+
+      destination
+    end
 
-      ActiveRecord::Base.connection_pool.with_connection do
-        # Remove a downloaded file if it is not unique
-        unless Record.unique_hash?(hash)
-          Miteru.logger.info "Don't download #{kit.url}. The same hash is already recorded. (SHA256: #{hash})."
-          FileUtils.rm downloaded_as
-          return
-        end
+    private
 
-        # Record a kit in DB
-        Record.create_by_kit_and_hash(kit, hash)
-        Miteru.logger.info "Download #{kit.url} as #{downloaded_as}"
+    def timeout
+      Miteru.config.download_timeout
+    end
+
+    def downloader
+      Down::Http.new(ssl_context:) { |client| client.timeout(timeout) }
+    end
+
+    def ssl_context
+      OpenSSL::SSL::SSLContext.new.tap do |ctx|
+        ctx.verify_mode = OpenSSL::SSL::VERIFY_NONE
       end
     end
 
-    def sha256(path)
-      return memo[path] if memo.key?(path)
+    def max_size
+      Miteru.config.file_max_size
+    end
 
+    def sha256(path)
       digest = Digest::SHA256.file(path)
-      hash = digest.hexdigest
-      memo[path] = hash
-      hash
+      digest.hexdigest
     end
+    memo_wise :sha256
   end
 end

data/lib/miteru/errors.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+require "http"
+
+module Miteru
+  class Error < StandardError; end
+
+  class FileSizeError < Error; end
+
+  class DownloadError < Error; end
+
+  class UniquenessError < Error; end
+
+  class StatusError < ::HTTP::Error
+    # @return [Integer]
+    attr_reader :status_code
+
+    # @return [String, nil]
+    attr_reader :body
+
+    #
+    # @param [String] msg
+    # @param [Integer] status_code
+    # @param [String, nil] body
+    #
+    def initialize(msg, status_code, body)
+      super(msg)
+
+      @status_code = status_code
+      @body = body
+    end
+
+    def detail
+      {status_code:, body:}
+    end
+  end
+end

data/lib/miteru/feeds/ayashige.rb

@@ -1,34 +1,23 @@
 # frozen_string_literal: true
 
-require "json"
-require "uri"
-
 module Miteru
   class Feeds
-    class Ayashige < Feed
+    class Ayashige < Base
       HOST = "ayashige.herokuapp.com"
-      URL = "https://#{HOST}"
+      URL = "https://#{HOST}".freeze
 
-      def urls
-        url = url_for("/api/v1/domains/")
-        res = JSON.parse(get(url))
+      def initialize(base_url = "https://ayashige.herokuapp.com")
+        super(base_url)
+      end
 
-        domains = res.map { |item| item["fqdn"] }
-        domains.map do |domain|
-          [
-            "https://#{domain}",
-            "http://#{domain}"
-          ]
-        end.flatten
-      rescue HTTPResponseError, HTTP::Error, JSON::ParserError => e
-        Miteru.logger.error "Failed to load ayashige feed (#{e})"
-        []
+      def urls
+        json.map { |item| item["fqdn"] }.map { |fqdn| "https://#{fqdn}" }
       end
 
       private
 
-      def url_for(path)
-        URI(URL + path)
+      def json
+        get_json "/api/v1/domains/"
       end
     end
   end

data/lib/miteru/feeds/base.rb

@@ -0,0 +1,141 @@
+# frozen_string_literal: true
+
+module Miteru
+  class Feeds
+    class Base < Service
+      IGNORE_EXTENSIONS = %w[.htm .html .php .asp .aspx .exe .txt].freeze
+
+      # @return [String]
+      attr_reader :base_url
+
+      # @return [Hash]
+      attr_reader :headers
+
+      #
+      # @param [String] base_url
+      #
+      def initialize(base_url)
+        super()
+
+        @base_url = base_url
+        @headers = {}
+      end
+
+      def source
+        @source ||= self.class.to_s.split("::").last
+      end
+
+      #
+      # Return URLs
+      #
+      # @return [Array<String>] URLs
+      #
+      def urls
+        raise NotImplementedError, "You must implement #{self.class}##{__method__}"
+      end
+
+      #
+      # Return decomposed URLs
+      #
+      # @return [Array<String>] Decomposed URLs
+      #
+      def decomposed_urls
+        urls.uniq.select { |url| url.start_with?("http://", "https://") }.map { |url| decompose(url) }.flatten.uniq
+      end
+
+      #
+      # @return [Array<Miteru::Website>]
+      #
+      def call
+        decomposed_urls.map { |url| Website.new(url, source:) }
+      end
+
+      class << self
+        def inherited(child)
+          super
+          Miteru.feeds << child
+        end
+      end
+
+      private
+
+      def timeout
+        Miteru.config.api_timeout
+      end
+
+      def directory_traveling?
+        Miteru.config.directory_traveling
+      end
+
+      #
+      # Validate extension of a URL
+      #
+      # @param [String] url
+      #
+      # @return [Boolean]
+      #
+      def invalid_extension?(url)
+        IGNORE_EXTENSIONS.any? { |ext| url.end_with? ext }
+      end
+
+      #
+      # Decompose a URL into URLs
+      #
+      # @param [String] url
+      #
+      # @return [Array<String>]
+      #
+      def decompose(url)
+        Try[URI::InvalidURIError] do
+          parsed = URI.parse(url)
+
+          base = "#{parsed.scheme}://#{parsed.hostname}"
+          return [base] unless directory_traveling?
+
+          segments = parsed.path.split("/")
+          return [base] if segments.empty?
+
+          urls = (0...segments.length).map { |idx| "#{base}#{segments[0..idx].join("/")}" }
+          urls.reject { |url| invalid_extension? url }
+        end.recover { [] }.value!
+      end
+
+      #
+      # @return [::HTTP::Client]
+      #
+      def http
+        @http ||= HTTP::Factory.build(headers:, timeout:)
+      end
+
+      #
+      # @param [String] path
+      #
+      # @return [URI]
+      #
+      def url_for(path)
+        URI.join base_url, path
+      end
+
+      #
+      # @param [String] path
+      # @param [Hash, nil] params
+      #
+      # @return [::HTTP::Response]
+      #
+      def get(path, params: nil)
+        http.get(url_for(path), params:)
+      end
+
+      #
+      # @param [String] path
+      # @param [Hash, nil] params
+      #
+      # @return [Hash]
+      #
+      def get_json(path, params: nil)
+        res = get(path, params:)
+        JSON.parse res.body.to_s
+      end
+    end
+  end
+end

data/lib/miteru/feeds/phishing_database.rb

@@ -1,19 +1,20 @@
 # frozen_string_literal: true
 
-require "json"
-require "uri"
-
 module Miteru
   class Feeds
-    class PhishingDatabase < Feed
-      URL = "https://raw.githubusercontent.com/mitchellkrogza/Phishing.Database/master/phishing-links-ACTIVE-NOW.txt"
+    class PhishingDatabase < Base
+      def initialize(base_url = "https://raw.githubusercontent.com")
+        super(base_url)
+      end
 
       def urls
-        body = get(URL)
-        body.to_s.lines.map(&:chomp)
-      rescue HTTPResponseError, HTTP::Error, JSON::ParserError => e
-        info "Failed to load phishing database feed (#{e})"
-        []
+        text.lines.map(&:chomp)
+      end
+
+      private
+
+      def text
+        get("/mitchellkrogza/Phishing.Database/master/phishing-links-ACTIVE-NOW.txt").body.to_s
       end
     end
   end

data/lib/miteru/feeds/urlscan.rb

@@ -1,35 +1,63 @@
 # frozen_string_literal: true
 
-require "urlscan"
-
 module Miteru
   class Feeds
-    class UrlScan < Feed
-      attr_reader :size
-
-      def initialize(size = 100)
-        @size = size
-        raise ArgumentError, "size must be less than 10,000" if size > 10_000
-      end
+    class UrlScan < Base
+      #
+      # @param [String] base_url
+      #
+      def initialize(base_url = "https://urlscan.io")
+        super(base_url)
 
-      def api
-        @api ||= ::UrlScan::API.new(Miteru.configuration.urlscan_api_key)
+        @headers = {"api-key": api_key}
       end
 
       def urls
-        urls_from_community_feed
-      rescue ::UrlScan::ResponseError => e
-        Miteru.logger.error "Failed to load urlscan.io feed (#{e})"
-        []
+        search_with_pagination.flat_map do |json|
+          (json["results"] || []).map { |result| result.dig("task", "url") }
+        end.uniq
       end
 
       private
 
-      def urls_from_community_feed
-        res = api.search("task.method:automatic", size: size)
+      def size
+        10_000
+      end
+
+      # @return [<Type>] <description>
+      #
+      def api_key
+        Miteru.config.urlscan_api_key
+      end
+
+      def q
+        "task.method:automatic AND date:#{Miteru.config.urlscan_date_condition}"
+      end
+
+      #
+      # @param [String, nil] search_after
+      #
+      # @return [Hash]
+      #
+      def search(search_after: nil)
+        get_json("/api/v1/search/", params: {q:, size:, search_after:}.compact)
+      end
+
+      def search_with_pagination
+        search_after = nil
+
+        Enumerator.new do |y|
+          loop do
+            res = search(search_after:)
+
+            y.yield res
+
+            has_more = res["has_more"]
+            break unless has_more
 
-        results = res["results"] || []
-        results.map { |result| result.dig("task", "url") }
+            search_after = res["results"].last["sort"].join(",")
+          end
+        end
       end
     end
   end

data/lib/miteru/feeds/urlscan_pro.rb

@@ -1,35 +1,37 @@
 # frozen_string_literal: true
 
-require "urlscan"
-
 module Miteru
   class Feeds
-    class UrlScanPro < Feed
-      def api
-        @api ||= ::UrlScan::API.new(Miteru.configuration.urlscan_api_key)
+    class UrlScanPro < Base
+      #
+      # @param [String] base_url
+      #
+      def initialize(base_url = "https://urlscan.io")
+        super(base_url)
+
+        @headers = {"api-key": api_key}
       end
 
       def urls
-        urls_from_pro_feed
-      rescue ::UrlScan::ResponseError => e
-        Miteru.logger.error "Failed to load urlscan.io pro feed (#{e})"
-        []
+        (json["results"] || []).map { |result| result["page_url"] }
       end
 
       private
 
-      def api_key?
-        Miteru.configuration.urlscan_api_key?
+      def api_key
+        Miteru.config.urlscan_api_key
       end
 
-      def urls_from_pro_feed
-        return [] unless api_key?
+      def q
+        "date:#{Miteru.config.urlscan_date_condition}"
+      end
+
+      def format
+        "json"
+      end
 
-        res = api.pro.phishfeed
-        results = res["results"] || []
-        results.map { |result| result["page_url"] }
-      rescue ArgumentError => _e
-        []
+      def json
+        get_json("/api/v1/pro/phishfeed", params: {q:, format:})
       end
     end
   end

data/lib/miteru/http.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+require "http"
+
+module Miteru
+  module HTTP
+    #
+    # Better error handling feature
+    #
+    class BetterError < ::HTTP::Feature
+      def wrap_response(response)
+        return response if response.status.success?
+
+        raise StatusError.new(
+          "Unsuccessful response code returned: #{response.code}",
+          response.code,
+          response.body.to_s
+        )
+      end
+
+      ::HTTP::Options.register_feature(:better_error, self)
+    end
+
+    #
+    # HTTP client factory
+    #
+    class Factory
+      class << self
+        USER_AGENT = "miteru/#{Miteru::VERSION}".freeze
+
+        #
+        # @param [Integer, nil] timeout
+        # @param [Hash] headers
+        # @param [Boolean] raise_exception
+        #
+        # @return [::HTTP::Client]
+        #
+        # @param [Object] raise_exception
+        def build(headers: {}, timeout: nil, raise_exception: true)
+          client = raise_exception ? ::HTTP.use(:better_error) : ::HTTP
+
+          headers["User-Agent"] ||= USER_AGENT
+
+          client = client.headers(headers)
+          client = client.timeout(timeout) unless timeout.nil?
+          client
+        end
+      end
+    end
+  end
+end