From e909bf2f8294f18f51e17cbffdfe1f6e3bc2c9fb Mon Sep 17 00:00:00 2001
From: Manabu Niseki <manabu.niseki@gmail.com>
Date: Sun, 3 Mar 2024 07:59:59 +0900
Subject: [PATCH] refactor: split enricher & client

---
 lib/mihari.rb                 |   1 +
 lib/mihari/clients/whois.rb   | 118 ++++++++++++++++++++++++++++++++++
 lib/mihari/enrichers/whois.rb |  97 ++--------------------------
 3 files changed, 125 insertions(+), 91 deletions(-)
 create mode 100644 lib/mihari/clients/whois.rb

diff --git a/lib/mihari.rb b/lib/mihari.rb
index 4b1ff774..d693d56d 100644
--- a/lib/mihari.rb
+++ b/lib/mihari.rb
@@ -244,6 +244,7 @@ def initialize_sentry
 require "mihari/clients/the_hive"
 require "mihari/clients/urlscan"
 require "mihari/clients/virustotal"
+require "mihari/clients/whois"
 require "mihari/clients/zoomeye"
 
 # Analyzers
diff --git a/lib/mihari/clients/whois.rb b/lib/mihari/clients/whois.rb
new file mode 100644
index 00000000..d7d16dc6
--- /dev/null
+++ b/lib/mihari/clients/whois.rb
@@ -0,0 +1,118 @@
+# frozen_string_literal: true
+
+require "whois-parser"
+
+module Mihari
+  module Clients
+    #
+    # Whois client
+    #
+    class Whois
+      # @return [Integer, nil]
+      attr_reader :timeout
+
+      # @return [::Whois::Client]
+      attr_reader :client
+
+      #
+      # @param [Integer, nil] timeout
+      #
+      def initialize(timeout: nil)
+        @timeout = timeout
+
+        @client = lambda do
+          return ::Whois::Client.new if timeout.nil?
+
+          ::Whois::Client.new(timeout:)
+        end.call
+      end
+
+      #
+      # Query IAIA Whois API
+      #
+      # @param [Mihari::Models::Artifact] artifact
+      #
+      # @param [Object] domain
+      def lookup(domain)
+        record = client.lookup(domain)
+        return if record.parser.available?
+
+        Models::WhoisRecord.new(
+          domain:,
+          created_on: get_created_on(record.parser),
+          updated_on: get_updated_on(record.parser),
+          expires_on: get_expires_on(record.parser),
+          registrar: get_registrar(record.parser),
+          contacts: get_contacts(record.parser)
+        )
+      end
+
+      private
+
+      #
+      # Get created_on
+      #
+      # @param [::Whois::Parser] parser
+      #
+      # @return [Date, nil]
+      #
+      def get_created_on(parser)
+        parser.created_on
+      rescue ::Whois::AttributeNotImplemented
+        nil
+      end
+
+      #
+      # Get updated_on
+      #
+      # @param [::Whois::Parser] parser
+      #
+      # @return [Date, nil]
+      #
+      def get_updated_on(parser)
+        parser.updated_on
+      rescue ::Whois::AttributeNotImplemented
+        nil
+      end
+
+      #
+      # Get expires_on
+      #
+      # @param [::Whois::Parser] parser
+      #
+      # @return [Date, nil]
+      #
+      def get_expires_on(parser)
+        parser.expires_on
+      rescue ::Whois::AttributeNotImplemented
+        nil
+      end
+
+      #
+      # Get registrar
+      #
+      # @param [::Whois::Parser] parser
+      #
+      # @return [Hash, nil]
+      #
+      def get_registrar(parser)
+        parser.registrar&.to_h
+      rescue ::Whois::AttributeNotImplemented
+        nil
+      end
+
+      #
+      # Get contacts
+      #
+      # @param [::Whois::Parser] parser
+      #
+      # @return [Array<Hash>, nil]
+      #
+      def get_contacts(parser)
+        parser.contacts.map(&:to_h)
+      rescue ::Whois::AttributeNotImplemented
+        nil
+      end
+    end
+  end
+end
diff --git a/lib/mihari/enrichers/whois.rb b/lib/mihari/enrichers/whois.rb
index f967def2..f00678ad 100644
--- a/lib/mihari/enrichers/whois.rb
+++ b/lib/mihari/enrichers/whois.rb
@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require "whois-parser"
-
 module Mihari
   module Enrichers
     #
@@ -18,22 +16,15 @@ class Whois < Base
       def call(artifact)
         return if artifact.domain.nil?
 
-        domain = PublicSuffix.domain(artifact.domain)
-        record = memoized_lookup(domain)
-        return if record.parser.available?
-
-        artifact.whois_record ||= Models::WhoisRecord.new(
-          domain:,
-          created_on: get_created_on(record.parser),
-          updated_on: get_updated_on(record.parser),
-          expires_on: get_expires_on(record.parser),
-          registrar: get_registrar(record.parser),
-          contacts: get_contacts(record.parser)
-        )
+        artifact.whois_record ||= memoized_lookup(PublicSuffix.domain(artifact.domain))
       end
 
       private
 
+      def client
+        @client ||= Clients::Whois.new(timeout:)
+      end
+
       #
       # @param [Mihari::Models::Artifact] artifact
       #
@@ -53,85 +44,9 @@ def supported_data_types
       # @return [Mihari::Models::WhoisRecord, nil]
       #
       def memoized_lookup(domain)
-        whois.lookup domain
+        client.lookup domain
       end
       memo_wise :memoized_lookup
-
-      #
-      # @return [::Whois::Client]
-      #
-      def whois
-        @whois ||= lambda do
-          return ::Whois::Client.new if timeout.nil?
-
-          ::Whois::Client.new(timeout:)
-        end.call
-      end
-
-      #
-      # Get created_on
-      #
-      # @param [::Whois::Parser] parser
-      #
-      # @return [Date, nil]
-      #
-      def get_created_on(parser)
-        parser.created_on
-      rescue ::Whois::AttributeNotImplemented
-        nil
-      end
-
-      #
-      # Get updated_on
-      #
-      # @param [::Whois::Parser] parser
-      #
-      # @return [Date, nil]
-      #
-      def get_updated_on(parser)
-        parser.updated_on
-      rescue ::Whois::AttributeNotImplemented
-        nil
-      end
-
-      #
-      # Get expires_on
-      #
-      # @param [::Whois::Parser] parser
-      #
-      # @return [Date, nil]
-      #
-      def get_expires_on(parser)
-        parser.expires_on
-      rescue ::Whois::AttributeNotImplemented
-        nil
-      end
-
-      #
-      # Get registrar
-      #
-      # @param [::Whois::Parser] parser
-      #
-      # @return [Hash, nil]
-      #
-      def get_registrar(parser)
-        parser.registrar&.to_h
-      rescue ::Whois::AttributeNotImplemented
-        nil
-      end
-
-      #
-      # Get contacts
-      #
-      # @param [::Whois::Parser] parser
-      #
-      # @return [Array<Hash>, nil]
-      #
-      def get_contacts(parser)
-        parser.contacts.map(&:to_h)
-      rescue ::Whois::AttributeNotImplemented
-        nil
-      end
     end
   end
 end