from typing import Dict, List
import asyncio
import datetime
from contextlib import asynccontextmanager
from pathlib import PureWindowsPath

from aiosmb.commons.interfaces.file import SMBFile
from aiosmb.wintypes.access_mask import *
from aiosmb.wintypes.fscc.FileAttributes import FileAttributes
from aiosmb.wintypes.fscc.structures.fileinfoclass import FileInfoClass
from aiosmb.protocol.smb2.commands.query_info import SecurityInfo
from aiosmb.protocol.smb2.commands import *


from aiosmb.connection import SMBConnection
from aiosmb.commons.connection.target import SMBTarget

class SMBDirectory:
	def __init__(self):
		#self.parent_share = None
		self.tree_id:int = None #thisdescribes the share itself, not the directory!
		self.fullpath:str = None
		self.unc_path:str = None
		self.parent_dir = None
		self.name:str = None
		self.creation_time:datetime.datetime = None
		self.last_access_time:datetime.datetime = None
		self.last_write_time:datetime.datetime = None
		self.change_time:datetime.datetime = None
		self.allocation_size:datetime.datetime = None
		self.attributes:FileAttributes = None
		self.file_id:int = None
		self.security_descriptor = None
		
		self.files:Dict[str, SMBFile] = {}
		self.subdirs:Dict[str, SMBDirectory] = {}

	@staticmethod
	def from_uncpath(unc_path:str):
		"""
		Creates SMBFile object from the UNC path supplied.
		Example uncpath: \\\\127.0.0.1\\C$\\temp\\test.exe
		"""
		unc = PureWindowsPath(unc_path)
		f = SMBDirectory()
		f.share_path = unc.drive
		f.fullpath = '\\'.join(unc.parts[1:])
		f.unc_path = unc_path
		
		return f
	
	@staticmethod
	def from_smbtarget(target:SMBTarget):
		"""
		Creates SMBDirectory object from the SMBUrl object
		"""
		if target.path is None:
			return None
		
		fpath = target.path.replace('/','\\')
		temp = '\\\\%s%s'
		unc = temp % (target.get_hostname_or_ip(), fpath)
		return SMBDirectory.from_uncpath(unc)

	@staticmethod
	def from_remotepath(connection:SMBConnection, remotepath:str):
		"""
		Creates SMBFile object from the connection and the remote path supplied.
		Example remotepath: \\C$\\temp\\test.exe
		"""
		temp = '\\\\%s\\%s'
		if remotepath[0] == '\\':
			temp = '\\\\%s%s'
		unc = temp % (connection.target.get_hostname_or_ip(), remotepath)
		return SMBDirectory.from_uncpath(unc)

	@staticmethod
	async def delete_unc(connection:SMBConnection, remotepath:str):
		"""
		Deletes a directory at a given path.
		"""
		try:
			remfile = SMBDirectory.from_uncpath(remotepath)
			tree_entry, err = await connection.tree_connect(remfile.share_path)
			if err is not None:
				raise err
			tree_id = tree_entry.tree_id

			desired_access = FileAccessMask.DELETE | FileAccessMask.FILE_READ_ATTRIBUTES
			share_mode = ShareAccess.FILE_SHARE_DELETE
			create_options = CreateOptions.FILE_DIRECTORY_FILE | CreateOptions.FILE_DELETE_ON_CLOSE 
			create_disposition = CreateDisposition.FILE_OPEN
			file_attrs = 0

			file_id, err = await connection.create(tree_id, remfile.fullpath, desired_access, share_mode, create_options, create_disposition, file_attrs, return_reply = False)
			if err is not None:
				raise err
			if file_id is not None:
				await connection.close(tree_id, file_id)

			await connection.tree_disconnect(tree_id)
			return True, None
		
		except Exception as e:
			return False, e

	@staticmethod
	async def create_remote(connection:SMBConnection, remotepath:str):
		try:
			remfile = SMBDirectory.from_remotepath(connection, remotepath)
			tree_entry, err = await connection.tree_connect(remfile.share_path)
			if err is not None:
				raise err
			tree_id = tree_entry.tree_id

			desired_access = FileAccessMask.FILE_READ_DATA | FileAccessMask.FILE_WRITE_DATA | FileAccessMask.FILE_READ_EA | FileAccessMask.FILE_WRITE_EA | FileAccessMask.FILE_READ_ATTRIBUTES | FileAccessMask.FILE_WRITE_ATTRIBUTES | FileAccessMask.READ_CONTROL | FileAccessMask.DELETE | FileAccessMask.SYNCHRONIZE
			share_mode = ShareAccess.FILE_SHARE_READ | ShareAccess.FILE_SHARE_WRITE | ShareAccess.FILE_SHARE_DELETE
			create_options = CreateOptions.FILE_DIRECTORY_FILE | CreateOptions.FILE_SYNCHRONOUS_IO_NONALERT
			create_disposition = CreateDisposition.FILE_CREATE
			file_attrs = 0

			file_id, err = await connection.create(tree_id, remfile.fullpath, desired_access, share_mode, create_options, create_disposition, file_attrs, return_reply = False)
			if err is not None:
				raise err
			if file_id is not None:
				await connection.close(tree_id, file_id)

			await connection.tree_disconnect(tree_id)
			return True, None

		except Exception as e:
			return False, e

	def get_share_path(self):
		unc = PureWindowsPath(self.unc_path)
		return unc.drive

	def get_subdir_paths(self, fullpath:bool = False):
		"""
		Returns a list of paths of all subdirectories in the directory
		"""
		paths = []
		for f in self.subdirs.values():
			if fullpath is True:
				paths.append(f.unc_path)
			else:
				paths.append(f.fullpath)
		return paths

	def get_file_paths(self, fullpath:bool = False):
		"""
		Returns a list of paths of all files in the directory
		"""
		paths = []
		for f in self.files.values():
			if fullpath is True:
				paths.append(f.unc_path)
			else:
				paths.append(f.fullpath)
		return paths

	async def delete(self, connection:SMBConnection):
		try:
			if self.file_id is not None:
				# if the directory is open, first we need to close it
				await connection.close(self.tree_id, self.file_id)
			return await SMBDirectory.delete_unc(connection, self.unc_path)
		except Exception as e:
			return False, e

	async def get_security_descriptor(self, connection:SMBConnection):
		if self.security_descriptor is None:
			file_id = None
			try:
				tree_id = self.tree_id
				if self.tree_id is None:
					tree_entry, err = await connection.tree_connect(self.share_path)
					if err is not None:
						raise err
					tree_id = tree_entry.tree_id

				desired_access = FileAccessMask.READ_CONTROL
				share_mode = ShareAccess.FILE_SHARE_READ
				create_options = CreateOptions.FILE_DIRECTORY_FILE | CreateOptions.FILE_SYNCHRONOUS_IO_NONALERT
				file_attrs = 0
				create_disposition = CreateDisposition.FILE_OPEN
				file_id, err = await connection.create(tree_id, self.fullpath, desired_access, share_mode, create_options, create_disposition, file_attrs)
				if err is not None:
					raise err
				
				self.security_descriptor, err = await connection.query_info(
					tree_id, 
					file_id,
					info_type = QueryInfoType.SECURITY, 
					information_class = FileInfoClass.NONE, 
					additional_information = SecurityInfo.ATTRIBUTE_SECURITY_INFORMATION | SecurityInfo.DACL_SECURITY_INFORMATION | SecurityInfo.OWNER_SECURITY_INFORMATION | SecurityInfo.GROUP_SECURITY_INFORMATION, 
					flags = 0, 
				)
				if err is not None:
					raise err
			except Exception as e:
				return None, e

			finally:
				if file_id is not None:
					await connection.close(tree_id, file_id)
				if tree_id is not None and self.tree_id is None:
					await connection.tree_disconnect(tree_id)


		return self.security_descriptor, None
	

	def get_console_output(self):
		lines = []
		for name in self.subdirs:
			directory = self.subdirs[name]
			entry = '%s\t%s\t%s\t%s' % ('drw-rw-rw-',  directory.allocation_size, directory.creation_time, directory.name)
			lines.append(entry)
		for name in self.files:
			directory = self.files[name]
			entry = '%s\t%s\t%s\t%s' % (' rw-rw-rw-',  directory.allocation_size, directory.creation_time, directory.name)
			lines.append(entry)
		return lines
		
	async def create_subdir(self, dir_name:str, connection:SMBConnection):
		try:
			should_close = False #dont close the tree_id only if the directory hasnt been connected to yet
			if not self.tree_id:
				should_close = True
				tree_entry, err = await connection.tree_connect(self.get_share_path())
				if err is not None:
					raise err
				self.tree_id = tree_entry.tree_id

			file_id = None
			newpath = dir_name
			if self.fullpath != '':
				newpath = '%s\\%s' % (self.fullpath, dir_name)
			try:
				file_id, err = await connection.create(
					self.tree_id, 
					newpath, 
					FileAccessMask.FILE_READ_DATA | FileAccessMask.FILE_WRITE_DATA | FileAccessMask.FILE_READ_EA | FileAccessMask.FILE_WRITE_EA | FileAccessMask.FILE_READ_ATTRIBUTES | FileAccessMask.FILE_WRITE_ATTRIBUTES | FileAccessMask.READ_CONTROL | FileAccessMask.DELETE | FileAccessMask.SYNCHRONIZE, 
					ShareAccess.FILE_SHARE_READ | ShareAccess.FILE_SHARE_WRITE | ShareAccess.FILE_SHARE_DELETE,
					CreateOptions.FILE_DIRECTORY_FILE | CreateOptions.FILE_SYNCHRONOUS_IO_NONALERT, 
					CreateDisposition.FILE_CREATE, 
					0
				)
				if err is not None:
					raise err
				return True, None
			finally:
				if file_id is not None:
					await connection.close(self.tree_id, file_id)
				if should_close is True:       
					await connection.tree_disconnect(self.tree_id)

		except Exception as e:
			return False, e
	
	async def delete_subdir(self, dir_name:str):
		raise Exception('delete subdir not implemented!')

	async def list_r(self, connection:SMBConnection, depth:int = 3, maxentries:int = None, fetch_dir_sd:bool = False, fetch_file_sd:bool = False, exclude_dir:List[str] = [], filter_cb=None):
		"""
		recursive list files and folders
		Beware this will clear out the lists of files/folders to save memory!
		"""
		if depth == 0:
			return
		depth -= 1
		ctr = 0

		try:
			async for obj, otype, err in self.list_gen(connection):
				await asyncio.sleep(0)
				if otype == 'dir' and fetch_dir_sd is True and obj.name not in exclude_dir:
					obj.tree_id = self.tree_id
					_, err = await obj.get_security_descriptor(connection)
					#if err is not None:
					#	print(err)
				
				if otype == 'file' and fetch_file_sd is True:
					obj.tree_id = self.tree_id
					_, err = await obj.get_security_descriptor(connection)
					#if err is not None:
					#	print(err)

				yield obj, otype, err
				
				ctr += 1
				if err is not None:
					break
				
				if ctr == maxentries:
					yield self, 'maxed', None
					break
				
				if otype == 'dir' and obj.name not in exclude_dir:
					if filter_cb is not None:
						res = await filter_cb('dir', obj)
						if res is False:
							continue
					obj.tree_id = self.tree_id
					async for e,t,err in obj.list_r(
							connection, 
							depth, 
							maxentries = maxentries, 
							exclude_dir = exclude_dir,
							fetch_dir_sd = fetch_dir_sd, 
							fetch_file_sd = fetch_file_sd,
							filter_cb = filter_cb):
						yield e,t,err
						# await asyncio.sleep(0)
		except Exception as e:
			yield None, None, e

	async def snapshots(self, connection, path = '\\'):
		"""
		Lists all snapshots of the share
		"""
		try:
			should_close = False
			if self.tree_id is None:
				should_close = True
				tree_entry, err = await connection.tree_connect(self.get_share_path())
				if err is not None:
					raise err
				self.tree_id = tree_entry.tree_id

			fpath = self.fullpath
			print(fpath)
			file_id, err = await connection.create(self.tree_id, '', FileAccessMask.READ_CONTROL|FileAccessMask.FILE_READ_DATA, ShareAccess.FILE_SHARE_READ, CreateOptions.FILE_DIRECTORY_FILE | CreateOptions.FILE_SYNCHRONOUS_IO_NONALERT, CreateDisposition.FILE_OPEN, 0)
			if err is not None:
				raise err
			snapshots, err = await connection.ioctl(self.tree_id, file_id, CtlCode.FSCTL_SRV_ENUMERATE_SNAPSHOTS, flags = IOCTLREQFlags.IS_FSCTL, MaxOutputResponse=16)
			if err is not None:
				return None, err

			# checking if we actually got all the snapshots
			if len(snapshots.SnapShots) < snapshots.NumberOfSnapShots:
				snapshots, err = await connection.ioctl(self.tree_id, file_id, CtlCode.FSCTL_SRV_ENUMERATE_SNAPSHOTS, flags = IOCTLREQFlags.IS_FSCTL, MaxOutputResponse=snapshots.SnapShotArraySize+12)
				if err is not None:
					return None, err

			return snapshots.SnapShots, None
		except Exception as e:
			return None, e
		finally:
			if file_id is not None:
				await connection.close(self.tree_id, file_id)
			if should_close is True:
				await connection.tree_disconnect(self.tree_id)

	async def list_gen(self, connection:SMBConnection):
		"""
		Lists all files and folders in the directory, yields the results as they arrive
		directory: SMBDirectory
		DOESN'T fill the SMBDirectory's data
		"""
		self.files = {}
		self.subdirs = {}
		
		desired_access = FileAccessMask.FILE_READ_DATA
		share_mode = ShareAccess.FILE_SHARE_READ
		create_options = CreateOptions.FILE_DIRECTORY_FILE | CreateOptions.FILE_SYNCHRONOUS_IO_NONALERT
		file_attrs = 0
		create_disposition = CreateDisposition.FILE_OPEN

	
		if not self.tree_id:
			tree_entry, err = await connection.tree_connect(self.get_share_path())
			if err is not None:
				yield self, 'dir', err
				return
			self.tree_id = tree_entry.tree_id
		
		file_id, err = await connection.create(self.tree_id, self.fullpath, desired_access, share_mode, create_options, create_disposition, file_attrs)
		if err is not None:
			yield self, 'dir', err
			return
		try:
			while True:
				await asyncio.sleep(0)
				fileinfos, err = await connection.query_directory(self.tree_id, file_id)
				if err is not None:
					raise err
				if not fileinfos:
					break
				for info in fileinfos:
					if info.FileAttributes & FileAttributes.FILE_ATTRIBUTE_DIRECTORY:
						dirname = info.FileName 
						if info.FileName in ['.','..']:
							continue
						subdir = SMBDirectory()
						subdir.tree_id = self.tree_id
						if self.fullpath != '':
							subdir.fullpath = '%s\\%s' % (self.fullpath, info.FileName)	
						else:
							subdir.fullpath = info.FileName
						subdir.unc_path = '%s\\%s' % (self.unc_path, info.FileName)
						subdir.parent_dir = self
						subdir.name = info.FileName
						subdir.creation_time = info.CreationTime
						subdir.last_access_time = info.LastAccessTime
						subdir.last_write_time = info.LastWriteTime
						subdir.change_time = info.ChangeTime
						subdir.allocation_size = info.AllocationSize
						subdir.attributes = info.FileAttributes
						
						yield subdir, 'dir', None
						
					else:
						file = SMBFile()
						file.tree_id = self.tree_id
						file.parent_dir = None
						if self.fullpath != '':
							file.fullpath = '%s\\%s' % (self.fullpath, info.FileName)
							file.unc_path = '%s\\%s' % (self.unc_path, info.FileName)
						else:
							file.fullpath = info.FileName
							file.unc_path = '%s\\%s' % (self.unc_path, info.FileName)
						file.name = info.FileName
						file.size = info.EndOfFile
						file.creation_time = info.CreationTime
						file.last_access_time = info.LastAccessTime
						file.last_write_time = info.LastWriteTime
						file.change_time = info.ChangeTime
						file.allocation_size = info.AllocationSize
						file.attributes = info.FileAttributes
						yield file, 'file', None
			
		except Exception as e:
			yield self, 'dir', e
			return
		finally:
			if file_id is not None:
				await connection.close(self.tree_id, file_id)

	async def list(self, connection:SMBConnection):
		"""
		Lists all files and folders in the directory
		directory: SMBDirectory
		fills the SMBDirectory's data
		"""
		self.files = {}
		self.subdirs = {}
		
		desired_access = FileAccessMask.FILE_READ_DATA
		share_mode = ShareAccess.FILE_SHARE_READ
		create_options = CreateOptions.FILE_DIRECTORY_FILE | CreateOptions.FILE_SYNCHRONOUS_IO_NONALERT
		file_attrs = 0
		create_disposition = CreateDisposition.FILE_OPEN

		
		if not self.tree_id:
			tree_entry, err = await connection.tree_connect(self.get_share_path())
			if err is not None:
				raise err
			self.tree_id = tree_entry.tree_id
		
		file_id, err = await connection.create(self.tree_id, self.fullpath, desired_access, share_mode, create_options, create_disposition, file_attrs)
		if err is not None:
			return False, err
		try:
			while True:
				fileinfos, err = await connection.query_directory(self.tree_id, file_id)
				if err is not None:
					raise err
				if not fileinfos:
					break
				for info in fileinfos:
					if info.FileAttributes & FileAttributes.FILE_ATTRIBUTE_DIRECTORY:
						dirname = info.FileName 
						if info.FileName in ['.','..']:
							continue
						subdir = SMBDirectory()
						subdir.tree_id = self.tree_id
						if self.fullpath != '':
							subdir.fullpath = '%s\\%s' % (self.fullpath, info.FileName)	
						else:
							subdir.fullpath = info.FileName
						subdir.unc_path = '%s\\%s' % (self.unc_path, info.FileName)
						subdir.parent_dir = self
						subdir.name = info.FileName
						subdir.creation_time = info.CreationTime
						subdir.last_access_time = info.LastAccessTime
						subdir.last_write_time = info.LastWriteTime
						subdir.change_time = info.ChangeTime
						subdir.allocation_size = info.AllocationSize
						subdir.attributes = info.FileAttributes
						
						self.subdirs[subdir.name] = subdir
						
					else:
						file = SMBFile()
						#file.parent_share = directory.parent_share
						file.tree_id = self.tree_id
						file.parent_dir = None
						if self.fullpath != '':
							file.fullpath = '%s\\%s' % (self.fullpath, info.FileName)
							file.unc_path = '%s\\%s' % (self.unc_path, info.FileName)
						else:
							file.fullpath = info.FileName
							file.unc_path = '%s\\%s' % (self.unc_path, info.FileName)
						file.name = info.FileName
						file.size = info.EndOfFile
						file.creation_time = info.CreationTime
						file.last_access_time = info.LastAccessTime
						file.last_write_time = info.LastWriteTime
						file.change_time = info.ChangeTime
						file.allocation_size = info.AllocationSize
						file.attributes = info.FileAttributes		
						self.files[file.name] = file
			return True, None
		except Exception as e:
			return False, e
		finally:
			if file_id is not None:
				await connection.close(self.tree_id, file_id)
		
	def __str__(self):
		t = '===== DIRECTORY ===== \r\n'
		for k in self.__dict__:
			if k.startswith('parent_'):
				continue
			if isinstance(self.__dict__[k], list):
				for item in self.__dict__[k]:
					t += '%s : %s\r\n' % (k, item)
			elif isinstance(self.__dict__[k], dict):
				for ks in self.__dict__[k]:
					t += '%s : %s\r\n' % (ks, self.__dict__[k][ks])
			else:
				t += '%s : %s\r\n' % (k, self.__dict__[k])
		
		return t

