ruby_smb 1.0.5 → 2.0.3

data/lib/ruby_smb/smb2/negotiate_context.rb

@@ -0,0 +1,108 @@
+module RubySMB
+  module SMB2
+
+    # An SMB2 PREAUTH_INTEGRITY_CAPABILITIES context struct as defined in
+    # [2.2.3.1.1 SMB2_PREAUTH_INTEGRITY_CAPABILITIES](https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/5a07bd66-4734-4af8-abcf-5a44ff7ee0e5)
+    class PreauthIntegrityCapabilities < BinData::Record
+      SHA_512 = 0x0001
+      HASH_ALGORITM_MAP = {
+        SHA_512 => 'SHA512'
+      }
+
+      endian  :little
+
+      uint16 :hash_algorithm_count, label: 'Hash Algorithm Count', initial_value: -> { hash_algorithms.size }
+      uint16 :salt_length,          label: 'Salt Length',          initial_value: -> { salt.num_bytes }
+      array  :hash_algorithms,      label: 'Hash Algorithms',      type: :uint16, initial_length: -> { hash_algorithm_count }
+      string :salt,                 label: 'Salt',                 read_length: -> { salt_length }
+    end
+
+    # An SMB2 ENCRYPTION_CAPABILITIES context struct as defined in
+    # [2.2.3.1.2 SMB2_ENCRYPTION_CAPABILITIES](https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/16693be7-2b27-4d3b-804b-f605bde5bcdd)
+    class EncryptionCapabilities < BinData::Record
+      AES_128_CCM = 0x0001
+      AES_128_GCM = 0x0002
+      ENCRYPTION_ALGORITHM_MAP = {
+        AES_128_CCM => 'AES-128-CCM',
+        AES_128_GCM => 'AES-128-GCM'
+      }
+
+      endian  :little
+
+      uint16 :cipher_count, label: 'Cipher Count', initial_value: -> { ciphers.size }
+      array  :ciphers,      label: 'Ciphers',      type: :uint16, initial_length: -> { cipher_count }
+    end
+
+    # An SMB2 COMPRESSION_CAPABILITIES context struct as defined in
+    # [2.2.3.1.3 SMB2_COMPRESSION_CAPABILITIES](https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/78e0c942-ab41-472b-b117-4a95ebe88271)
+    class CompressionCapabilities < BinData::Record
+      # Flags
+      # Chained compression is not supported.
+      SMB2_COMPRESSION_CAPABILITIES_FLAG_NONE    = 0x00000000
+      # Chained compression is supported on this connection.
+      SMB2_COMPRESSION_CAPABILITIES_FLAG_CHAINED = 0x00000001
+
+      # Compression Algorithms
+      NONE         = 0x0000
+      LZNT1        = 0x0001
+      LZ77         = 0x0002
+      LZ77_Huffman = 0x0003
+      Pattern_V1   = 0x0004
+      COMPRESSION_ALGORITHM_MAP = {
+        NONE         => 'NONE',
+        LZNT1        => 'LZNT1',
+        LZ77         => 'LZ77',
+        LZ77_Huffman => 'LZ77_Huffman',
+        Pattern_V1   => 'Pattern_V1'
+      }
+
+      endian  :little
+
+      uint16 :compression_algorithm_count, label: 'Compression Algorithm Count', initial_value: -> { compression_algorithms.size }
+      uint16 :padding,                     label: 'Padding',                     initial_value: 0
+      uint32 :flags,                       label: 'Flags'
+      array  :compression_algorithms,      label: 'Compression Algorithms',      type: :uint16, initial_length: -> { compression_algorithm_count }
+    end
+
+    # An SMB2 NETNAME_NEGOTIATE_CONTEXT_ID context struct as defined in
+    # [2.2.3.1.4 SMB2_NETNAME_NEGOTIATE_CONTEXT_ID](https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/ca6726bd-b9cf-43d9-b0bc-d127d3c993b3)
+    class NetnameNegotiateContextId < BinData::Record
+      endian  :little
+
+      stringz16 :net_name, label: 'Net Name'
+    end
+
+
+    # An SMB2 NEGOTIATE_CONTEXT struct as defined in
+    # [2.2.3.1 SMB2 NEGOTIATE_CONTEXT Request Values](https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/15332256-522e-4a53-8cd7-0bd17678a2f7)
+    class NegotiateContext < BinData::Record
+      # The NegotiateContext Data field contains a list of preauthentication integrity hash functions as well as an optional salt value, as specified in section 2.2.3.1.1.
+      SMB2_PREAUTH_INTEGRITY_CAPABILITIES  = 0x0001
+      # The NegotiateContext Data field contains a list of encryption algorithms, as specified in section 2.2.3.1.2.
+      SMB2_ENCRYPTION_CAPABILITIES         = 0x0002
+      # The NegotiateContext Data field contains a list of compression algorithms, as specified in section 2.2.3.1.3.
+      SMB2_COMPRESSION_CAPABILITIES        = 0x0003
+      # The NegotiateContext Data field contains the server name to which the client connects.
+      SMB2_NETNAME_NEGOTIATE_CONTEXT_ID    = 0x0005
+
+      endian  :little
+
+      string :pad,          label: 'Padding',     length: -> { pad_length }
+      uint16 :context_type, label: 'Context Type'
+      uint16 :data_length,  label: 'Data Length', initial_value: -> { data.num_bytes }
+      uint32 :reserved,     label: 'Reserved',    initial_value: 0
+      choice :data,         label: 'Data',        selection: -> { context_type } do
+        preauth_integrity_capabilities SMB2_PREAUTH_INTEGRITY_CAPABILITIES, label: 'Preauthentication Integrity Capabilities'
+        encryption_capabilities        SMB2_ENCRYPTION_CAPABILITIES,        label: 'Encryption Capabilities'
+        compression_capabilities       SMB2_COMPRESSION_CAPABILITIES,       label: 'Compression Capabilities'
+        netname_negotiate_context_id   SMB2_NETNAME_NEGOTIATE_CONTEXT_ID,   label: 'Netname Negotiate Context ID'
+      end
+
+      def pad_length
+        offset = pad.abs_offset % 8
+        (8 - offset) % 8
+      end
+
+    end
+  end
+end

