Merge pull request #3 from taketo1113/option-ipv4-ipv6

Add ipv4/ipv6 options

.gitignore

@@ -8,6 +8,7 @@
 /tmp/
 
 /Gemfile.lock
+.ruby-version
 
 # rspec failure tracking
 .rspec_status

README.md

@@ -28,6 +28,7 @@ If bundler is not being used to manage dependencies, install the gem by executin
     $ gem install ddig
 
 ## Usage
+### Ruby
 
 ```ruby
 ddig = Ddig.lookup('dns.google', nameservers: ['8.8.8.8', '2001:4860:4860::8888'])
@@ -120,6 +121,26 @@ doh.aaaa
 ```
 
 ### CLI
+- Usage
+```
+# ddig --help
+Usage: ddig [options] hostname
+    -t, --type={all|do53|dot}        resolve type (default: all)
+        --udp                        use resolve type of udp(do53)
+        --dot                        use resolve type of dot
+        --doh-h1                     use resolve type of doh (http/1.1)
+        --doh-path=doh-path          doh service path
+    -4, --ipv4                       use IPv4 query transport only
+    -6, --ipv6                       use IPv6 query transport only
+    -@ipaddress|doh-hostname,        nameserver
+        --nameserver
+    -p, --port=port                  port
+        --format={text|json}         output format (default: text)
+
+    -v, --verbose                    run verbosely
+    -h, --help                       show this help message.
+        --version                    show version.
+```
 
 - UDP(Do53)
 ```sh

lib/ddig.rb

@@ -2,6 +2,7 @@
 
 require_relative "ddig/version"
 require_relative "ddig/nameserver"
+require_relative "ddig/ip"
 require_relative "ddig/resolver/do53"
 require_relative "ddig/resolver/dot"
 require_relative "ddig/resolver/doh_h1"
@@ -10,16 +11,19 @@
 module Ddig
   class Error < StandardError; end
 
-  def self.lookup(hostname, nameservers: nil)
+  def self.lookup(hostname, nameservers: nil, use_ipv4: nil, use_ipv6: nil)
     @hostname = hostname
     @nameservers = nameservers
+    @use_ipv4 = use_ipv4
+    @use_ipv6 = use_ipv6
 
     @nameserver = Ddig::Nameserver.new(nameservers: @nameservers)
+    @ip = Ddig::Ip.new(use_ipv4: @use_ipv4, use_ipv6: @use_ipv6)
 
-    @do53_ipv4 = Ddig::Resolver::Do53.new(hostname: @hostname, nameservers: @nameserver.servers, ip: :ipv4).lookup
-    @do53_ipv6 = Ddig::Resolver::Do53.new(hostname: @hostname, nameservers: @nameserver.servers, ip: :ipv6).lookup
+    @do53_ipv4 = Ddig::Resolver::Do53.new(hostname: @hostname, nameservers: @nameserver.servers, ip: :ipv4).lookup unless @ip.ip_type == :ipv6
+    @do53_ipv6 = Ddig::Resolver::Do53.new(hostname: @hostname, nameservers: @nameserver.servers, ip: :ipv6).lookup unless @ip.ip_type == :ipv4
 
-    @ddr = Ddig::Ddr.new(nameservers: @nameservers)
+    @ddr = Ddig::Ddr.new(nameservers: @nameservers, ip: @ip.ip_type)
 
     {
       do53: {

lib/ddig/cli.rb

@@ -28,6 +28,8 @@ def parse_options
         opts.on("--dot", "use resolve type of dot") { |v| @options[:type] = 'dot' }
         opts.on("--doh-h1", "use resolve type of doh (http/1.1)") { |v| @options[:type] = 'doh_h1' }
         opts.on("--doh-path=doh-path", "doh service path") { |v| @options[:doh_path] = v }
+        opts.on("-4", "--ipv4", "use IPv4 query transport only") { |v| @options[:ipv4] = v }
+        opts.on("-6", "--ipv6", "use IPv6 query transport only") { |v| @options[:ipv6] = v }
         opts.on("-@", "--nameserver=ipaddress|doh-hostname", "nameserver") { |v| @options[:nameserver] = v }
         opts.on("-p", "--port=port", "port") { |v| @options[:port] = v }
         opts.on("--format={text|json}", "output format (default: text)") { |v| @options[:format] = v }
@@ -45,6 +47,11 @@ def parse_options
     end
 
     def exec
+      if @options[:ipv4] || @options[:ipv6]
+        @use_ipv4 = @options[:ipv4] || false
+        @use_ipv6 = @options[:ipv6] || false
+      end
+
       case @options[:type]
       when "all"
         resolve_all
@@ -58,7 +65,7 @@ def exec
     end
 
     def resolve_all
-      @ddig = Ddig.lookup(@hostname, nameservers: @options[:nameserver])
+      @ddig = Ddig.lookup(@hostname, nameservers: @options[:nameserver], use_ipv4: @use_ipv4, use_ipv6: @use_ipv6)
 
       if @options[:format] == 'json'
         # TODO: to_json
@@ -69,7 +76,13 @@ def resolve_all
     end
 
     def resolve_do53
-      do53 = Ddig::Resolver::Do53.new(hostname: @hostname, nameservers: @options[:nameserver]).lookup
+      ip = Ddig::Ip.new(use_ipv4: @use_ipv4, use_ipv6: @use_ipv6)
+      do53 = Ddig::Resolver::Do53.new(hostname: @hostname, nameservers: @options[:nameserver], ip: ip.ip_type).lookup
+
+      if do53.nil?
+        puts "Error: Could not lookup wit nameserver: #{@options[:nameserver]}"
+        exit
+      end
 
       do53.a.each do |address|
         rr_type = 'A'

lib/ddig/ddr.rb

@@ -68,35 +68,39 @@ def discover_designated_resolvers
           do53_v6 = ::Ddig::Resolver::Do53.new(hostname: target, nameservers: [unencrypted_resolver], ip: :ipv6).lookup
 
           # ipv4
-          unless do53_v4.nil? || do53_v4.a.nil?
-            do53_v4.a.each do |address|
-              designated_resolver = ::Ddig::Ddr::DesignatedResolver.new(unencrypted_resolver: unencrypted_resolver, target: target, protocol: protocol, port: port, dohpath: dohpath, address: address.to_s, ip: :ipv4)
-              @designated_resolvers << designated_resolver
+          unless @ip == :ipv6
+            unless do53_v4.nil? || do53_v4.a.nil?
+              do53_v4.a.each do |address|
+                designated_resolver = ::Ddig::Ddr::DesignatedResolver.new(unencrypted_resolver: unencrypted_resolver, target: target, protocol: protocol, port: port, dohpath: dohpath, address: address.to_s, ip: :ipv4)
+                @designated_resolvers << designated_resolver
+              end
             end
-          end
-          unless do53_v6.nil? || do53_v6.a.nil?
-            do53_v6.a.each do |address|
-              designated_resolver = ::Ddig::Ddr::DesignatedResolver.new(unencrypted_resolver: unencrypted_resolver, target: target, protocol: protocol, port: port, dohpath: dohpath, address: address.to_s, ip: :ipv4)
-              @designated_resolvers << designated_resolver
+            unless do53_v6.nil? || do53_v6.a.nil?
+              do53_v6.a.each do |address|
+                designated_resolver = ::Ddig::Ddr::DesignatedResolver.new(unencrypted_resolver: unencrypted_resolver, target: target, protocol: protocol, port: port, dohpath: dohpath, address: address.to_s, ip: :ipv4)
+                @designated_resolvers << designated_resolver
+              end
             end
           end
 
           # ipv6
-          unless do53_v4.nil? || do53_v4.aaaa.nil?
-            do53_v4.aaaa.each do |address|
-              designated_resolver = ::Ddig::Ddr::DesignatedResolver.new(unencrypted_resolver: unencrypted_resolver, target: target, protocol: protocol, port: port, dohpath: dohpath, address: address.to_s, ip: :ipv6)
-              @designated_resolvers << designated_resolver
+          unless @ip == :ipv4
+            unless do53_v4.nil? || do53_v4.aaaa.nil?
+              do53_v4.aaaa.each do |address|
+                designated_resolver = ::Ddig::Ddr::DesignatedResolver.new(unencrypted_resolver: unencrypted_resolver, target: target, protocol: protocol, port: port, dohpath: dohpath, address: address.to_s, ip: :ipv6)
+                @designated_resolvers << designated_resolver
+              end
             end
-          end
-          unless do53_v6.nil? || do53_v6.aaaa.nil?
-            do53_v6.aaaa.each do |address|
-              designated_resolver = ::Ddig::Ddr::DesignatedResolver.new(unencrypted_resolver: unencrypted_resolver, target: target, protocol: protocol, port: port, dohpath: dohpath, address: address.to_s, ip: :ipv6)
-              @designated_resolvers << designated_resolver
+            unless do53_v6.nil? || do53_v6.aaaa.nil?
+              do53_v6.aaaa.each do |address|
+                designated_resolver = ::Ddig::Ddr::DesignatedResolver.new(unencrypted_resolver: unencrypted_resolver, target: target, protocol: protocol, port: port, dohpath: dohpath, address: address.to_s, ip: :ipv6)
+                @designated_resolvers << designated_resolver
+              end
             end
           end
 
           # ipv4hint
-          unless ipv4hint.nil?
+          unless ipv4hint.nil? || @ip == :ipv6
             ipv4hint.each do |address|
               ip = :ipv4
               designated_resolver = ::Ddig::Ddr::DesignatedResolver.new(unencrypted_resolver: unencrypted_resolver, target: target, protocol: protocol, port: port, dohpath: dohpath, address: address.to_s, ip: ip)
@@ -105,7 +109,7 @@ def discover_designated_resolvers
           end
 
           # ipv6hint
-          unless ipv6hint.nil?
+          unless ipv6hint.nil? || @ip == :ipv4
             ipv6hint.each do |address|
               ip = :ipv6
               designated_resolver = ::Ddig::Ddr::DesignatedResolver.new(unencrypted_resolver: unencrypted_resolver, target: target, protocol: protocol, port: port, dohpath: dohpath, address: address.to_s, ip: ip)

lib/ddig/ip.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module Ddig
+  class Ip
+    attr_reader :ip_type
+
+    def initialize(use_ipv4: nil, use_ipv6: nil)
+      @use_ipv4 = use_ipv4
+      @use_ipv6 = use_ipv6
+
+      set_ip_type
+    end
+
+    def set_ip_type
+      if @use_ipv4.nil? && self.class.enable_ipv4?
+        @use_ipv4 = true
+      end
+
+      if @use_ipv6.nil? && self.class.enable_ipv6?
+        @use_ipv6 = true
+      end
+
+      if @use_ipv4 && @use_ipv6
+        @ip_type = :all
+      elsif @use_ipv4
+        @ip_type = :ipv4
+      elsif @use_ipv6
+        @ip_type = :ipv6
+      end
+    end
+
+    def self.enable_ipv4?
+      ip_list.any? { |addr| addr.ipv4? }
+    end
+
+    def self.enable_ipv6?
+      ip_list.any? { |addr| addr.ipv6? }
+    end
+
+    private
+
+    def self.ip_list
+      Socket.ip_address_list.map do |addrinfo|
+        if RUBY_VERSION < '3.1'
+          # for ipaddr gem <= v1.2.2, Not support zone identifiers in IPv6 addresses
+          addr = IPAddr.new(addrinfo.ip_address.split('%').first)
+        else
+          addr = IPAddr.new(addrinfo.ip_address)
+        end
+
+        if !addr.loopback? && !addr.link_local?
+          addr
+        end
+      end.compact
+    end
+  end
+end

spec/ddig/ip_spec.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+RSpec.describe Ddig::Ip do
+  it "#enable_ipv4?" do
+    skip "IPv4 enabled: #{Ddig::Ip.enable_ipv4?}" unless enable_ipv4?
+    expect(Ddig::Ip.enable_ipv4?).to be true
+  end
+
+  it "#enable_ipv6?" do
+    skip "IPv6 enabled: #{Ddig::Ip.enable_ipv6?}" unless enable_ipv6?
+    expect(Ddig::Ip.enable_ipv6?).to be true
+  end
+
+  context "ip_type" do
+    it "return :all with use_ipv4 & use_ipv6 is true" do
+      use_ipv4 = true
+      use_ipv6 = true
+
+      expect(Ddig::Ip.new(use_ipv4: use_ipv4, use_ipv6: use_ipv6).ip_type).to be :all
+    end
+
+    it "return :ipv4 with use_ipv4: true & use_ipv6: false" do
+      use_ipv4 = true
+      use_ipv6 = false
+
+      expect(Ddig::Ip.new(use_ipv4: use_ipv4, use_ipv6: use_ipv6).ip_type).to be :ipv4
+    end
+
+    it "return :all with use_ipv4: false & use_ipv6: true" do
+      use_ipv4 = false
+      use_ipv6 = true
+
+      expect(Ddig::Ip.new(use_ipv4: use_ipv4, use_ipv6: use_ipv6).ip_type).to be :ipv6
+    end
+
+    context "use_ipv4/use_ipv6: nil" do
+      before(:each) do
+        @use_ipv4 = nil
+        @use_ipv6 = nil
+      end
+
+      it "return ip_type via enable_ipv4?/enable_ipv6?" do
+        if enable_ipv4? && enable_ipv6?
+          ip_type = :all
+        elsif enable_ipv4?
+          ip_type = :ipv4
+        elsif enable_ipv6?
+          ip_type = :ipv6
+        end
+
+        expect(Ddig::Ip.new(use_ipv4: @use_ipv4, use_ipv6: @use_ipv6).ip_type).to be ip_type
+      end
+    end
+  end
+end

spec/support/ip_helper.rb

@@ -0,0 +1,11 @@
+module Ddig
+  module Ipv6Helper
+    def enable_ipv4?
+      Ddig::Ip.enable_ipv4?
+    end
+
+    def enable_ipv6?
+      Ddig::Ip.enable_ipv6?
+    end
+  end
+end