Code review

lib/ruby_smb/dcerpc/winreg.rb

@@ -71,6 +71,8 @@ module Winreg
 
       BUFFER_SIZE = 1024
 
+      RegValue = Struct.new(:type, :data)
+
       # Open the registry root key and return a handle for it. The key can be
       # either a long format (e.g. HKEY_LOCAL_MACHINE) or a short format
       # (e.g. HKLM)
@@ -132,11 +134,10 @@ def open_key(handle, sub_key)
       # @param handle [Ndr::NdrContextHandle] the handle for the key
       # @param value_name [String] the name of the value
       # @param value_name [Boolean] also return the data type if set to true
-      # @return [String] the data of the value entry
-      # @return [Array] if `type` is true, an array containing the data type and the actual data of the value entry
+      # @return [RegValue] a RegValue struct containing the data type and the actual data of the value entry
       # @raise [RubySMB::Dcerpc::Error::InvalidPacket] if the response is not a QueryValueResponse packet
       # @raise [RubySMB::Dcerpc::Error::WinregError] if the response error status is not ERROR_SUCCESS
-      def query_value(handle, value_name, type: false)
+      def query_value(handle, value_name)
         query_value_request_packet = RubySMB::Dcerpc::Winreg::QueryValueRequest.new(hkey: handle, lp_value_name: value_name)
         query_value_request_packet.lp_type = 0
         query_value_request_packet.lpcb_data = 0
@@ -166,11 +167,7 @@ def query_value(handle, value_name, type: false)
             "#{WindowsError::Win32.find_by_retval(query_value_response.error_status.value).join(',')}"
         end
 
-        if type
-          [query_value_response.lp_type, query_value_response.data]
-        else
-          query_value_response.data
-        end
+        RegValue.new(query_value_response.lp_type, query_value_response.data)
       end
 
       # Close the handle to the registry key.
@@ -363,8 +360,8 @@ def read_registry_key_value(key, value_name, bind: true)
         root_key, sub_key = key.gsub(/\//, '\\').split('\\', 2)
         root_key_handle = open_root_key(root_key)
         subkey_handle = open_key(root_key_handle, sub_key)
-        value = query_value(subkey_handle, value_name)
-        value
+        reg_value = query_value(subkey_handle, value_name)
+        reg_value.data
       ensure
         close_key(subkey_handle) if subkey_handle
         close_key(root_key_handle) if root_key_handle

lib/ruby_smb/dcerpc/winreg/query_value_response.rb

@@ -25,9 +25,11 @@ def initialize_instance
         def data
           bytes = lp_data.to_a.pack('C*')
           case lp_type
+          when 0 # 0 is undefined type, let's consider an array of bytes
+            bytes
           when 1,2
             bytes.force_encoding('utf-16le').strip
-          when 0,3 # 0 is undefined type, let's consider an array of bytes
+          when 3
             bytes
           when 4
             bytes.unpack('V').first

spec/lib/ruby_smb/dcerpc/winreg_spec.rb

@@ -124,12 +124,13 @@
   describe '#query_value' do
     let(:handle)                      { double('Handle') }
     let(:value_name)                  { double('Value Name') }
-    let(:query_value_request_packet) { double('Query Value Request Packet #1') }
-    let(:lp_data)                    { double('LpData #2') }
+    let(:query_value_request_packet)  { double('Query Value Request Packet #1') }
+    let(:lp_data)                     { double('LpData #2') }
     let(:response1)                   { double('Response #1') }
     let(:response2)                   { double('Response #2') }
     let(:query_value_response1)       { double('Query Value Response #1') }
     let(:query_value_response2)       { double('Query Value Response #2') }
+    let(:lp_type)                     { double('Type') }
     let(:data)                        { double('Data') }
     let(:lpcb_data)                   { double('LpcbData') }
     let(:max_count)                   { 5 }
@@ -156,8 +157,9 @@
       allow(described_class::QueryValueResponse).to receive(:read).with(response2).and_return(query_value_response2)
       allow(query_value_response1).to receive(:error_status).and_return(WindowsError::Win32::ERROR_SUCCESS)
       allow(query_value_response2).to receive_messages(
-        :error_status => WindowsError::Win32::ERROR_SUCCESS,
-        :data         => data
+        error_status: WindowsError::Win32::ERROR_SUCCESS,
+        lp_type:      lp_type,
+        data:         data
       )
       allow(query_value_response1).to receive(:lpcb_data).and_return(lpcb_data)
       allow(lpcb_data).to receive(:to_i).and_return(max_count)
@@ -226,18 +228,9 @@
     end
 
     it 'returns the expected response data' do
-      expect(winreg.query_value(handle, value_name)).to eq(data)
+      expect(winreg.query_value(handle, value_name)).to eq(RubySMB::Dcerpc::Winreg::RegValue.new(lp_type, data))
     end
 
-    context 'when the data type is also required' do
-      let(:lp_type) { double('Type') }
-      before :example do
-        allow(query_value_response2).to receive(:lp_type).and_return(lp_type)
-      end
-      it 'returns the expected response type and data' do
-        expect(winreg.query_value(handle, value_name, type: true)).to eq([lp_type, data])
-      end
-    end
   end
 
   describe '#close_key' do