diff --git a/README.md b/README.md index c7fd2ec..2e7d159 100644 --- a/README.md +++ b/README.md @@ -44,17 +44,17 @@ Ronin::DNS::Proxy.run('127.0.0.1', 2346) do |server| server.add_rule :A, 'updates.example.com', :ServFail # define a dynamic rule - server.add_rule :CNAME, /^www\./, ->(type,name,transaction) { + server.add_rule(:CNAME, /^www\./) do |type,name,transaction| # append '.hax' to the domain name names = name.split('.').push('hax') transaction.respond!(names) - } + end # return MX records - server.add_rule :MX, 'example.com', ->(type,name,transaction) { + server.add_rule(:MX, 'example.com') do |type,name,transaction| transaction.respond!(10, Resolv::DNS::Name.create('email.evil.com' )) - } + end end ``` diff --git a/lib/ronin/dns/proxy/rule.rb b/lib/ronin/dns/proxy/rule.rb index 5da703f..9de2368 100644 --- a/lib/ronin/dns/proxy/rule.rb +++ b/lib/ronin/dns/proxy/rule.rb @@ -55,10 +55,30 @@ class Rule # @param [String, Array, Symbol, #call] result # The result to return. # - def initialize(type,name,result) + # @yield [type, name, transaction] + # If no result argument is given, the given block will be passed the + # DNS query's type, name, and transaction object. + # + # @yieldparam [Symbol] type + # The query type. + # + # @yieldparam [String] name + # The queried host name. + # + # @yieldparam [Async::DNS::Transaction] transaction + # The DNS query transaction object. + # + # @raise [ArgumentError] + # Must specify a `result` argument or a block. + # + def initialize(type,name,result=nil,&block) + unless (result || block) + raise(ArgumentError,"must specify a result value or a block") + end + @type = type @name = name - @result = result + @result = result || block end # diff --git a/lib/ronin/dns/proxy/server.rb b/lib/ronin/dns/proxy/server.rb index 837ecdd..651e086 100644 --- a/lib/ronin/dns/proxy/server.rb +++ b/lib/ronin/dns/proxy/server.rb @@ -117,7 +117,7 @@ def initialize(host,port, nameservers: Ronin::Support::Network::DNS.nameservers, # @param [String, Regexp] name # The record name that the rule will match against. # - # @param [String, Array, Symbol, Proc] result + # @param [String, Array, Symbol, #call] result # The result to respond with. It can be a String, or an Array of # Strings, or an error code: # @@ -129,8 +129,18 @@ def initialize(host,port, nameservers: Ronin::Support::Network::DNS.nameservers, # * `:Refused` - The operation was refused by the server. # * `:NotAuth` - The server is not authoritive for the zone. # - # If a `Proc` is given, then it will be called with the query type, - # query name, and the DNS query transaction object. + # @yield [type, name, transaction] + # If no result argument is given, the given block will be passed the + # DNS query's type, name, and transaction object. + # + # @yieldparam [Symbol] type + # The query type. + # + # @yieldparam [String] name + # The queried host name. + # + # @yieldparam [Async::DNS::Transaction] transaction + # The DNS query transaction object. # # @example override the IP address for a domain: # server.add_rule :A, 'example.com', '10.0.0.42' @@ -145,17 +155,21 @@ def initialize(host,port, nameservers: Ronin::Support::Network::DNS.nameservers, # server.add_rule :TXT, /^spf\./, "v=spf1 include:10.0.0.1 ~all" # # @example define a dynamic rule: - # server.add_rule :CNAME, /^www\./, ->(type,name,transaction) { + # server.add_rule(:CNAME, /^www\./) do |type,name,transaction| # # append '.hax' to the domain name # names = name.split('.').push('hax') # # transaction.respond!(names) - # } + # end # # @api public # - def add_rule(record_type,name,result) - @rules << Rule.new(record_type,name,result) + def add_rule(record_type,name,result=nil,&block) + unless (result || block) + raise(ArgumentError,"must specify a result value or a block") + end + + @rules << Rule.new(record_type,name,result,&block) end # Mapping of Resolv resource classes to Symbols. diff --git a/spec/rule_spec.rb b/spec/rule_spec.rb index 518a439..f9218fb 100644 --- a/spec/rule_spec.rb +++ b/spec/rule_spec.rb @@ -20,6 +20,30 @@ it "must set #result" do expect(subject.result).to eq(result) end + + context "when no result is given" do + context "but a block is given" do + let(:block) do + proc { |type,name,transaction| + transaction.respond!('foo') + } + end + + subject { described_class.new(type,name,&block) } + + it "must set #result to the block" do + expect(subject.result).to be(block) + end + end + + context "but a block is not given" do + it do + expect { + described_class.new(type,name) + }.to raise_error(ArgumentError,"must specify a result value or a block") + end + end + end end describe "#matches?" do diff --git a/spec/server_spec.rb b/spec/server_spec.rb index 43ec7f2..50cf374 100644 --- a/spec/server_spec.rb +++ b/spec/server_spec.rb @@ -63,16 +63,48 @@ let(:record_name) { 'example.com' } let(:record_result) { '10.0.0.1' } - before do - subject.add_rule :TXT, 'foo.example.com', '1.2.3.4' - subject.add_rule record_type, record_name, record_result + context "when type, name, and result arguments are given" do + before do + subject.add_rule :TXT, 'foo.example.com', '1.2.3.4' + subject.add_rule record_type, record_name, record_result + end + + it "must append a new Ronin::DNS::Proxy::Rule object to #rules with the type, name, and result arguments" do + expect(subject.rules.last).to be_kind_of(Ronin::DNS::Proxy::Rule) + expect(subject.rules.last.type).to eq(record_type) + expect(subject.rules.last.name).to eq(record_name) + expect(subject.rules.last.result).to eq(record_result) + end end - it "must append a new Ronin::DNS::Proxy::Rule object to #rules" do - expect(subject.rules.last).to be_kind_of(Ronin::DNS::Proxy::Rule) - expect(subject.rules.last.type).to eq(record_type) - expect(subject.rules.last.name).to eq(record_name) - expect(subject.rules.last.result).to eq(record_result) + context "when no result argument is given" do + context "but a block is given" do + let(:block) do + proc { |type,name,transaction| + transaction.respond!('foo') + } + end + + before do + subject.add_rule :TXT, 'foo.example.com', '1.2.3.4' + subject.add_rule(record_type,record_name,&block) + end + + it "must set the rule's result to the given block" do + expect(subject.rules.last).to be_kind_of(Ronin::DNS::Proxy::Rule) + expect(subject.rules.last.type).to eq(record_type) + expect(subject.rules.last.name).to eq(record_name) + expect(subject.rules.last.result).to be(block) + end + end + + context "and no block is given" do + it do + expect { + subject.add_rule(record_type,record_name) + }.to raise_error(ArgumentError,"must specify a result value or a block") + end + end end end