data/lib/ruby_smb/smb2/packet.rb

@@ -30,6 +30,8 @@ module RubySMB
       require 'ruby_smb/smb2/packet/write_response'
       require 'ruby_smb/smb2/packet/ioctl_request'
       require 'ruby_smb/smb2/packet/ioctl_response'
+      require 'ruby_smb/smb2/packet/transform_header'
+      require 'ruby_smb/smb2/packet/compression_transform_header'
     end
   end
 end

data/lib/ruby_smb/smb2/packet/compression_transform_header.rb

@@ -0,0 +1,41 @@
+module RubySMB
+  module SMB2
+    module Packet
+      # An SMB2 COMPRESSION_TRANSFORM_HEADER Packet as defined in
+      # [2.2.42 SMB2 COMPRESSION_TRANSFORM_HEADER](https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/1d435f21-9a21-4f4c-828e-624a176cf2a0)
+      class CompressionTransformHeader < BinData::Record
+        endian :little
+
+        bit32            :protocol,                         label: 'Protocol ID Field',      initial_value: 0xFC534D42
+        uint32           :original_compressed_segment_size, label: 'Original Compressed Segment Size'
+        uint16           :compression_algorithm,            label: 'Compression Algorithm'
+        uint16           :flags,                            label: 'Flags'
+        uint32           :offset,                           label: 'Offset / Length'
+      end
+
+      # An SMB2 SMB2_COMPRESSION_TRANSFORM_HEADER_PAYLOAD Packet as defined in
+      # [2.2.42.1 SMB2_COMPRESSION_TRANSFORM_HEADER_PAYLOAD](https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/8898e8e7-f1b2-47f5-a525-2ce5bad6db64)
+      class Smb2CompressionPayloadHeader < BinData::Record
+        endian :little
+        hide   :reserved
+
+        uint16           :algorithm_id,                     label: 'Algorithm ID'
+        uint16           :reserved
+        uint32           :payload_length,                   label: 'Compressed Payload Length'
+      end
+
+      # An SMB2 SMB2_COMPRESSION_PATTERN_PAYLOAD_V1 Packet as defined in
+      # [2.2.42.2 SMB2_COMPRESSION_PATTERN_PAYLOAD_V1](https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/f6859837-395a-4d0a-8971-1fc3919e2d09)
+      class Smb2CompressionPatternPayloadV1 < BinData::Record
+        endian :little
+        hide   :reserved1, :reserved2
+
+        uint8            :pattern,                          label: 'Pattern'
+        uint8            :reserved1
+        uint16           :reserved2
+        uint32           :repetitions,                      label: 'Repetitions'
+      end
+    end
+  end
+end
+