async def smb_mkdir(path:str, connection = None):
	"""
	Creates a new directory on the remote SMB server.
	"""
	if path.lower().startswith('smb'):
			from aiosmb.commons.connection.factory import SMBConnectionFactory
			factory = SMBConnectionFactory.from_url(path)
			if connection is None:
				connection = factory.get_connection()
				async with connection:
					_, err = await connection.login()
					if err is not None:
						raise err
					smbdir = factory.get_directory()
					_, err = await smbdir.create_remote(connection, smbdir.fullpath)
					if err is not None:
						raise err
					return
			else:
				smbdir = factory.get_directory()
				_, err = await smbdir.create_remote(connection, smbdir.fullpath)
				if err is not None:
					raise err
				return
	else:
		smbdir = SMBDirectory.from_remotepath(connection, path)
		_, err = await smbdir.create_remote(connection, smbdir.fullpath)
		if err is not None:
			raise err


async def smb_rmdir(path:str, connection = None):
	"""
	Removes a directory from the SMB server
	"""
	if path.lower().startswith('smb'):
			from aiosmb.commons.connection.factory import SMBConnectionFactory
			factory = SMBConnectionFactory.from_url(path)
			if connection is None:
				connection = factory.get_connection()
				async with connection:
					_, err = await connection.login()
					if err is not None:
						raise err
					smbdir = factory.get_directory()
					_, err = await smbdir.delete(connection)
					if err is not None:
						raise err
					return
			else:
				smbdir = factory.get_directory()
				_, err = await smbdir.delete(connection)
				if err is not None:
					raise err
				return
	else:
		smbdir = SMBDirectory.from_remotepath(connection, path)
		_, err = await smbdir.delete(connection)
		if err is not None:
			raise err

