import enum
import io
import asyncio
import datetime
import traceback
import typing
from dataclasses import dataclass, field
from pathlib import Path

from aiosmb import logger
from aiosmb.commons.exceptions import *
from aiosmb.protocol.smb.command_codes import SMBCommand
from aiosmb.wintypes.ntstatus import NTStatus
from aiosmb.protocol.smb.message import SMBMessage
from aiosmb.protocol.smb.commands import *
from aiosmb.protocol.smb2.message import SMB2Message
from aiosmb.protocol.smb2.commands.tree_connect import *
from aiosmb.protocol.smb2.commands import *
from aiosmb.protocol.smb2.headers import *
from aiosmb.protocol.smb2.command_codes import *
from aiosmb.protocol.common import *
from aiosmb.wintypes.dtyp.constrcuted_security.guid import *
from aiosmb.protocol.smb2.commands.sessionsetup import *
from aiosmb.protocol.smb2.commands.create import *
from aiosmb.wintypes.fscc.FileAttributes import FileAttributes
from aiosmb.wintypes.fscc.structures.fileinfoclass import FileInfoClass


def convert_to_ft(linux_filetime:int):
	return int((linux_filetime + 11644473600) * 10000000)

class SMBConnectionStatus(enum.Enum):
	NEGOTIATING = 'NEGOTIATING'
	SESSIONSETUP = 'SESSIONSETUP'
	RUNNING = 'RUNNING'
	CLOSED = 'CLOSED'

@dataclass
class FileEntry:
	file_id:int = None
	path:Path = None
	handle:io.FileIO = None

@dataclass
class TreeEntry:
	file_entries:dict[int, FileEntry] = field(default_factory=dict)
	hostname:str = None
	path:Path = None
	tree_id:int = None

@dataclass
class DirEnumState:
	entries:list[str] = field(default_factory=list)
	index:int = 0
	pattern:str = '*'
	info_class:FileInfoClass = None