data/lib/ruby_smb/smb2/packet/error_packet.rb

@@ -11,10 +11,8 @@ module RubySMB
         uint16       :structure_size,      label: 'Structure Size', initial_value: 9
         uint8        :error_context_count, label: 'ErrorContextCount'
         uint8        :reserved
-        uint32       :byte_count,          label: 'Byte Count of ErrorData',
-                                           initial_value: -> { error_data.num_bytes == 1 ? 0 : error_data.num_bytes }
-        string       :error_data,          label: 'Error Data', initial_value: "\x00",
-                                           read_length: -> { byte_count == 0 ? 1 : byte_count }
+        uint32       :byte_count,          label: 'Byte Count of ErrorData'
+        string       :error_data,          label: 'Error Data', read_length: -> { byte_count }
 
         def valid?
           return smb2_header.protocol == RubySMB::SMB2::SMB2_PROTOCOL_ID &&

data/lib/ruby_smb/smb2/packet/negotiate_request.rb

@@ -1,3 +1,5 @@
+require 'ruby_smb/smb2/negotiate_context'
+
 module RubySMB
   module SMB2
     module Packet
@@ -8,23 +10,30 @@ module RubySMB
 
         endian              :little
         smb2_header         :smb2_header
-        uint16              :structure_size,      label: 'Structure Size', initial_value: 36
-        uint16              :dialect_count,       label: 'Dialect Count'
-        smb2_security_mode  :security_mode
-        uint16              :reserved1, label: 'Reserved', initial_value: 0
-        smb2_capabilities   :capabilities
-        string              :client_guid,         label: 'Client GUID',        length: 16
-        file_time           :client_start_time,   label: 'Client Start Time',  initial_value: 0
-        array               :dialects,            label: 'Dialects',           type: :uint16, read_until: :eof
-
-        # Adds a dialect to the Dialects array and increments the dialect count
+        uint16              :structure_size,           label: 'Structure Size', initial_value: 36
+        uint16              :dialect_count,            label: 'Dialect Count', initial_value: -> { dialects.size }
+        smb2_security_mode  :security_mode,            label: 'Security Mode'
+        uint16              :reserved1,                label: 'Reserved', initial_value: 0
+        smb2_capabilities   :capabilities,             label: 'Capabilities'
+        string              :client_guid,              label: 'Client GUID', length: 16
+        struct              :negotiate_context_info,   label: 'Negotiate Context Info', onlyif: -> { has_negotiate_context? } do
+          uint32            :negotiate_context_offset, label: 'Negotiate Context Offset', initial_value: -> { negotiate_context_list.abs_offset }
+          uint16            :negotiate_context_count,  label: 'Negotiate Context Count', initial_value: -> { negotiate_context_list.size }
+          uint16            :reserved2,                label: 'Reserved', initial_value: 0
+        end
+        file_time           :client_start_time,        label: 'Client Start Time', initial_value: 0, onlyif: -> { !has_negotiate_context? }
+        array               :dialects,                 label: 'Dialects', type: :uint16, initial_length: -> { dialect_count }
+        string              :pad,                      label: 'Padding', length: -> { pad_length(self.dialects) }, onlyif: -> { has_negotiate_context? }
+        array               :negotiate_context_list,   label: 'Negotiate Context List', type: :negotiate_context, onlyif: -> { has_negotiate_context? }, read_until: :eof
+
+        # Adds a dialect to the Dialects array
         #
         # @param [Fixnum] the numeric code for the dialect you wish to add
         # @return [Array<Fixnum>] the array of all currently selected dialects
+        # @raise [ArgumentError] if the dialect is not an Integer
         def add_dialect(dialect)
