From 97250e3ff61f83ab0751dcf79843e2e65bdae7c5 Mon Sep 17 00:00:00 2001
From: Manabu Niseki <manabu.niseki@gmail.com>
Date: Wed, 7 Oct 2020 09:34:26 +0900
Subject: [PATCH] feat: add --record_type option

Add --record__type option to set the record type to query
---
 README.md                 |  6 +++---
 lib/rogue_one/cli.rb      | 16 +++++++++++++---
 lib/rogue_one/detector.rb | 37 ++++++++++++++++++++++++++++++-------
 3 files changed, 46 insertions(+), 13 deletions(-)

diff --git a/README.md b/README.md
index 8c0c5be..8b8457a 100644
--- a/README.md
+++ b/README.md
@@ -39,16 +39,16 @@ Usage:
   rogue_one report [DNS_SERVER]
 
 Options:
+  [--custom-list=CUSTOM_LIST]    # A path to a custom list of domains
   [--default-list=DEFAULT_LIST]  # A default list of top 100 domains (Alexa or Fortune)
                                  # Default: alexa
-  [--custom-list=CUSTOM_LIST]    # A path to a custom list of domains
+  [--record-type=RECORD_TYPE]    # A type of the DNS resource to check
+                                 # Default: A
   [--threshold=N]                # Threshold value for determining malicious or not
   [--verbose], [--no-verbose]
 
 Show a report of a given DNS server
 
-Show a report of a given DNS server
-
 $ rogue_one report 1.1.1.1
 {
   "verdict": "benign one",
diff --git a/lib/rogue_one/cli.rb b/lib/rogue_one/cli.rb
index 96fc624..1578b39 100644
--- a/lib/rogue_one/cli.rb
+++ b/lib/rogue_one/cli.rb
@@ -12,19 +12,29 @@ def exit_on_failure?
     end
 
     desc "report [DNS_SERVER]", "Show a report of a given DNS server"
-    method_option :default_list, type: :string, default: "alexa", desc: "A default list of top 100 domains (Alexa or Fortune)"
     method_option :custom_list, type: :string, desc: "A path to a custom list of domains"
+    method_option :default_list, type: :string, default: "alexa", desc: "A default list of top 100 domains (Alexa or Fortune)"
+    method_option :record_type, type: :string, default: "A", desc: "A type of the DNS resource to check"
     method_option :threshold, type: :numeric, desc: "Threshold value for determining malicious or not"
     method_option :verbose, type: :boolean
     def report(dns_server)
       with_error_handling do
         Ping.pong? dns_server
 
-        default_list = options["default_list"].downcase
         custom_list = options["custom_list"]
+        default_list = options["default_list"].downcase
+        record_type = options["record_type"].upcase
         threshold = options["threshold"]
         verbose = options["verbose"]
-        detector = Detector.new(target: dns_server, default_list: default_list, custom_list: custom_list, threshold: threshold, verbose: verbose)
+
+        detector = Detector.new(
+          custom_list: custom_list,
+          default_list: default_list,
+          record_type: record_type,
+          target: dns_server,
+          threshold: threshold,
+          verbose: verbose,
+        )
         puts JSON.pretty_generate(detector.report)
       end
     end
diff --git a/lib/rogue_one/detector.rb b/lib/rogue_one/detector.rb
index c49243a..522938e 100644
--- a/lib/rogue_one/detector.rb
+++ b/lib/rogue_one/detector.rb
@@ -11,18 +11,26 @@
 
 module RogueOne
   class Detector
-    attr_reader :target
-    attr_reader :default_list
     attr_reader :custom_list
-    attr_reader :verbose
+    attr_reader :default_list
     attr_reader :max_concurrency
+    attr_reader :record_type
+    attr_reader :target
+    attr_reader :verbose
 
     GOOGLE_PUBLIC_DNS = "8.8.8.8"
 
-    def initialize(target:, default_list: "alexa", custom_list: nil, threshold: nil, verbose: false)
+    def initialize(target:,
+                   custom_list: nil,
+                   default_list: "alexa",
+                   record_type: "A",
+                   threshold: nil,
+                   verbose: false)
       @target = target
-      @default_list = default_list
+
       @custom_list = custom_list
+      @default_list = default_list
+      @record_type = record_type.upcase.to_sym
       @threshold = threshold
       @verbose = verbose
 
@@ -59,7 +67,10 @@ def threshold
     def meta
       return nil unless verbose
 
-      { threshold: threshold }
+      {
+        record_type: record_type,
+        threshold: threshold,
+      }
     end
 
     def landing_pages
@@ -135,6 +146,7 @@ def read_domains(path)
 
     def bulk_resolve(resolver, domains)
       results = []
+
       Async do
         barrier = Async::Barrier.new
         semaphore = Async::Semaphore.new(max_concurrency, parent: barrier)
@@ -143,7 +155,7 @@ def bulk_resolve(resolver, domains)
           semaphore.async do
             addresses = []
             begin
-              addresses = resolver.addresses_for(domain, Resolv::DNS::Resource::IN::A, { retries: 1 }).map(&:to_s)
+              addresses = resolver.addresses_for(domain, dns_resource_by_record_type, { retries: 1 }).map(&:to_s)
             rescue Async::DNS::ResolutionFailure
               # do nothing
             end
@@ -161,5 +173,16 @@ def normal_resolver
     def target_resolver
       Async::DNS::Resolver.new([[:udp, target, 53], [:tcp, target, 53]])
     end
+
+    def dns_resource_by_record_type
+      @dns_resource_by_record_type ||= dns_resources.dig(record_type)
+    end
+
+    def dns_resources
+      {
+        A: Resolv::DNS::Resource::IN::A,
+        AAAA: Resolv::DNS::Resource::IN::AAAA,
+      }
+    end
   end
 end