class SMBServerConnection:
	def __init__(self, settings, connection, connection_id, shutdown_evt = asyncio.Event()):
		self.settings = settings
		self.gssapi = self.settings.create_new_gssapi(connection)
		self.connection_id = connection_id
		
		#######DONT CHANGE THIS
		#use this for smb2 > self.supported_dialects = [NegotiateDialects.WILDCARD, NegotiateDialects.SMB202, NegotiateDialects.SMB210]
		#self.supported_dialects = [NegotiateDialects.SMB202, NegotiateDialects.SMB210]
		self.supported_dialects = self.settings.preferred_dialects #[NegotiateDialects.WILDCARD, NegotiateDialects.SMB311]
		#######
		
		self.connection = connection
		self.incoming_task = None
		
		self.selected_dialect = None
		self.signing_required = self.settings.RequireSigning
		self.encryption_required = False
		self.last_treeid = 10
		
		self.status = SMBConnectionStatus.NEGOTIATING
		
		self.OutstandingResponsesEvent = {}
		self.OutstandingRequests = {}
		self.OutstandingResponses = {}

		self.pending_table = {}
		
		self.SequenceWindow = 0
		self.MaxTransactSize = self.settings.MaxTransactSize
		self.MaxReadSize = self.settings.MaxReadSize
		self.MaxWriteSize = self.settings.MaxWriteSize
		self.ServerGuid = self.settings.ServerGuid

		self.SupportsFileLeasing = False
		self.SupportsMultiCredit = False
		
		self.SupportsDirectoryLeasing = False
		self.SupportsMultiChannel = False
		self.SupportsPersistentHandles = False
		self.SupportsEncryption = False
		self.ClientCapabilities = 0
		self.ServerCapabilities = NegotiateCapabilities.LARGE_MTU #NegotiateCapabilities.DFS | NegotiateCapabilities.LARGE_MTU
		self.ClientSecurityMode = 0
		self.ServerSecurityMode = NegotiateSecurityMode.SMB2_NEGOTIATE_SIGNING_ENABLED
		if self.signing_required is True:
			self.ServerSecurityMode |= NegotiateSecurityMode.SMB2_NEGOTIATE_SIGNING_REQUIRED
		
		self.SessionId = 0
		self.SessionKey = None

		self.last_treeid = 1

		self.__treeid_to_entry = {}

	async def run(self):
		self.incoming_task = asyncio.create_task(self.__handle_smb_in())
		return self.incoming_task
	
	async def stop(self):
		try:
			await self.print('[INF] Stopping connection')
			if self.incoming_task is not None:
				self.incoming_task.cancel()
			if self.connection is not None:
				await self.connection.close()
		except Exception as e:
			traceback.print_exc()
			await self.print('[ERR] %s' % e)

	async def print(self, msg):
		if self.settings.log_callback is not None:
			await self.settings.log_callback(f'[SMBSERVERCONN][{self.connection_id}] {msg}')

	def get_new_tree_id(self):
		self.last_treeid += 1
		return self.last_treeid

	async def __handle_smb_in(self):
		"""
		Waits from SMB message bytes from the transport in_queue, and fills the connection table.
		This function started automatically when calling connect.
		Pls don't touch it.
		"""
		try:
			async for msg_data in self.connection.read():
				if msg_data is None:
					break
				
				if msg_data[0] < 252:
					raise Exception('Unknown SMB packet type %s' % msg_data[0])

				if msg_data[0] == 0xFD:
					raise Exception('SMB2_COMPRESSED_DATA_HEADER is not implemented yet')

				if msg_data[0] == 0xFC:
					raise Exception('SMB2_COMPRESSED_TRANSFORM_HEADER is not implemented yet')
					
				if msg_data[0] == 0xFE:
					#version2
					try:
						msg = SMB2Message.from_bytes(msg_data)
						await self.print('[PACKET] %s' % msg)
					except Exception as e:
						raise Exception('Failed to parse SMB data. Probably not implemented feature.')

					if self.status == SMBConnectionStatus.NEGOTIATING:
						await self.negotiate(msg)
					elif self.status == SMBConnectionStatus.SESSIONSETUP:
						await self.session_setup(msg)
					
					elif self.status == SMBConnectionStatus.RUNNING:
						if msg.header.Command == SMB2Command.TREE_CONNECT:
							await self.tree_connect(msg)
						elif msg.header.Command == SMB2Command.CREATE:
							await self.create(msg)
						elif msg.header.Command == SMB2Command.IOCTL:
							await self.ioctl(msg)
						elif msg.header.Command == SMB2Command.ECHO:
							await self.echo(msg)
						elif msg.header.Command == SMB2Command.QUERY_DIRECTORY:
							await self.query_directory(msg)
						elif msg.header.Command == SMB2Command.CLOSE:
							await self.close(msg)
						else:
							print('SMB2 command not implemented! %s' % msg.header.Command)
							continue
				if msg_data[0] == 0xFF:
					#version1
					msg = SMBMessage.from_bytes(msg_data)
					await self.print('[PACKET] %s' % msg)
					if self.status == SMBConnectionStatus.NEGOTIATING:
						if msg.header.Command == SMBCommand.SMB_COM_NEGOTIATE:
							await self.negotiate(msg)
					else:
						print('SMBv1 message recieved! This is unexpected! %s' % msg.header.Command)
						continue

				logger.log(1, '__handle_smb_in got new message with Id %s' % msg.header.MessageId)

				if msg.header.MessageId in self.pending_table:
					await self.pending_table[msg.header.MessageId].stop()
					del self.pending_table[msg.header.MessageId]
				
		except asyncio.CancelledError:
			return
		except:
			traceback.print_exc()
		finally:
			await self.stop()
	
	async def sendSMB(self, msg, compression_cb = None):
		"""
		Sends an SMB message to teh remote endpoint.
		msg: SMB2Message or SMBMessage
		Returns: MessageId integer
		"""
		try:
			if self.status == SMBConnectionStatus.NEGOTIATING:
				if isinstance(msg, SMBMessage):
					message_id = 0
					self.SequenceWindow += 1
				else:
					msg.header.CreditCharge = 1
					if msg.header.CreditReq is None:
						msg.header.CreditReq = 1
					msg.header.MessageId = self.SequenceWindow
					message_id = self.SequenceWindow
					self.SequenceWindow += 1

				self.OutstandingResponsesEvent[message_id] = asyncio.Event()
				await self.connection.write(msg.to_bytes())
				return message_id, msg, None
					
			if msg.header.Command is not SMB2Command.CANCEL:
				msg.header.MessageId = self.SequenceWindow
				self.SequenceWindow += 1
			
			msg.header.SessionId = self.SessionId
			
			if not msg.header.CreditCharge:
				msg.header.CreditCharge = 1

			
			
			if self.status != SMBConnectionStatus.SESSIONSETUP and msg.header.CreditReq is None:
				msg.header.CreditReq = 127
			
			message_id = msg.header.MessageId
			#print(msg)

			#creating an event for outstanding response
			self.OutstandingResponsesEvent[message_id] = asyncio.Event()
			await self.connection.write(msg.to_bytes())
			
			return message_id, msg, None
		except Exception as e:
			traceback.print_exc()
			return None, None, e
		
	async def negotiate(self, req):
		try:
			if (isinstance(req, SMBMessage) and 'SMB 2.???' in req.command.Dialects) or isinstance(req, SMB2Message):
					#if NegotiateDialects.WILDCARD in self.supported_dialects:
					reply = SMB2Message()
					reply.command = NEGOTIATE_REPLY()
					reply.command.SecurityMode = self.ServerSecurityMode
					reply.command.DialectRevision = NegotiateDialects.SMB202
					reply.command.NegotiateContextCount = 0
					reply.command.ServerGuid = self.ServerGuid
					reply.command.Capabilities = self.ServerCapabilities
					reply.command.MaxTransactSize = self.MaxTransactSize
					reply.command.MaxReadSize = self.MaxReadSize
					reply.command.MaxWriteSize = self.MaxWriteSize
					reply.command.SystemTime = datetime.datetime.utcnow()
					reply.command.ServerStartTime = datetime.datetime.utcnow()
					reply.command.SecurityBuffer = self.gssapi.get_mechtypes_list()
					reply.command.NegotiateContextOffset = 0
					reply.command.NegotiateContextList = []
					
					reply.header = SMB2Header_SYNC()
					reply.header.Command  = SMB2Command.NEGOTIATE
					reply.header.Flags = SMB2HeaderFlag.SMB2_FLAGS_SERVER_TO_REDIR
					reply.header.CreditCharge = 1
					await self.sendSMB(reply)

					self.selected_dialect = NegotiateDialects.SMB202

					
					self.SupportsMultiChannel = NegotiateCapabilities.MULTI_CHANNEL in self.ServerCapabilities
					self.SupportsFileLeasing = NegotiateCapabilities.LEASING in self.ServerCapabilities
					self.SupportsMultiCredit = NegotiateCapabilities.LARGE_MTU in self.ServerCapabilities
					
					self.status = SMBConnectionStatus.SESSIONSETUP
			return
		except Exception as e:
			traceback.print_exc()

	async def session_setup(self, msg):
		err = None
		try:
			maxiter = 5
			if self.SessionId == 0:
				self.SessionId += 1
			while maxiter > 0:
				reply = SMB2Message()
				reply.command = SESSION_SETUP_REPLY()
				reply.command.SessionFlags = 0
				try:
					await self.print('[INF] Calling authenticate_server')
					#reply.command.Buffer, to_continue, err  = await self.gssapi.authenticate(msg.command.Buffer)
					reply.command.Buffer, to_continue, err  = await self.gssapi.authenticate_server(msg.command.Buffer)
					await self.print('[INF] Authenticate results: %s, %s, %s' % (reply.command.Buffer, to_continue, err))
					
					if err is not None:
						raise err
					
					if to_continue is False and reply.command.Buffer is None:
						reply = SMB2Message()
						reply.command = SESSION_SETUP_REPLY()
						reply.command.SessionFlags = SessionFlags.SMB2_SESSION_FLAG_IS_GUEST

						reply.command.Buffer, err  = await self.gssapi.authenticate_server_finished()

						if err is not None:
							raise err

						reply.header = SMB2Header_SYNC()
						reply.header.Command  = SMB2Command.SESSION_SETUP
						reply.header.Flags = SMB2HeaderFlag.SMB2_FLAGS_SERVER_TO_REDIR
						#reply.header.Flags |= SMB2HeaderFlag.SMB2_FLAGS_SIGNED
						reply.header.SessionId = self.SessionId
						
						reply.header.CreditCharge = 1
						reply.header.CreditReq = 1
						reply.header.Status = NTStatus.SUCCESS
						#reply.header.Flags |= SMB2HeaderFlag.SMB2_FLAGS_SIGNED
						
						await self.sendSMB(reply)
						#self.SessionKey = self.gssapi.get_session_key()[:16]
						self.status = SMBConnectionStatus.RUNNING
						return False, None
					
				except Exception as e:
					#TODO: Clear this up, kerberos lib needs it's own exceptions!
					if str(e).find('Preauth') != -1:
						raise SMBKerberosPreauthFailed()
					else:
						raise e
						#raise SMBKerberosPreauthFailed()
				
				if to_continue is True:
					reply.header = SMB2Header_SYNC()
					reply.header.Command  = SMB2Command.SESSION_SETUP
					reply.header.Flags = SMB2HeaderFlag.SMB2_FLAGS_SERVER_TO_REDIR
					reply.header.SessionId = self.SessionId
					
					reply.header.CreditCharge = 1
					reply.header.CreditReq = 1
					reply.header.Status = NTStatus.MORE_PROCESSING_REQUIRED
					await self.sendSMB(reply)
				
				if to_continue is False:
					reply.header = SMB2Header_SYNC()
					reply.header.Command  = SMB2Command.SESSION_SETUP
					reply.header.Flags = SMB2HeaderFlag.SMB2_FLAGS_SERVER_TO_REDIR
					reply.header.SessionId = self.SessionId
					
					reply.header.CreditCharge = 1
					reply.header.CreditReq = 1
					reply.header.Status = NTStatus.SUCCESS
					reply.header.Flags |= SMB2HeaderFlag.SMB2_FLAGS_SIGNED
					reply.header.CreditReq = 127

					await self.sendSMB(reply)

					#self.SessionKey = self.gssapi.get_session_key()[:16]
					self.status = SMBConnectionStatus.RUNNING

				return True, None
		except Exception as e:
			await self.print('[ERR] %s' % e)

			reply.header = SMB2Header_SYNC()
			reply.header.Command  = SMB2Command.SESSION_SETUP
			reply.header.Flags = SMB2HeaderFlag.SMB2_FLAGS_SERVER_TO_REDIR
			reply.header.Flags |= SMB2HeaderFlag.SMB2_FLAGS_SIGNED
			reply.header.SessionId = self.SessionId
			reply.header.CreditCharge = 1
			reply.header.CreditReq = 127
			reply.header.Status = NTStatus.ACCESS_DENIED
			await self.sendSMB(reply)
			return False, None # none because relay doesnt need to know the error


	async def tree_connect(self, msg:SMB2Message):
		try:
			command = typing.cast(TREE_CONNECT_REQ, msg.command)
			# command.Path like \\hostname\sharename
			reply = SMB2Message()
			reply.header = SMB2Header_SYNC()
			reply.header.Command = SMB2Command.TREE_CONNECT
			reply.header.Flags = SMB2HeaderFlag.SMB2_FLAGS_SERVER_TO_REDIR
			reply.header.SessionId = self.SessionId
			reply.header.CreditCharge = 1
			reply.header.CreditReq = msg.header.CreditReq
			reply.header.Status = NTStatus.SUCCESS
			reply.command = TREE_CONNECT_REPLY()
			# Parse sharename
			try:
				_,_,_, share_name, *rest = command.Path.split('\\')
				print('TREE_CONNECT: %s' % share_name)
				print(rest)
			except Exception:
				share_name = command.Path

			if share_name.upper().endswith('IPC$'):
				reply.command.ShareType = ShareType.PIPE
			else:
				reply.command.ShareType = ShareType.DISK

			reply.command.ShareFlags = ShareFlags.SMB2_SHAREFLAG_MANUAL_CACHING
			reply.command.Capabilities = TreeCapabilities.SMB2_SHARE_CAP_NONE
			reply.command.MaximalAccess = FileAccessMask.GENERIC_ALL

			# Map sharename to on-disk path from settings.shares
			mapped = None
			for k, v in self.settings.shares.items():
				if k.lower() == share_name.lower():
					mapped = Path(v)
					break
			if mapped is None:
				reply.header.Status = NTStatus.BAD_NETWORK_NAME
				#await self.sendSMB(reply)
				await self.sendSMB2(msg, SMB2Command.TREE_CONNECT, reply.command, NTStatus.BAD_NETWORK_NAME)
				return
			tree_id = self.get_new_tree_id()
			reply.header.TreeId = tree_id
			self.__treeid_to_entry[tree_id] = TreeEntry(hostname=None, path=mapped, tree_id=tree_id)

			await self.sendSMB(reply)
		except Exception as e:
			traceback.print_exc()
	
	async def ioctl(self, msg:SMB2Message):
		try:
			command = typing.cast(IOCTL_REQ, msg.command)
			if command.CtlCode == CtlCode.FSCTL_DFS_GET_REFERRALS:
				reply = SMB2Message()
				reply.header = SMB2Header_SYNC()
				reply.header.Command = SMB2Command.IOCTL
				reply.header.Flags = SMB2HeaderFlag.SMB2_FLAGS_SERVER_TO_REDIR
				reply.header.SessionId = self.SessionId
				reply.header.CreditCharge = 1
				reply.header.CreditReq = msg.header.CreditReq
				reply.header.Status = NTStatus.FS_DRIVER_REQUIRED
				reply.header.messageId = msg.header.MessageId
				# we're currently not handling the DFS_GET_REFERRALS
				reply.command = IOCTL_REPLY()
				reply.command.CtlCode = command.CtlCode
				reply.command.FileId = command.FileId
				reply.command.InputOffset = command.InputOffset
				reply.command.InputCount = command.InputCount
				reply.command.OutputOffset = 0
				reply.command.OutputCount = 0
				reply.command.Flags = command.Flags
				reply.command.Reserved2 = command.Reserved2
				reply.command.Buffer = b''
				#await self.sendSMB(reply)
				await self.sendSMB2(msg, SMB2Command.IOCTL, reply.command, NTStatus.FS_DRIVER_REQUIRED)

		except Exception as e:
			traceback.print_exc()
		
	async def create(self, msg:SMB2Message):
		try:
			command = typing.cast(CREATE_REQ, msg.command)
			handle = None
			if command.Name is None:
				# this means we want to open the share which was specified in tree connect
				share = self.__treeid_to_entry[msg.header.TreeId].path
				print('OPENING SHARE %s' % share.name)
				# this must be a directory
				if not (share.exists() and share.is_dir()):
					raise Exception('Share %s does not exist' % share.name)
				
				reply = SMB2Message()
				reply.header = SMB2Header_SYNC()
				reply.header.Command = SMB2Command.CREATE
				reply.header.Flags = SMB2HeaderFlag.SMB2_FLAGS_SERVER_TO_REDIR
				reply.header.SessionId = self.SessionId
				reply.header.CreditCharge = 1
				reply.header.CreditReq = msg.header.CreditReq
				reply.header.Status = NTStatus.SUCCESS
				reply.command = CREATE_REPLY()
				reply.command.OplockLevel = OplockLevel.SMB2_OPLOCK_LEVEL_NONE
				reply.command.Flags = CreateOptions.FILE_WRITE_THROUGH
				reply.command.CreateAction = CreateAction.FILE_CREATED
				reply.command.CreationTime = convert_to_ft(share.stat().st_ctime)
				reply.command.LastAccessTime = convert_to_ft(share.stat().st_atime)
				reply.command.LastWriteTime = convert_to_ft(share.stat().st_mtime)
				reply.command.ChangeTime = convert_to_ft(share.stat().st_mtime)
				reply.command.AllocationSize = 0
				reply.command.EndofFile = 0
				reply.command.FileAttributes = FileAttributes.FILE_ATTRIBUTE_DIRECTORY
				reply.command.FileId = 0
				reply.command.CreateContextsOffset = 0
				reply.command.CreateContextsLength = 0
			else:
				# this can be a file or a pipe or a printer etc.
				# we only handle files for now
				# first we need to translate the UNC path to a local path
				# this happens by getting the share name from the tree entry (via tree_id)
				if msg.header.TreeId not in self.__treeid_to_entry:
					raise Exception('Tree ID %s not found' % msg.header.TreeId)
				
				te = self.__treeid_to_entry[msg.header.TreeId]
				if te is None:
					raise Exception('Tree entry for Tree ID %s not found' % msg.header.TreeId)
				
				share = te.path

				file_name = command.Name
				local_path = share / Path(file_name).relative_to(share)
				if not local_path.exists():
					raise Exception('File %s does not exist' % local_path)
				if not local_path.is_file():
					raise Exception('File %s is not a file' % local_path)
				if local_path.is_file():
					file_id = 0
					handle = open(local_path, 'rb') #this must change based on the mode
					# store the handle in the file entry
					fe = FileEntry(path=local_path, handle=handle, file_id=file_id)
					te.file_entries[reply.command.FileId] = fe

				fstat = local_path.stat()
				reply = SMB2Message()
				reply.header = SMB2Header_SYNC()
				reply.header.Command = SMB2Command.CREATE
				reply.header.Flags = SMB2HeaderFlag.SMB2_FLAGS_SERVER_TO_REDIR
				reply.header.SessionId = self.SessionId
				reply.header.CreditCharge = 1
				reply.header.CreditReq = msg.header.CreditReq
				reply.header.Status = NTStatus.SUCCESS
				reply.command = CREATE_REPLY()
				reply.command.OplockLevel = OplockLevel.SMB2_OPLOCK_LEVEL_NONE
				reply.command.Flags = CreateOptions.FILE_WRITE_THROUGH
				reply.command.CreateAction = CreateAction.FILE_CREATED
				reply.command.CreationTime = convert_to_ft(fstat.st_ctime)
				reply.command.LastAccessTime = convert_to_ft(fstat.st_atime)
				reply.command.LastWriteTime = convert_to_ft(fstat.st_mtime)
				reply.command.ChangeTime = convert_to_ft(fstat.st_mtime)
				reply.command.AllocationSize = fstat.st_size
				reply.command.EndofFile = fstat.st_size
				reply.command.FileAttributes = FileAttributes.FILE_ATTRIBUTE_NORMAL
				reply.command.FileId = fe.file_id
				reply.command.CreateContextsOffset = 0
				reply.command.CreateContextsLength = 0
			#await self.sendSMB(reply)
			await self.sendSMB2(msg, SMB2Command.CREATE, reply.command, NTStatus.SUCCESS)
		except Exception as e:
			traceback.print_exc()
	
	async def query_directory(self, msg:SMB2Message):
		try:
			command = typing.cast(QUERY_DIRECTORY_REQ, msg.command)
			# Validate tree and determine base path
			if msg.header.TreeId not in self.__treeid_to_entry:
				raise Exception('Tree ID %s not found' % msg.header.TreeId)
			te = self.__treeid_to_entry[msg.header.TreeId]
			share = te.path
			# FileId 0 means the directory opened by CREATE against the tree root/share
			# We support listing the root of the share for now
			base_path = share
			# Determine pattern
			pattern = command.FileName if command.FileName else '*'
			# Special-case: when client sends FileName of "\\" or "." in root, list all
			if pattern in ['\\\x00', '\\', '.', '.\x00']:
				pattern = '*'
			# Load or initialize enumeration state per TreeId
			if not hasattr(self, '_dir_enum_state'):
				self._dir_enum_state = {}
			state_key = (msg.header.TreeId, command.FileId)
			if QueryDirectoryFlag.SMB2_RESTART_SCANS in command.Flags or state_key not in self._dir_enum_state or QueryDirectoryFlag.SMB2_REOPEN in command.Flags:
				# Build fresh list according to pattern
				entries = []
				for child in base_path.iterdir():
					name = child.name
					try:
						from aiosmb.commons.utils.glob2re import glob2re
						import re
						if not re.match(glob2re(pattern), name):
							continue
					except Exception:
						if pattern not in ['*', None]:
							continue
					entries.append(name)
				self._dir_enum_state[state_key] = DirEnumState(entries=entries, index=0, pattern=pattern, info_class=command.FileInformationClass)
			state = self._dir_enum_state[state_key]
			# Build from current index forward until OutputBufferLength filled
			entries = []
			for name in state.entries[state.index:]:
				child = base_path / name
				st = child.stat()
				# Build directory info record based on requested class
				# NextEntryOffset filled later when concatenating
				creation_time = convert_to_ft(int(st.st_ctime))
				last_access_time = convert_to_ft(int(st.st_atime))
				last_write_time = convert_to_ft(int(st.st_mtime))
				change_time = convert_to_ft(int(st.st_mtime))
				allocation_size = 0 if child.is_dir() else st.st_size
				end_of_file = 0 if child.is_dir() else st.st_size
				attrs = FileAttributes.FILE_ATTRIBUTE_DIRECTORY if child.is_dir() else FileAttributes.FILE_ATTRIBUTE_NORMAL
				filename_bytes = name.encode('utf-16-le')
				fixed = b''
				fixed += (0).to_bytes(4, 'little', signed=False) # NextEntryOffset placeholder
				fixed += (0).to_bytes(4, 'little', signed=False) # FileIndex
				fixed += creation_time.to_bytes(8, 'little', signed=False)
				fixed += last_access_time.to_bytes(8, 'little', signed=False)
				fixed += last_write_time.to_bytes(8, 'little', signed=False)
				fixed += change_time.to_bytes(8, 'little', signed=False)
				fixed += end_of_file.to_bytes(8, 'little', signed=True)
				fixed += allocation_size.to_bytes(8, 'little', signed=True)
				fixed += attrs.value.to_bytes(4, 'little', signed=False)
				if command.FileInformationClass == FileInfoClass.FileFullDirectoryInformation:
					fixed += len(filename_bytes).to_bytes(4, 'little', signed=False)
					fixed += (0).to_bytes(4, 'little', signed=False) # EaSize
					entry = fixed + filename_bytes
				elif command.FileInformationClass == FileInfoClass.FileDirectoryInformation:
					# Same as Full but without EaSize
					fixed += len(filename_bytes).to_bytes(4, 'little', signed=False)
					entry = fixed + filename_bytes
				elif command.FileInformationClass == FileInfoClass.FileBothDirectoryInformation:
					# Full + ShortName fields
					fixed += len(filename_bytes).to_bytes(4, 'little', signed=False)
					fixed += (0).to_bytes(4, 'little', signed=False) # EaSize
					fixed += (0).to_bytes(1, 'little', signed=False) # ShortNameLength
					fixed += (0).to_bytes(1, 'little', signed=False) # Reserved
					fixed += b'\x00' * 24 # ShortName (12 WCHAR)
					fixed += (0).to_bytes(2, 'little', signed=False) # Reserved2
					entry = fixed + filename_bytes
				elif command.FileInformationClass == FileInfoClass.FileIdFullDirectoryInformation:
					# Add FileId (8 bytes) then FileNameLength and EaSize
					file_id = getattr(st, 'st_ino', 0)
					fixed += file_id.to_bytes(8, 'little', signed=False)
					fixed += len(filename_bytes).to_bytes(4, 'little', signed=False)
					fixed += (0).to_bytes(4, 'little', signed=False) # EaSize
					entry = fixed + filename_bytes
				elif command.FileInformationClass == FileInfoClass.FileIdBothDirectoryInformation:
					# Both + FileId
					fixed += len(filename_bytes).to_bytes(4, 'little', signed=False)
					fixed += (0).to_bytes(4, 'little', signed=False) # EaSize
					fixed += (0).to_bytes(1, 'little', signed=False) # ShortNameLength
					fixed += (0).to_bytes(1, 'little', signed=False) # Reserved
					fixed += b'\x00' * 24 # ShortName (12 WCHAR)
					fixed += (0).to_bytes(2, 'little', signed=False) # Reserved2
					file_id = getattr(st, 'st_ino', 0)
					fixed += file_id.to_bytes(8, 'little', signed=False)
					entry = fixed + filename_bytes
				else:
					# Unsupported info class, skip for now
					continue
				
				if len(entry) % 8 != 0:
					entry += b'\x00' * (8 - len(entry) % 8)
				entries.append(entry)

			# Concatenate entries with proper NextEntryOffset (8-byte aligned) and enforce OutputBufferLength
			max_len = command.OutputBufferLength if command.OutputBufferLength else 65535
			buffer = b''
			emitted = 0
			for i, entry in enumerate(entries):
				cur_len = len(entry)
				# compute 8-byte alignment padding for this entry when followed by another
				padded_len = (cur_len + 7) & ~7
				pad_len = padded_len - cur_len
				# determine buffer start position for this entry
				start_pos = len(buffer)
				# If this entry alone doesn't fit, stop
				if start_pos + cur_len > max_len:
					break
				# Decide if we can chain a next entry: we need room for this entry + its padding + at least the next entry header
				can_chain_next = False
				if i + 1 < len(entries):
					next_len = len(entries[i+1])
					if start_pos + cur_len + pad_len + next_len <= max_len:
						can_chain_next = True
				# Append current entry
				buffer += entry
				emitted += 1
				if can_chain_next:
					# write NextEntryOffset (padded length) and append padding
					buffer = bytearray(buffer)
					buffer[start_pos:start_pos+4] = padded_len.to_bytes(4, 'little', signed=False)
					buffer = bytes(buffer)
					if pad_len:
						buffer += b'\x00' * pad_len
				else:
					# last entry in this chunk: leave NextEntryOffset as 0 and do not pad
					pass
			# Advance enumeration index by number of emitted entries
			state.index += emitted

			reply = SMB2Message()
			reply.header = SMB2Header_SYNC()
			reply.header.Command = SMB2Command.QUERY_DIRECTORY
			reply.header.Flags = SMB2HeaderFlag.SMB2_FLAGS_SERVER_TO_REDIR
			reply.header.SessionId = self.SessionId
			reply.header.TreeId = msg.header.TreeId
			reply.header.CreditCharge = 1
			reply.header.CreditReq = msg.header.CreditReq if msg.header.CreditReq else 1
			
			# Check if we have no data to return (either no buffer or all entries consumed)
			if len(buffer) == 0 or (emitted == 0 and state.index >= len(state.entries)):
				# No entries matched, nothing fits, or enumeration is complete
				reply.header.Status = NTStatus.NO_MORE_FILES
				reply.command = ERROR_REPLY()
				await self.sendSMB2(msg, SMB2Command.QUERY_DIRECTORY, reply.command, NTStatus.NO_MORE_FILES)
			else:
				reply.header.Status = NTStatus.SUCCESS
				reply.command = QUERY_DIRECTORY_REPLY()
				reply.command.Data = buffer
				await self.sendSMB2(msg, SMB2Command.QUERY_DIRECTORY, reply.command, NTStatus.SUCCESS)
			#await self.sendSMB(reply)
			
		except Exception as e:
			traceback.print_exc()

	async def close(self, msg:SMB2Message):
		try:
			command = typing.cast(CLOSE_REQ, msg.command)
			reply = SMB2Message()
			reply.header = SMB2Header_SYNC()
			reply.header.Command = SMB2Command.CLOSE
			reply.header.Flags = SMB2HeaderFlag.SMB2_FLAGS_SERVER_TO_REDIR
			reply.header.SessionId = self.SessionId
			reply.header.TreeId = msg.header.TreeId
			reply.header.CreditCharge = 1
			reply.header.CreditReq = msg.header.CreditReq
			reply.header.Status = NTStatus.SUCCESS
			reply.command = CLOSE_REPLY()
			reply.command.Flags = CloseFlag.NONE
			reply.command.Reserved = 0
			# Try to find the file entry by TreeId and FileId
			fe = None
			if msg.header.TreeId in self.__treeid_to_entry:
				te = self.__treeid_to_entry[msg.header.TreeId]
				if command.FileId in te.file_entries:
					fe = te.file_entries.pop(command.FileId)
			if fe and fe.handle:
				try:
					fe.handle.close()
				except Exception:
					pass
			# Fill attributes if requested
			if CloseFlag.SMB2_CLOSE_FLAG_POSTQUERY_ATTRIB in command.Flags and fe is not None:
				st = fe.path.stat()
				reply.command.CreationTime = convert_to_ft(int(st.st_ctime))
				reply.command.LastAccessTime = convert_to_ft(int(st.st_atime))
				reply.command.LastWriteTime = convert_to_ft(int(st.st_mtime))
				reply.command.ChangeTime = convert_to_ft(int(st.st_mtime))
				reply.command.AllocationSize = 0 if fe.path.is_dir() else st.st_size
				reply.command.EndofFile = 0 if fe.path.is_dir() else st.st_size
				reply.command.FileAttributes = (FileAttributes.FILE_ATTRIBUTE_DIRECTORY if fe.path.is_dir() else FileAttributes.FILE_ATTRIBUTE_NORMAL).value
				reply.command.Flags = CloseFlag.SMB2_CLOSE_FLAG_POSTQUERY_ATTRIB
			await self.sendSMB(reply)
		except Exception as e:
			traceback.print_exc()
	
	
	async def echo(self, msg:SMB2Message):
		try:
			command = typing.cast(ECHO_REQ, msg.command)
			reply = SMB2Message()
			reply.header = SMB2Header_SYNC()
			reply.header.Command = SMB2Command.ECHO
			reply.header.Flags = SMB2HeaderFlag.SMB2_FLAGS_SERVER_TO_REDIR
			reply.header.SessionId = self.SessionId
			reply.header.CreditCharge = msg.header.CreditCharge if msg.header.CreditCharge else 1
			reply.header.CreditReq = msg.header.CreditReq if msg.header.CreditReq else 1
			reply.header.Status = NTStatus.SUCCESS
			reply.command = ECHO_REPLY()
			await self.sendSMB(reply)
		except Exception as e:
			traceback.print_exc()
	
	async def sendSMB2(self, orgiginator:SMB2Message, command:SMB2Command, reply, status = NTStatus.SUCCESS):
		print(f"[DEBUG] sending SMB2 message with MessageId {orgiginator.header.MessageId}")
		replymsg = SMB2Message()
		replymsg.header = SMB2Header_SYNC()
		replymsg.header.MessageId = orgiginator.header.MessageId
		replymsg.header.Command = command
		replymsg.header.Flags = SMB2HeaderFlag.SMB2_FLAGS_SERVER_TO_REDIR
		replymsg.header.SessionId = self.SessionId
		# Fix: CreditCharge should match the request's CreditCharge (credits consumed)
		replymsg.header.CreditCharge = orgiginator.header.CreditCharge
		# Fix: CreditReq should be the credits we're granting back to the client
		replymsg.header.CreditReq = orgiginator.header.CreditReq if orgiginator.header.CreditReq else 1
		replymsg.header.Status = status
		replymsg.command = reply
		await self.sendSMB(replymsg)