-          return ArgumentError, 'Must be a number' unless dialect.is_a? Integer
-          self.dialect_count += 1
-          dialects << dialect
+          raise ArgumentError, 'Must be a number' unless dialect.is_a? Integer
+          self.dialects << dialect
         end
 
         # Takes an array of dialects and sets it on the packet. Also updates
@@ -35,12 +44,40 @@ module RubySMB
         # @return [Array<Fixnum>] the current value of the dialects array
         def set_dialects(add_dialects = [])
           self.dialects = []
-          self.dialect_count = 0
           add_dialects.each do |dialect|
             add_dialect(dialect)
           end
           dialects
         end
+
+        # Adds a Negotiate Context to the #negotiate_context_list
+        #
+        # @param [NegotiateContext] the Negotiate Context structure you wish to add
+        # @return [Array<Fixnum>] the array of all currently added Negotiate Contexts
+        # @raise [ArgumentError] if the dialect is not a NegotiateContext structure
+        def add_negotiate_context(nc)
+          raise ArgumentError, 'Must be a NegotiateContext' unless nc.is_a? NegotiateContext
+          previous_element = negotiate_context_list.last || negotiate_context_list
+          pad_length = pad_length(previous_element)
+          self.negotiate_context_list << nc
+          self.negotiate_context_list.last.pad = "\x00" * pad_length
+          self.negotiate_context_list
+        end
+
+
+        private
+
+        # Determines the correct length for the padding, so that the next
+        # field is 8-byte aligned.
+        def pad_length(prev_element)
+          offset = (prev_element.abs_offset + prev_element.to_binary_s.length) % 8
+          (8 - offset) % 8
+        end
+
+        # Return true if the dialect version requires Negotiate Contexts
+        def has_negotiate_context?
+          dialects.any? { |dialect| dialect.to_i == 0x0311 }
+        end
       end
     end
   end

data/lib/ruby_smb/smb2/packet/negotiate_response.rb

@@ -1,3 +1,5 @@
+require 'ruby_smb/smb2/negotiate_context'
+
 module RubySMB
   module SMB2
     module Packet
@@ -8,11 +10,12 @@ module RubySMB
 
         endian              :little
         smb2_header         :smb2_header
-        uint16              :structure_size, label: 'Structure Size', initial_value: 65
+        uint16              :structure_size,            label: 'Structure Size', initial_value: 65
         smb2_security_mode  :security_mode
         uint16              :dialect_revision,          label: 'Dialect Revision'
-        uint16              :negotiate_context_count,   label: 'Negotiate Context Count',      initial_value: 0
-        string              :server_guid,               label: 'Server GUID',                  length: 16
+        uint16              :negotiate_context_count,   label: 'Negotiate Context Count', initial_value: -> { negotiate_context_list.size }, onlyif: -> { has_negotiate_context? }
+        uint16              :reserved1,                 label: 'Reserved', initial_value: 0, onlyif: -> { !has_negotiate_context? }
+        string              :server_guid,               label: 'Server GUID', length: 16
         smb2_capabilities   :capabilities
         uint32              :max_transact_size,         label: 'Max Transaction Size'
         uint32              :max_read_size,             label: 'Max Read Size'
@@ -21,13 +24,56 @@ module RubySMB
         file_time           :server_start_time,         label: 'Server Start Time'
         uint16              :security_buffer_offset,    label: 'Offset to Security Buffer'
         uint16              :security_buffer_length,    label: 'Security Buffer Length', initial_value: -> { security_buffer.length }
-        uint32              :negotiate_context_offset,  label: 'Offset to Negotiate Context'
+        uint32              :negotiate_context_offset,  label: 'Offset to Negotiate Context', onlyif: -> { has_negotiate_context? }
+        uint32              :reserved2,                 label: 'Reserved', initial_value: 0, onlyif: -> { !has_negotiate_context? }
         string              :security_buffer,           label: 'Security Buffer', read_length: :security_buffer_length
+        string              :pad,                       label: 'Padding', length: -> { pad_length(self.security_buffer) }, onlyif: -> { has_negotiate_context? }
+        array               :negotiate_context_list,    label: 'Negotiate Context List', initial_length: -> { negotiate_context_count }, type: :negotiate_context, onlyif: -> { has_negotiate_context? }
 
         def initialize_instance
           super
           smb2_header.flags.reply = 1
         end
