Better handling of TYPE#/CLASS#, tests

bug/bundler_fix
Jon Hart 2014-11-03 11:07:08 -08:00
parent 7b5be67d65
commit e94f5fc635
5 changed files with 35 additions and 75 deletions

View File

@ -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 /^(?<class_num>\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 /^(?<class_num>\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 /^(?<type_num>\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 /^(?<type_num>\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

View File

@ -89,18 +89,15 @@ module Net # :nodoc:
# *PRIVATE* method
def new_from_string(cls)
case cls
when /^CLASS(?<class_num>\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

View File

@ -167,11 +167,8 @@ module Net # :nodoc:
# *PRIVATE* method
def new_from_string(type)
case type
when /^TYPE(?<type_num>\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

View File

@ -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

View File

@ -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