Better handling of TYPE#/CLASS#, tests

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('TYPE', [true, 'The query type (name, # or TYPE#)', 'A']),
-      OptString.new('CLASS', [true, 'The query class (name or #)', 'IN'])
+      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
+    Net::DNS::RR::Classes.new(query_class).to_s
-    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
   end
 
   def query_class_num
-    return @query_class_num if @query_class_num
+    Net::DNS::RR::Classes.new(query_class).to_i
-    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
   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
+    Net::DNS::RR::Types.new(query_type).to_s
-    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
   end
 
   def query_type_num
-    return @query_type_num if @query_type_num
+    Net::DNS::RR::Types.new(query_type).to_i
-    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
   end
 end
 end

lib/net/dns/rr/classes.rb

@@ -89,18 +89,15 @@ module Net # :nodoc:
         # *PRIVATE* method
         def new_from_string(cls)
           case cls
-          when /^CLASS(?<class_num>\d+)$/
+          when /^CLASS(\d+)$/
-            class_num = class_num.to_i
+            new_from_num(Regexp.last_match(1).to_i)
-            raise ClassArgumentError, "Invalid class #{class_num}" if class_num > 0xFFFF
-            @str = cls
-            @num = class_num
           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

lib/net/dns/rr/types.rb

@@ -167,11 +167,8 @@ module Net # :nodoc:
         # *PRIVATE* method
         def new_from_string(type)
           case type
-          when /^TYPE(?<type_num>\d+)$/
+          when /^TYPE(\d+)$/
-            type_num = type_num.to_i
+            new_from_num(Regexp.last_match(1).to_i)
-            raise TypeArgumentError, "Invalid type #{type_num}" if type_num > 0xFFFF
-            @str = type
-            @num = type_num
           else
             # String with name of type
             if Types.has_key? type

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

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