+
+        # Find the first Negotiate Context structure that matches the given
+        # context type
+        #
+        # @param [Integer] the Negotiate Context structure you wish to add
+        # @return [NegotiateContext] the Negotiate Context structure or nil if
+        # not found
+        def find_negotiate_context(type)
+          negotiate_context_list.find { |nc| nc.context_type == type }
+        end
+
+        # Adds a Negotiate Context to the #negotiate_context_list
+        #
+        # @param [NegotiateContext] the Negotiate Context structure you wish to add
+        # @return [Array<Fixnum>] the array of all currently added Negotiate Contexts
+        # @raise [ArgumentError] if the dialect is not a NegotiateContext structure
+        def add_negotiate_context(nc)
+          raise ArgumentError, 'Must be a NegotiateContext' unless nc.is_a? NegotiateContext
+          previous_element = negotiate_context_list.last || negotiate_context_list
+          pad_length = pad_length(previous_element)
+          self.negotiate_context_list << nc
+          self.negotiate_context_list.last.pad = "\x00" * pad_length
+          self.negotiate_context_list
+        end
+
+
+        private
+
+        # Determines the correct length for the padding, so that the next
+        # field is 8-byte aligned.
+        def pad_length(prev_element)
+          offset = (prev_element.abs_offset + prev_element.to_binary_s.length) % 8
+          (8 - offset) % 8
+        end
+
+        # Return true if the dialect version requires Negotiate Contexts
+        def has_negotiate_context?
+          dialect_revision == 0x0311
+        end
+
       end
     end
   end

data/lib/ruby_smb/smb2/packet/transform_header.rb

@@ -0,0 +1,84 @@
+module RubySMB
+  module SMB2
+    module Packet
+      # An SMB2 TRANSFORM_HEADER Packet as defined in
+      # [2.2.41 SMB2 TRANSFORM_HEADER](https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/d6ce2327-a4c9-4793-be66-7b5bad2175fa)
+      class TransformHeader < BinData::Record
+        endian :little
+        hide   :reserved0
+
+        endian           :little
+        bit32            :protocol,              label: 'Protocol ID Field',      initial_value: 0xFD534D42
+        string           :signature,             label: 'Signature', length: 16
+        string           :nonce,                 label: 'Nonce',     length: 16
+        uint32           :original_message_size, label: 'Original Message Size'
+        uint16           :reserved0
+        uint16           :flags,                 label: 'Flags / Encryption Algorithm'
+        uint64           :session_id,            label: 'Session ID'
+        array            :encrypted_data,        label: 'Encrypted Data', type: :uint8, read_until: :eof
+
+        def decrypt(key, algorithm: 'AES-128-GCM')
+          auth_data = self.to_binary_s[20...52]
+          encrypted_data = self.encrypted_data.to_ary.pack('C*')
+
+          case algorithm
+          when 'AES-128-CCM'
+            cipher = OpenSSL::CCM.new('AES', key, 16)
+            unencrypted_data = cipher.decrypt(encrypted_data + self.signature, self.nonce[0...11], auth_data)
+            unless unencrypted_data.length > 0
+              raise OpenSSL::Cipher::CipherError  # raised for consistency with GCM mode
+            end
+          when 'AES-128-GCM'
+            cipher = OpenSSL::Cipher.new(algorithm).decrypt
+            cipher.key = key
+            cipher.iv = self.nonce[0...12]
+            cipher.auth_data = auth_data
+            cipher.auth_tag = self.signature
+            unencrypted_data = cipher.update(encrypted_data)
+            cipher.final # raises OpenSSL::Cipher::CipherError on signature failure
+          else
+            raise ArgumentError.new('Invalid algorithm, must be either AES-128-CCM or AES-128-GCM')
+          end
+
+          unencrypted_data[0...self.original_message_size]
+        rescue Exception => e
+          raise RubySMB::Error::EncryptionError, "Error while decrypting with '#{algorithm}' (#{e.class}: #{e})"
+        end
+
+        def encrypt(unencrypted_data, key, algorithm: 'AES-128-GCM')
+          if unencrypted_data.is_a? BinData::Record
+            unencrypted_data = unencrypted_data.to_binary_s
+          end
+
+          self.original_message_size.assign(unencrypted_data.length)
+
+          case algorithm
+          when 'AES-128-CCM'
+            cipher = OpenSSL::CCM.new('AES', key, 16)
+            random_iv = OpenSSL::Random.random_bytes(11)
+            self.nonce.assign(random_iv)
+            result = cipher.encrypt(unencrypted_data, random_iv, self.to_binary_s[20...52])
+            encrypted_data = result[0...-16]
+            auth_tag = result[-16..-1]
+          when 'AES-128-GCM'
+            cipher = OpenSSL::Cipher.new(algorithm).encrypt
+            cipher.iv_len = 12
+            cipher.key = key
+            self.nonce.assign(cipher.random_iv)
+            cipher.auth_data = self.to_binary_s[20...52]
+            encrypted_data = cipher.update(unencrypted_data) + cipher.final
+            auth_tag = cipher.auth_tag
+          else
+            raise ArgumentError.new('Invalid algorithm, must be either AES-128-CCM or AES-128-GCM')
+          end
+
+          self.encrypted_data.assign(encrypted_data.bytes)
+          self.signature.assign(auth_tag)
+          nil
+        rescue Exception => e
+          raise RubySMB::Error::EncryptionError, "Error while encrypting with '#{algorithm}' (#{e.class}: #{e})"
+        end
+      end
+    end
+  end
+end

