# -*- coding: binary -*-

module Msf
  module Exploit::Remote::SMB::Server
    module Share
      module InformationLevel
        module Query

          # Handles a TRANS2_QUERY_FILE_INFORMATION transaction request with SMB_QUERY_FILE_BASIC_INFO
          # Information Level.
          #
          # @param c [Socket] The client sending the request.
          # @param fid [Integer] The file identifier which the client is requesting info from.
          # @return [Integer] The number of bytes returned to the client as response.
          def smb_cmd_trans_query_file_info_basic(c, fid)
            smb = @state[c]

            if fid == smb[:file_id].to_i
              attrib = CONST::SMB_EXT_FILE_ATTR_NORMAL
            elsif fid.nil? || fid == 0 || fid == smb[:dir_id].to_i # empty fid
              attrib = CONST::SMB_EXT_FILE_ATTR_DIRECTORY
            else
              return smb_error(CONST::SMB_COM_TRANSACTION2, c, CONST::SMB_STATUS_OBJECT_NAME_NOT_FOUND, true)
            end

            send_info_basic_res(c, { file_attributes: attrib })
          end

          # Handles a TRANS2_QUERY_FILE_INFORMATION transaction request with SMB_QUERY_FILE_STANDARD_INFO
          # Information Level.
          #
          # @param c [Socket] The client sending the request.
          # @param fid [Integer] The file identifier which the client is requesting info from.
          # @return [Integer] The number of bytes returned to the client as response.
          def smb_cmd_trans_query_file_info_standard(c, fid)
            contents = get_file_contents(client: c)

            send_info_standard_res(c, {
              allocation_size: 1048576,
              number_links: 1,
              delete_pending: 0,
              directory: 0,
              end_of_file: contents.length
            })
          end

          # Handles a TRANS2_QUERY_PATH_INFORMATION transaction request with SMB_QUERY_FILE_BASIC_INFO
          # Information Level.
          #
          # @param c [Socket] The client sending the request.
          # @param path [String] The path which the client is requesting info from.
          # @return [Integer] The number of bytes returned to the client as response.
          def smb_cmd_trans_query_path_info_basic(c, path)
            if path && path.ends_with?(file_name.downcase)
              attrib = CONST::SMB_EXT_FILE_ATTR_NORMAL
            elsif path && folder_name && path.ends_with?(folder_name.downcase)
              attrib = CONST::SMB_EXT_FILE_ATTR_DIRECTORY
            elsif path.nil? || path.empty? || path == "\x00" || path == "\\" # empty path
              attrib = CONST::SMB_EXT_FILE_ATTR_DIRECTORY
            else
              return smb_error(CONST::SMB_COM_TRANSACTION2, c, CONST::SMB_STATUS_OBJECT_NAME_NOT_FOUND, true)
            end

            send_info_basic_res(c, { file_attributes: attrib })
          end

          # Handles a TRANS2_QUERY_PATH_INFORMATION transaction request with SMB_QUERY_FILE_STANDARD_INFO
          # Information Level.
          #
          # @param c [Socket] The client sending the request.
          # @param path [String] The path which the client is requesting info from.
          # @return [Integer] The number of bytes returned to the client as response.
          def smb_cmd_trans_query_path_info_standard(c, path)
            contents = get_file_contents(client: c)

            if path && path.include?(file_name.downcase)
              attrib = 0 # File attributes => file
            elsif path && folder_name && path.ends_with?(folder_name.downcase)
              attrib = 1 # File attributes => directory
            elsif path.nil? || path.empty? || path == "\x00" || path == "\\" # empty path
              attrib = 1 # File attributes => directory
            else
              return smb_error(CONST::SMB_COM_TRANSACTION2, c, CONST::SMB_STATUS_OBJECT_NAME_NOT_FOUND, true)
            end

            send_info_standard_res(c, {
              allocation_size: 1048576,
              number_links: 1,
              delete_pending: 0,
              directory: attrib,
              end_of_file: contents.length
            })
          end

          # Handles a TRANS2_QUERY_PATH_INFORMATION transaction request with SMB_QUERY_FILE_NETWORK_INFO
          # Information Level.
          #
          # @param c [Socket] The client sending the request.
          # @param path [String] The path which the client is requesting info from.
          # @return [Integer] The number of bytes returned to the client as response.
          def smb_cmd_trans_query_path_info_network(c, path)
            contents = get_file_contents(client: c)

            if path && path.include?(file_name.downcase)
              attrib = CONST::SMB_EXT_FILE_ATTR_NORMAL
            elsif path && folder_name && path.ends_with?(folder_name.downcase)
              attrib = CONST::SMB_EXT_FILE_ATTR_DIRECTORY
            elsif path.nil? || path.empty? || path == "\x00" || path == "\\" # empty path
              attrib = CONST::SMB_EXT_FILE_ATTR_DIRECTORY
            else
              return smb_error(CONST::SMB_COM_TRANSACTION2, c, CONST::SMB_STATUS_OBJECT_NAME_NOT_FOUND, true)
            end

            send_info_network_res(c, {
              allocation_size: 1048576,
              end_of_file: contents.length,
              file_attributes: attrib
            })
          end

          # Builds and sends an TRANS2_QUERY_PATH_INFORMATION response with SMB_QUERY_FILE_BASIC_INFO
          # information level.
          #
          # @param c [Socket] The client to answer.
          # @param opts [Hash{Symbol => <Integer, String>}] Response custom values.
          # @option opts [Integer] :file_attributes The extended file attributes of the file.
          # @return [Integer] The number of bytes returned to the client as response.
          def send_info_basic_res(c, opts = {})
            file_attributes = opts[:file_attributes] || 0

            trans2_params = CONST::SMB_TRANS2_QUERY_PATH_INFORMATION_RES_PARAMETERS.make_struct
            trans2_params.v['EaErrorOffset'] = 0

            query_path_info = CONST::SMB_QUERY_FILE_BASIC_INFO_HDR.make_struct
            query_path_info.v['loCreationTime'] = lo
            query_path_info.v['hiCreationTime'] = hi
            query_path_info.v['loLastAccessTime'] = lo
            query_path_info.v['hiLastAccessTime'] = hi
            query_path_info.v['loLastWriteTime'] = lo
            query_path_info.v['hiLastWriteTime'] = hi
            query_path_info.v['loLastChangeTime'] = lo
            query_path_info.v['hiLastChangeTime'] = hi
            query_path_info.v['ExtFileAttributes'] = file_attributes

            send_trans2_res(c, trans2_params, query_path_info)
          end

          # Builds and sends an TRANS2_QUERY_PATH_INFORMATION response with SMB_QUERY_FILE_STANDARD_INFO
          # information level.
          #
          # @param c [Socket] The client to answer.
          # @param opts [Hash{Symbol => <Integer, String>}] Response custom values.
          # @option opts [Integer] :allocation_size The number of bytes that are allocated to the file.
          # @option opts [Integer] :number_links The number of hard links to the file.
          # @option opts [Integer] :delete_pending Indicates whether there is a delete action pending for the file.
          # @option opts [Integer] :directory Indicates whether the file is a directory.
          # @option opts [Integer] :end_of_file The offset from the start to the end of the file.
          # @return [Integer] The number of bytes returned to the client as response.
          def send_info_standard_res(c, opts = {})
            allocation_size = opts[:allocation_size] || 0
            number_links = opts[:number_links] || 0
            delete_pending = opts[:delete_pending] || 0
            directory = opts[:directory] || 0
            end_of_file = opts[:end_of_file] || 0

            trans2_params = CONST::SMB_TRANS2_QUERY_PATH_INFORMATION_RES_PARAMETERS.make_struct
            trans2_params.v['EaErrorOffset'] = 0

            query_path_info = CONST::SMB_QUERY_FILE_STANDARD_INFO_HDR.make_struct
            query_path_info.v['AllocationSize'] = allocation_size
            query_path_info.v['EndOfFile'] = end_of_file
            query_path_info.v['NumberOfLinks'] = number_links
            query_path_info.v['DeletePending'] = delete_pending
            query_path_info.v['Directory'] = directory

            send_trans2_res(c, trans2_params, query_path_info)
          end

          # Builds and sends an TRANS2_QUERY_PATH_INFORMATION response with SMB_QUERY_FILE_NETWORK_INFO
          # information level.
          #
          # @param c [Socket] The client to answer.
          # @param opts [Hash{Symbol => <Integer, String>}] Response custom values.
          # @option opts [Integer] :allocation_size The number of bytes that are allocated to the file.
          # @option opts [Integer] :end_of_file The offset from the start to the end of the file.
          # @option opts [Integer] :file_attributes The file attributes.
          # @return [Integer] The number of bytes returned to the client as response.
          def send_info_network_res(c, opts= {})
            allocation_size = opts[:allocation_size] || 0
            end_of_file = opts[:end_of_file] || 0
            file_attributes = opts[:file_attributes] || 0

            pkt = CONST::SMB_TRANS_RES_PKT.make_struct
            smb_set_defaults(c, pkt)

            trans2_params = CONST::SMB_TRANS2_QUERY_PATH_INFORMATION_RES_PARAMETERS.make_struct
            trans2_params.v['EaErrorOffset'] = 0

            query_path_info = CONST::SMB_QUERY_FILE_NETWORK_INFO_HDR.make_struct
            query_path_info.v['loCreationTime'] = lo
            query_path_info.v['hiCreationTime'] = hi
            query_path_info.v['loLastAccessTime'] = lo
            query_path_info.v['hiLastAccessTime'] = hi
            query_path_info.v['loLastWriteTime'] = lo
            query_path_info.v['hiLastWriteTime'] = hi
            query_path_info.v['loLastChangeTime'] = lo
            query_path_info.v['hiLastChangeTime'] = hi
            query_path_info.v['AllocationSize'] = allocation_size
            query_path_info.v['EndOfFile'] = end_of_file
            query_path_info.v['ExtFileAttributes'] = file_attributes

            send_trans2_res(c, trans2_params, query_path_info)
          end
        end
      end
    end
  end
end
