From e94f5fc635a696fdb8872fad9aff5c1b3635b82a Mon Sep 17 00:00:00 2001 From: Jon Hart Date: Mon, 3 Nov 2014 11:07:08 -0800 Subject: [PATCH] Better handling of TYPE#/CLASS#, tests --- lib/msf/core/auxiliary/llmnr.rb | 80 ++++++----------------------- lib/net/dns/rr/classes.rb | 9 ++-- lib/net/dns/rr/types.rb | 7 +-- spec/lib/net/dns/rr/classes_spec.rb | 7 +++ spec/lib/net/dns/rr/types_spec.rb | 7 +++ 5 files changed, 35 insertions(+), 75 deletions(-) diff --git a/lib/msf/core/auxiliary/llmnr.rb b/lib/msf/core/auxiliary/llmnr.rb index 9e6cff64a7..c6384ef8ab 100644 --- a/lib/msf/core/auxiliary/llmnr.rb +++ b/lib/msf/core/auxiliary/llmnr.rb @@ -23,8 +23,8 @@ module Auxiliary::LLMNR OptAddressRange.new('RHOSTS', [true, 'The multicast address or CIDR range of targets to query', '224.0.0.252']), Opt::RPORT(5355), OptString.new('NAME', [true, 'The name to query', 'localhost']), - OptString.new('TYPE', [true, 'The query type (name or #)', 'A']), - OptString.new('CLASS', [true, 'The query class (name or #)', 'IN']) + OptString.new('TYPE', [true, 'The query type (name, # or TYPE#)', 'A']), + OptString.new('CLASS', [true, 'The query class (name, # or CLASS#)', 'IN']) ], self.class) end @@ -34,83 +34,35 @@ module Auxiliary::LLMNR end def query_class - datastore['CLASS'].upcase + if datastore['CLASS'] =~ /^\d+$/ + datastore['CLASS'].to_i + else + datastore['CLASS'].upcase + end end def query_class_name - return @query_class_name if @query_class_name - if /^(?\d+)$/ =~ query_class - class_num = class_num.to_i - raise ClassArgumentError, "LLMNR RR class #{query_class} out of range" if class_num > 0x7FFF - begin - @query_class_name = Net::DNS::RR::Classes.to_str(class_num) - rescue ClassArgumentError - @query_class_name = "CLASS#{class_num}" - end - else - unless Net::DNS::RR::Classes.valid?(query_class) - raise ClassArgumentError, "LLMNR RR class #{query_class} invalid" - end - @query_class_name = query_class - end - - @query_class_name + Net::DNS::RR::Classes.new(query_class).to_s end def query_class_num - return @query_class_num if @query_class_num - if /^(?\d+)$/ =~ query_class - class_num = class_num.to_i - raise ClassArgumentError, "LLMNR RR class #{query_class} out of range" if class_num > 0x7FFF - @query_class_num = class_num - else - unless Net::DNS::RR::Classes.valid?(query_class) - raise ClassArgumentError, "LLMNR RR class #{query_class} invalid" - end - @query_class_num = Net::DNS::RR::Classes::Classes[query_class] - end - - @query_class_num + Net::DNS::RR::Classes.new(query_class).to_i end def query_type - datastore['TYPE'].upcase + if datastore['TYPE'] =~ /^\d+$/ + datastore['TYPE'].to_i + else + datastore['TYPE'].upcase + end end def query_type_name - return @query_type_name if @query_type_name - if /^(?\d+)$/ =~ query_type - type_num = type_num.to_i - raise TypeArgumentError, "LLMNR RR type #{query_type} out of range" if type_num > 0xFFFF - begin - @query_type_name = Net::DNS::RR::Types.to_str(type_num) - rescue TypeArgumentError - @query_type_name = "TYPE#{type_num}" - end - else - unless Net::DNS::RR::Types.valid?(query_type) - raise TypeArgumentError, "LLMNR RR type #{query_type} invalid" - end - @query_type_name = query_type - end - - @query_type_name + Net::DNS::RR::Types.new(query_type).to_s end def query_type_num - return @query_type_num if @query_type_num - if /^(?\d+)$/ =~ query_type - type_num = type_num.to_i - raise TypeArgumentError, "LLMNR RR type #{query_type} out of range" if type_num > 0xFFFF - @query_type_num = type_num - else - unless Net::DNS::RR::Types.valid?(query_type) - raise TypeArgumentError, "LLMNR RR type #{query_type} invalid" - end - @query_type_num = Net::DNS::RR::Types::Types[query_type] - end - - @query_type_num + Net::DNS::RR::Types.new(query_type).to_i end end end diff --git a/lib/net/dns/rr/classes.rb b/lib/net/dns/rr/classes.rb index 3f59f6384c..2ec1aaf1dd 100644 --- a/lib/net/dns/rr/classes.rb +++ b/lib/net/dns/rr/classes.rb @@ -89,18 +89,15 @@ module Net # :nodoc: # *PRIVATE* method def new_from_string(cls) case cls - when /^CLASS(?\d+)$/ - class_num = class_num.to_i - raise ClassArgumentError, "Invalid class #{class_num}" if class_num > 0xFFFF - @str = cls - @num = class_num + when /^CLASS(\d+)$/ + new_from_num(Regexp.last_match(1).to_i) else # String with name of class if Classes.has_key? cls @str = cls @num = Classes[cls] else - raise ClassArgumentError, "Unknown cls #{cls}" + raise ClassArgumentError, "Unknown class #{cls}" end end end diff --git a/lib/net/dns/rr/types.rb b/lib/net/dns/rr/types.rb index 28c10f80f7..e4507d23c2 100644 --- a/lib/net/dns/rr/types.rb +++ b/lib/net/dns/rr/types.rb @@ -167,11 +167,8 @@ module Net # :nodoc: # *PRIVATE* method def new_from_string(type) case type - when /^TYPE(?\d+)$/ - type_num = type_num.to_i - raise TypeArgumentError, "Invalid type #{type_num}" if type_num > 0xFFFF - @str = type - @num = type_num + when /^TYPE(\d+)$/ + new_from_num(Regexp.last_match(1).to_i) else # String with name of type if Types.has_key? type diff --git a/spec/lib/net/dns/rr/classes_spec.rb b/spec/lib/net/dns/rr/classes_spec.rb index 85dfa24c00..7b2fc8cb90 100644 --- a/spec/lib/net/dns/rr/classes_spec.rb +++ b/spec/lib/net/dns/rr/classes_spec.rb @@ -46,10 +46,17 @@ describe Net::DNS::RR::Classes do rr_class.send(:initialize, 'CH') expect(rr_class.to_i).to eql(3) expect(rr_class.to_s).to eql('CH') + rr_class.send(:initialize, 'CLASS9') + expect(rr_class.to_i).to eql(9) + expect(rr_class.to_s).to eql('CLASS9') + rr_class.send(:initialize, 'CLASS1') + expect(rr_class.to_i).to eql(1) + expect(rr_class.to_s).to eql('IN') end it 'raises when the supplied RR class String is invalid' do expect { rr_class.send(:initialize, 'cats') }.to raise_error(ClassArgumentError) + expect { rr_class.send(:initialize, 'CLASS123456') }.to raise_error(ClassArgumentError) end end end diff --git a/spec/lib/net/dns/rr/types_spec.rb b/spec/lib/net/dns/rr/types_spec.rb index a010c0889e..728525faad 100644 --- a/spec/lib/net/dns/rr/types_spec.rb +++ b/spec/lib/net/dns/rr/types_spec.rb @@ -46,10 +46,17 @@ describe Net::DNS::RR::Types do rr_type.send(:initialize, 'SRV') expect(rr_type.to_i).to eql(33) expect(rr_type.to_s).to eql('SRV') + rr_type.send(:initialize, 'TYPE12') + expect(rr_type.to_i).to eql(12) + expect(rr_type.to_s).to eql('PTR') + rr_type.send(:initialize, 'TYPE123') + expect(rr_type.to_i).to eql(123) + expect(rr_type.to_s).to eql('TYPE123') end it 'raises when the supplied RR type String is invalid' do expect { rr_type.send(:initialize, 'cats') }.to raise_error(TypeArgumentError) + expect { rr_type.send(:initialize, 'TYPE123456') }.to raise_error(TypeArgumentError) end end end