Add support for multiple tlv types when logging

lib/rex/post/meterpreter/command_mapper.rb

@@ -9,6 +9,7 @@ module Post
 module Meterpreter
 
 class CommandMapper
+  @@cached_tlv_types = {}
 
   # Get the numeric command ID for the specified command name.
   #
@@ -97,6 +98,30 @@ def self.get_commands(*extensions)
     commands
   end
 
+  def self.create_tlv_types_cache
+    @@cached_tlv_types = {}
+
+    available_modules = [
+      ::Rex::Post::Meterpreter,
+      *::Rex::Post::Meterpreter::ExtensionMapper.get_extension_klasses
+    ].uniq
+
+    available_modules.each do |clazz|
+      clazz.constants.each do |const|
+        next unless const.to_s.start_with?('TLV_TYPE_') || const.to_s.start_with?('PACKET_')
+
+        if @@cached_tlv_types[clazz.const_get(const)].nil?
+          @@cached_tlv_types[clazz.const_get(const)] = const
+        else
+          @@cached_tlv_types[clazz.const_get(const)] = [@@cached_tlv_types[clazz.const_get(const)], const].flatten
+        end
+      end
+    end
+  end
+
+  def self.find_cached_tlv_type_by_value(value)
+    @@cached_tlv_types[value]
+  end
 end
 
 end

lib/rex/post/meterpreter/packet.rb

@@ -284,8 +284,6 @@ def self.generate_command_id_map_csharp
 #
 ###
 class Tlv
-  @@cached_tlv_types = {}
-
   attr_accessor :type, :value, :compress
 
   HEADER_SIZE = 8
@@ -316,11 +314,19 @@ def initialize(type, value = nil, compress=false)
     end
   end
 
-  def _tlv_type_string(type)
-    # Try to regenerate the cache if cache is empty, or the type we are looking for is not present in the cache.
-    regenerate_tlv_types_cache if @@cached_tlv_types.nil? || @@cached_tlv_types.key(type).nil?
+  def _tlv_type_string(value)
+    tlv_type = ::Rex::Post::Meterpreter::CommandMapper.find_cached_tlv_type_by_value(value)
+    ::Rex::Post::Meterpreter::CommandMapper.create_tlv_types_cache if tlv_type.nil?
+    tlv_type ||= ::Rex::Post::Meterpreter::CommandMapper.find_cached_tlv_type_by_value(value)
+
+    return if tlv_type.nil?
+
+    return_str = tlv_type.to_s.gsub('TLV_TYPE_', '').gsub('PACKET_TYPE_', '')
 
-    return @@cached_tlv_types.key(type).to_s.gsub('TLV_TYPE_', '').gsub('_', '-') unless @@cached_tlv_types.key(type).nil?
+    return return_str if tlv_type.is_a?(Symbol)
+    # In the off-chance we have multiple TLV types which have the same value,
+    # output them in the format `type=oneOf(A B C D)`
+    return "oneOf(#{return_str.gsub(/[:\[\] ]/, '').gsub(',', ' ')})" if tlv_type.is_a?(Array)
 
     nil
   end
@@ -338,7 +344,12 @@ def inspect
       when TLV_META_TYPE_COMPLEX; "COMPLEX"
       else; 'unknown-meta-type'
       end
-    stype = _tlv_type_string(type)
+
+    stype = case type
+      when PACKET_TYPE_REQUEST; 'Request'
+      when PACKET_TYPE_RESPONSE; 'Response'
+      else; _tlv_type_string(type)
+      end
     stype ||= "unknown-#{type}"
 
     val = value.inspect
@@ -503,23 +514,6 @@ def ntohq(value)
     htonq(value)
   end
 
-  def regenerate_tlv_types_cache
-    @@cached_tlv_types = {}
-
-    available_modules = [
-      ::Rex::Post::Meterpreter,
-      *::Rex::Post::Meterpreter::ExtensionMapper.get_extension_klasses
-    ].uniq
-
-    available_modules.each do |clazz|
-      clazz.constants.each do |const|
-        next unless const.to_s.start_with?('TLV_TYPE_') || const.to_s.start_with?('PACKET_')
-
-        @@cached_tlv_types[const] = clazz.const_get(const)
-      end
-    end
-  end
-
 end
 
 ###

spec/lib/rex/post/meterpreter/packet_spec.rb

@@ -130,7 +130,7 @@
       end
 
       it "should show the correct type and meta type in inspect" do
-        tlv_to_s = "#<Rex::Post::Meterpreter::Tlv type=COMMAND-ID      meta=INT        value=1001 command=stdapi_fs_chdir>"
+        tlv_to_s = "#<Rex::Post::Meterpreter::Tlv type=COMMAND_ID      meta=INT        value=1001 command=stdapi_fs_chdir>"
         expect(tlv.inspect).to eq tlv_to_s
       end
     end
@@ -147,7 +147,7 @@
       end
 
       it "should show the correct type and meta type in inspect" do
-        tlv_to_s = "#<Rex::Post::Meterpreter::Tlv type=COMMAND-ID      meta=INT        value=31337 command=unknown>"
+        tlv_to_s = "#<Rex::Post::Meterpreter::Tlv type=COMMAND_ID      meta=INT        value=31337 command=unknown>"
         expect(tlv.inspect).to eq tlv_to_s
       end
     end