data/lib/ruby_smb/smb2/packet/tree_connect_request.rb

@@ -1,22 +1,108 @@
 module RubySMB
   module SMB2
     module Packet
+
+
+      # An SMB2 RemotedIdentityTreeConnectContext Packet as defined in
+      # [2.2.9.2.1 SMB2_REMOTED_IDENTITY_TREE_CONNECT Context](https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/ee7ff411-93e0-484f-9f73-31916fee4cb8)
+      # TODO: implement helper methods to add each Remote Identity element
+      class RemotedIdentityTreeConnectContext < BinData::Record
+        endian :little
+        uint16 :ticket_type,       label: 'Ticket Type', initial_value: 0x0001
+        uint16 :ticket_size,       label: 'Ticket Size', initial_value: -> { num_bytes }
+        uint16 :user,              label: 'User'
+        uint16 :user_name,         label: 'User Name'
+        uint16 :domain,            label: 'Domain'
+        uint16 :groups,            label: 'Groups'
+        uint16 :restricted_groups, label: 'Restricted Groups'
+        uint16 :privileges,        label: 'Privileges'
+        uint16 :primary_group,     label: 'Primary Group'
+        uint16 :owner,             label: 'Owner'
+        uint16 :default_dacl,      label: 'Default DACL'
+        uint16 :device_groups,     label: 'Device Groups'
+        uint16 :user_claims,       label: 'User Claims'
+        uint16 :device_claims,     label: 'Device Claims'
+        string :ticket_info,       label: 'Ticket Info', read_length: -> { ticket_size - ticket_info.rel_offset}
+      end
+
+      # An SMB2 TreeConnectContext Packet as defined in
+      # [2.2.9.2 SMB2 TREE_CONNECT_CONTEXT Request Values](https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/06eaaabc-caca-4776-9daf-82439e90dacd)
+      class TreeConnectContext < BinData::Record
+
+        # Context Types
+
+        # This value is reserved.
+        SMB2_RESERVED_TREE_CONNECT_CONTEXT_ID = 0x0000
+        # The Data field contains remoted identity tree connect context data as
+        # specified in section [2.2.9.2.1](https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/ee7ff411-93e0-484f-9f73-31916fee4cb8)
+        SMB2_REMOTED_IDENTITY_TREE_CONNECT_CONTEXT_ID = 0x0001
+
+        endian :little
+        uint16 :context_type, label: 'Context Type'
+        uint16 :data_length,  label: 'Data Length', initial_value: -> { data.to_binary_s.size }
+        uint32 :reserved,     label: 'Reserved'
+        choice :data,         label: 'Data', selection: -> { context_type } do
+          remoted_identity_tree_connect_context SMB2_REMOTED_IDENTITY_TREE_CONNECT_CONTEXT_ID, label: 'Remoted Identity Tree Connect Context'
+        end
+
+      end
+
+      # An SMB2 TreeConnectRequestExtension Packet as defined in
+      # [2.2.9.1 SMB2 TREE_CONNECT Request Extension](https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/9ca7328b-b6ca-41a7-9773-0fa237261b76)
+      class TreeConnectRequestExtension < BinData::Record
+        endian :little
+        uint32   :tree_connect_context_offset, label: 'Tree Connect Context Offset', initial_value: -> { tree_connect_contexts.rel_offset }
+        uint16   :tree_connect_context_count,  label: 'Tree Connect Context Count', initial_value: -> { tree_connect_contexts.size }
+        string   :reserved,                    label: 'Reserved', length: 10
+        string16 :path,                        label: 'Path Buffer'
+        array    :tree_connect_contexts,       label: 'Tree Connect Contexts', type: :tree_connect_context, initial_length: -> { tree_connect_context_count }
+      end
+
       # An SMB2 TreeConnectRequest Packet as defined in
       # [2.2.9 SMB2 TREE_CONNECT Request](https://msdn.microsoft.com/en-us/library/cc246567.aspx)
       class TreeConnectRequest < RubySMB::GenericPacket
         COMMAND = RubySMB::SMB2::Commands::TREE_CONNECT
 