async def smb_walk(path:str, connection = None, fullpath = False):
	"""
	Enumerates files and subdirectories of a directory from the SMB server
	"""
	async def _walk(smbdir:SMBDirectory, connection, fullpath):
		_, err = await smbdir.list(connection)
		if err is None:
			yield smbdir, smbdir.get_subdir_paths(fullpath), smbdir.get_file_paths(fullpath)
			for subdir in smbdir.subdirs:
				async for x in _walk(smbdir.subdirs[subdir], connection, fullpath):
					yield x
		
		
	if path.lower().startswith('smb'):
			from aiosmb.commons.connection.factory import SMBConnectionFactory
			factory = SMBConnectionFactory.from_url(path)
			if connection is None:
				connection = factory.get_connection()
				async with connection:
					_, err = await connection.login()
					if err is not None:
						raise err
					smbdir = factory.get_directory()
			else:
				smbdir = factory.get_directory()
	else:
		smbdir = SMBDirectory.from_remotepath(connection, path)
	
	async for smbdir, subdirs, files in _walk(smbdir, connection, fullpath):
		yield smbdir, subdirs, files
		

async def amain():
	url = 'smb2+ntlm-password://TEST\\Administrator:Passw0rd!1@10.10.10.2/C$'
	#await smb_mkdir(url)
	
	from aiosmb.commons.connection.factory import SMBConnectionFactory
	connection = SMBConnectionFactory.from_url(url).get_connection()
	_, err = await connection.login()
	if err is not None:
		raise err

	async for smbdir, subdirs, files in smb_walk(url, connection, fullpath=True):
		print(smbdir.fullpath)
		print(subdirs)
		print(files)
		print('')
	#await smb_mkdir('C$\\temp\\temp2', connection=connection)
	#await smb_rmdir('C$\\temp\\temp2', connection=connection)
	#async with smb_open('C$/temp/repodata.json', 'rb', connection=connection) as f:
	#	while True:
	#		data, err = await f.read(1024)
	#		print(data)
	#		print(len(data))
	#		if data == b'':
	#			break
	


if __name__ == '__main__':
	import asyncio
	asyncio.run(amain())

