Implement changes from PR feedback

rapid7 · Apr 22, 2020 · c920ca7 · c920ca7
1 parent 2331948
commit c920ca7
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 11 deletions.
diff --git a/lib/msf/util/dot_net_deserialization.rb b/lib/msf/util/dot_net_deserialization.rb
@@ -59,6 +59,9 @@ def self.generate(cmd, gadget_chain: DEFAULT_GADGET_CHAIN, formatter: DEFAULT_FO
   # Take the specified serialized blob and encapsulate it with the specified
   # formatter.
   #
+  # @param stream [Msf::Util::DotNetDeserialization::Types::SerializedStream]
+  #   The serialized stream representing the gadget chain to format into a
+  #   string.
   # @param formatter [Symbol] The formatter to use to encapsulate the serialized
   #   data blob.
   # @return [String]
@@ -81,6 +84,8 @@ def self.generate_formatted(stream, formatter: DEFAULT_FORMATTER)
   # the OS command. The chosen gadget chain must be compatible with the target
   # application.
   #
+  # @param cmd [String] The operating system command to execute. It will
+  #   automatically be prefixed with "cmd /c" by the gadget chain.
   # @param gadget_chain [Symbol] The gadget chain to use for execution.
   # @return [Types::SerializedStream]
   def self.generate_gadget_chain(cmd, gadget_chain: DEFAULT_GADGET_CHAIN)

diff --git a/lib/msf/util/dot_net_deserialization/formatters/los_formatter.rb b/lib/msf/util/dot_net_deserialization/formatters/los_formatter.rb
@@ -4,6 +4,8 @@ module DotNetDeserialization
 module Formatters
 module LosFormatter
 
+  TOKEN_BINARY_SERIALIZED = 50
+
   #
   # Limited Object Stream Types
   #
@@ -20,8 +22,7 @@ class ObjectStateFormatter < BinData::Record
 
   def self.generate(stream)
     stream = stream.to_binary_s
-    # token: 50 =  Token_BinarySerialized
-    formatted  = ObjectStateFormatter.new(token: 50).to_binary_s
+    formatted  = ObjectStateFormatter.new(token: TOKEN_BINARY_SERIALIZED).to_binary_s
     formatted << DotNetDeserialization.encode_7bit_int(stream.length)
     formatted << stream
   end

diff --git a/lib/msf/util/dot_net_deserialization/types.rb b/lib/msf/util/dot_net_deserialization/types.rb
@@ -65,14 +65,14 @@ def self.from_values(values)
     def get_object(id)
       id = id.value if id.is_a? BinData::BasePrimitive
 
-      @objects = @objects || {}
+      @objects ||= {}
       @objects[id]
     end
 
     def set_object(id, object)
       id = id.value if id.is_a? BinData::BasePrimitive
 
-      @objects = @objects || {}
+      @objects ||= {}
       @objects[id] = object
     end
   end

diff --git a/lib/msf/util/dot_net_deserialization/types/primitives.rb b/lib/msf/util/dot_net_deserialization/types/primitives.rb
@@ -60,7 +60,7 @@ class EnumArray < BinData::Array
     def assign(values)
       if values.is_a? ::Array
         enum = eval_parameter(:enum)
-        values.map! { |value| (value.is_a? Symbol) ? enum.fetch(value) : value }
+        values = values.map { |value| (value.is_a? Symbol) ? enum.fetch(value) : value }
       end
       super(values)
     end
@@ -178,6 +178,8 @@ def selection_routine(index)
 
     module Factory
       def from_member_values(class_info:, member_type_info:, member_values:, **kwargs)
+        raise ::ArgumentError, 'Invalid class_info type' unless class_info.is_a? Types::General::ClassInfo
+        raise ::ArgumentError, 'Invalid member_type_info type' unless member_type_info.is_a? Types::General::MemberTypeInfo
         raise ::ArgumentError, 'Invalid member count' unless class_info.member_count == member_values.length
 
         kwargs[:class_info] = class_info

diff --git a/lib/msf/util/dot_net_deserialization/types/record_values.rb b/lib/msf/util/dot_net_deserialization/types/record_values.rb
@@ -51,8 +51,7 @@ def member_type_info
       object.member_type_info
     end
 
-    include Primitives::MemberValues::Factory
-    self.singleton_class.include Primitives::MemberValues::Factory
+    extend Primitives::MemberValues::Factory
   end
 
   class ClassWithMembersAndTypes < BinData::Record
@@ -64,8 +63,7 @@ class ClassWithMembersAndTypes < BinData::Record
     int32                  :library_id
     member_values          :member_values, class_info: -> { class_info }, member_type_info: -> { member_type_info }
 
-    include Primitives::MemberValues::Factory
-    self.singleton_class.include Primitives::MemberValues::Factory
+    extend Primitives::MemberValues::Factory
   end
 
   class MemberReference < BinData::Record
@@ -114,8 +112,7 @@ class SystemClassWithMembersAndTypes < BinData::Record
     member_type_info       :member_type_info, member_count: -> { class_info.member_count }
     member_values          :member_values, class_info: -> { class_info }, member_type_info: -> { member_type_info }
 
-    include Primitives::MemberValues::Factory
-    self.singleton_class.include Primitives::MemberValues::Factory
+    extend Primitives::MemberValues::Factory
   end
 end
 end