+        # Flags (SMB 3.1.1 only)
+
+        # The client has previously connected to the specified cluster share
+        # using the SMB dialect of the connection on which the request is received.
+        SMB2_TREE_CONNECT_FLAG_CLUSTER_RECONNECT = 0x0001
+        # The client can handle synchronous share redirects via a Share Redirect
+        # error context response as specified in section [2.2.2.2.2](https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/f3073a8b-9f0f-47c0-91e5-ec3be9a49f37).
+        SMB2_TREE_CONNECT_FLAG_REDIRECT_TO_OWNER = 0x0002
+        # A tree connect request extension, as specified in section
+        # [2.2.9.1](https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-smb2/9ca7328b-b6ca-41a7-9773-0fa237261b76),
+        # is present, starting at the Buffer field of this tree connect request.
+        SMB2_TREE_CONNECT_FLAG_EXTENSION_PRESENT = 0x0003
+
         endian       :little
         smb2_header  :smb2_header
         uint16       :structure_size, label: 'Structure Size', initial_value: 9
+        # The flags field is only used by SMB 3.1.1, it must be 0 for other versions
         uint16       :flags,          label: 'Flags',          initial_value: 0x00
-        uint16       :path_offset,    label: 'Path Offset',    initial_value: 0x48
-        uint16       :path_length,    label: 'Path Length',    initial_value: -> { path.length }
-        string       :path,           label: 'Path Buffer'
-
-        def encode_path(path)
-          self.path = path.encode('utf-16le')
+        # if SMB2_TREE_CONNECT_FLAG_EXTENSION_PRESENT flag is set, #path_offset
+        # will have to be updated with the correct offset of the path name,
+        # which is located in the TreeConnect Context.
+        uint16       :path_offset,    label: 'Path Offset',    initial_value: -> do
+          if flags == SMB2_TREE_CONNECT_FLAG_EXTENSION_PRESENT
+            tree_connect_request_extension.path.abs_offset
+          else
+            path.abs_offset
+          end
+        end
+        # if SMB2_TREE_CONNECT_FLAG_EXTENSION_PRESENT flag is set, #path_length
+        # will have to be updated with the correct full share path name,
+        # which is located in the TreeConnect Context.
+        uint16       :path_length,    label: 'Path Length',    initial_value: -> do
+          if flags == SMB2_TREE_CONNECT_FLAG_EXTENSION_PRESENT
+            tree_connect_request_extension.path.to_binary_s.length
+          else
+            path.to_binary_s.length
+          end
         end
+        string16     :path,           label: 'Path Buffer',    onlyif: -> { flags != SMB2_TREE_CONNECT_FLAG_EXTENSION_PRESENT }
+        tree_connect_request_extension :tree_connect_request_extension, label: 'Tree Connect Request Extension', onlyif: -> { flags == SMB2_TREE_CONNECT_FLAG_EXTENSION_PRESENT }
       end
     end
   end