Added unit tests

Author: smashery

lib/rex/proto/dns/cached_resolver.rb

@@ -21,6 +21,7 @@ class CachedResolver < Resolver
     #
     # @return [nil]
     def initialize(config = {})
+      dns_cache_no_start = config.delete(:dns_cache_no_start)
       super(config)
       self.cache = Rex::Proto::DNS::Cache.new
       # Read hostsfile into cache
@@ -72,7 +73,7 @@ def initialize(config = {})
         end
       end
       # TODO: inotify or similar on hostsfile for live updates? Easy-button functionality
-      self.cache.start unless config[:dns_cache_no_start]
+      self.cache.start unless dns_cache_no_start
       return
     end
 

spec/lib/rex/proto/dns/custom_nameserver_provider_spec.rb

@@ -0,0 +1,114 @@
+# -*- coding:binary -*-
+require 'spec_helper'
+require 'net/dns'
+
+
+RSpec.describe Rex::Proto::DNS::CustomNameserverProvider do
+  def packet_for(name)
+    packet = Net::DNS::Packet.new(name, Net::DNS::A, Net::DNS::IN)
+    Rex::Proto::DNS::Packet.encode_drb(packet)
+  end
+
+  let(:base_nameserver) do
+    '1.2.3.4'
+  end
+
+  let(:ruleless_nameserver) do
+    '1.2.3.5'
+  end
+
+  let(:ruled_nameserver) do
+    '1.2.3.6'
+  end
+
+  let(:ruled_nameserver2) do
+    '1.2.3.7'
+  end
+
+  let(:ruled_nameserver3) do
+    '1.2.3.8'
+  end
+
+  let (:config) do
+    {:dns_cache_no_start => true}
+  end
+
+  subject(:many_ruled_provider) do
+    dns_resolver = Rex::Proto::DNS::CachedResolver.new(config)
+    dns_resolver.extend(Rex::Proto::DNS::CustomNameserverProvider)
+    dns_resolver.nameservers = [base_nameserver]
+    dns_resolver.add_nameserver([], ruleless_nameserver, nil)
+    dns_resolver.add_nameserver(['*.metasploit.com'], ruled_nameserver, nil)
+    dns_resolver.add_nameserver(['*.metasploit.com'], ruled_nameserver2, nil)
+    dns_resolver.add_nameserver(['*.notmetasploit.com'], ruled_nameserver3, nil)
+
+    dns_resolver
+  end
+
+  subject(:ruled_provider) do
+    dns_resolver = Rex::Proto::DNS::CachedResolver.new(config)
+    dns_resolver.extend(Rex::Proto::DNS::CustomNameserverProvider)
+    dns_resolver.nameservers = [base_nameserver]
+    dns_resolver.add_nameserver([], ruleless_nameserver, nil)
+    dns_resolver.add_nameserver(['*.metasploit.com'], ruled_nameserver, nil)
+
+    dns_resolver
+  end
+
+  subject(:ruleless_provider) do
+    dns_resolver = Rex::Proto::DNS::CachedResolver.new(config)
+    dns_resolver.extend(Rex::Proto::DNS::CustomNameserverProvider)
+    dns_resolver.nameservers = [base_nameserver]
+    dns_resolver.add_nameserver([], ruleless_nameserver, nil)
+
+    dns_resolver
+  end
+
+  subject(:empty_provider) do
+    dns_resolver = Rex::Proto::DNS::CachedResolver.new(config)
+    dns_resolver.extend(Rex::Proto::DNS::CustomNameserverProvider)
+    dns_resolver.nameservers = [base_nameserver]
+
+    dns_resolver
+  end
+
+  context 'When no nameserver is configured' do
+   it 'The resolver base is returned' do
+     packet = packet_for('subdomain.metasploit.com')
+     ns = empty_provider.nameservers_for_packet(packet)
+     expect(ns).to eq([[base_nameserver, {}]])
+   end
+  end
+
+  context 'When a base nameserver is configured' do
+   it 'The base nameserver is returned' do
+     packet = packet_for('subdomain.metasploit.com')
+     ns = ruleless_provider.nameservers_for_packet(packet)
+     expect(ns).to eq([[ruleless_nameserver, {}]])
+   end
+  end
+
+  context 'When a nameserver rule is configured and a rule entry matches' do
+   it 'The correct nameserver is returned' do
+     packet = packet_for('subdomain.metasploit.com')
+     ns = ruled_provider.nameservers_for_packet(packet)
+     expect(ns).to eq([[ruled_nameserver, {}]])
+    end
+  end
+
+  context 'When a nameserver rule is configured and no rule entry is applicable' do
+   it 'The base nameserver is returned when no rule entry' do
+     packet = packet_for('subdomain.notmetasploit.com')
+     ns = ruled_provider.nameservers_for_packet(packet)
+     expect(ns).to eq([[ruleless_nameserver, {}]])
+   end
+  end
+
+  context 'When many rules are configured' do
+   it 'Returns multiple entries if multiple rules match' do
+     packet = packet_for('subdomain.metasploit.com')
+     ns = many_ruled_provider.nameservers_for_packet(packet)
+     expect(ns).to eq([[ruled_nameserver, {}], [ruled_nameserver2, {}]])
+   end
+  end
+end