Update dependencies in setup.py and enhance error messages in SMBConnection. Add new user management methods in SMBMachine and implement snapshot listing in SMBShare and SMBDirectory. Improve directory query handling in SMBServerConnection and add user management commands in SMBClient for better usability.

Author: SkelSec

aiosmb/commons/connection/factory.py

@@ -90,6 +90,20 @@ def get_credential(self):
 		"""Returns a new SPNEGOCredential object with the credential from the factory"""
 		return SPNEGOCredential([copy.deepcopy(self.credential)]).build_context()
 
+	async def test_connection(self):
+		ntlm_data = None
+		try:
+			connection = self.get_connection()
+			async with connection:
+				_, err = await connection.login()
+				ntlm_data = connection.get_extra_info()
+				if err is not None:
+					raise err
+
+				return True, ntlm_data, None
+		except Exception as e:
+			return False, ntlm_data, e
+
 	@staticmethod
 	def from_components(ip_or_hostname:str, username:str, secret:str, secrettype:str = 'password', 
 							domain:str = None, port:str = 445, dialect:str = 'smb', dcip:str = None, proxies = None, authproto:str = 'ntlm',

aiosmb/commons/interfaces/directory.py

@@ -316,6 +316,42 @@ async def list_r(self, connection:SMBConnection, depth:int = 3, maxentries:int =
 		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):
 		"""

aiosmb/commons/interfaces/file.py

@@ -342,6 +342,8 @@ async def open(self, connection:SMBConnection, mode:str = 'r', share_mode = None
 	async def download(self, connection:SMBConnection ,local_path:str, share_mode = ShareAccess.FILE_SHARE_READ | ShareAccess.FILE_SHARE_WRITE):
 		"""Downloads the file to the local path. File must not be open."""
 		try:
+			# if treeid is already present, we won't be closing it after download
+			closetree = True if self.tree_id is None else False
 			if self.is_pipe:
 				raise Exception('Cannot download a pipe!')
 			if self.__connection is not None:
@@ -365,10 +367,10 @@ async def download(self, connection:SMBConnection ,local_path:str, share_mode =
 
 			return str(local_path.absolute()), None
 		except Exception as e:
-			return False, e
+			return None, e #this should return "None, e" and NOT "False, e"
 		finally:
 			await self.close()
-			if self.tree_id is not None:
+			if self.tree_id is not None and closetree is True:
 				await connection.tree_disconnect(self.tree_id)
 
 	async def open_pipe(self, connection:SMBConnection, mode:str):

aiosmb/commons/interfaces/machine.py

@@ -203,6 +203,98 @@ async def list_group_members(self, domain_name:str, group_name:str) -> AsyncGene
 		except Exception as e:
 			yield None, None, None, e
 
+	async def add_new_user(self, domain_name:str, user_name:str, password:str):
+		try:
+			async with samrrpc_from_smb(self.connection, auth_level=self.force_rpc_auth) as samrpc:
+				if domain_name is None or domain_name == '' or domain_name == 'Builtin':
+					dhandle, err = await samrpc.openBuiltinDomain()
+				else:
+					dhandle, err = await samrpc.open_domain_by_name(domain_name)
+				if err is not None:
+					raise err
+				
+				_, err = await samrpc.create_user(dhandle, user_name, password)
+				if err is not None:
+					raise err
+				return True, None
+
+		except Exception as e:
+			return None, e
+	
+	async def enable_user(self, domain_name:str, user_name:str) -> Awaitable[Tuple[str, Union[Exception, None]]]:
+		try:
+			async with samrrpc_from_smb(self.connection, auth_level=self.force_rpc_auth) as samrpc:
+				async with samrpc.get_user_handle(domain_name, user_name) as uhandle:
+					return await samrpc.enable_account(uhandle)
+
+		except Exception as e:
+			import traceback
+			traceback.print_exc()
+			return None, e
+
+	async def disable_user(self, domain_name:str, user_name:str) -> Awaitable[Tuple[str, Union[Exception, None]]]:
+		try:
+			async with samrrpc_from_smb(self.connection, auth_level=self.force_rpc_auth) as samrpc:
+				async with samrpc.get_user_handle(domain_name, user_name) as uhandle:
+					return await samrpc.disable_account(uhandle)
+		except Exception as e:
+			return None, e
+	
+	async def delete_user(self, domain_name:str, user_name:str) -> Awaitable[Tuple[str, Union[Exception, None]]]:
+		try:
+			async with samrrpc_from_smb(self.connection, auth_level=self.force_rpc_auth) as samrpc:
+				async with samrpc.get_user_handle(domain_name, user_name) as uhandle:
+					return await samrpc.disable_account(uhandle)
+
+		except Exception as e:
+			return None, e
+
+	async def change_user_password(self, domain_name:str, user_name:str, old_password:str, new_password:str) -> Awaitable[Tuple[str, Union[Exception, None]]]:
+		try:
+			async with samrrpc_from_smb(self.connection, auth_level=self.force_rpc_auth) as samrpc:
+				async with samrpc.get_user_handle(domain_name, user_name) as uhandle:
+					return await samrpc.change_account_password(uhandle, old_password, new_password)
+
+		except Exception as e:
+			return None, e
+
+	async def change_user_password_nt(self, domain_name:str, user_name:str, old_pw_hash:str, new_password:str) -> Awaitable[Tuple[str, Union[Exception, None]]]:
+		try:
+			async with samrrpc_from_smb(self.connection, auth_level=self.force_rpc_auth) as samrpc:
+				async with samrpc.get_user_handle(domain_name, user_name) as uhandle:
+					return await samrpc.change_account_password_nt(uhandle, old_pw_hash, new_password)
+
+		except Exception as e:
+			return None, e
+	
+	async def change_user_password_new(self, domain_name:str, user_name:str, new_password:str) -> Awaitable[Tuple[str, Union[Exception, None]]]:
+		try:
+			async with samrrpc_from_smb(self.connection, auth_level=self.force_rpc_auth) as samrpc:
+				async with samrpc.get_user_handle(domain_name, user_name) as uhandle:
+					return await samrpc.change_account_password_4(uhandle, new_password)
+
+		except Exception as e:
+			return None, e
+
+
+	async def change_user_password_internal(self, domain_name:str, user_name:str, new_password:str) -> Awaitable[Tuple[str, Union[Exception, None]]]:
+		try:
+			async with samrrpc_from_smb(self.connection, auth_level=self.force_rpc_auth) as samrpc:
+				async with samrpc.get_user_handle(domain_name, user_name) as uhandle:
+					return await samrpc.change_account_password_internal(uhandle, new_password)
+
+		except Exception as e:
+			return None, e
+	
+	async def list_users_allowed_to_replicate(self, domain_name:str):
+		try:
+			async with samrrpc_from_smb(self.connection, auth_level=self.force_rpc_auth) as samrpc:
+				async for name, sid, rid, err in samrpc.list_users_allowed_to_replicate(domain_name):
+					yield name, sid, rid, err
+		except Exception as e:
+			yield None, None, None, err
+			
+
 	async def add_sid_to_group(self, domain_name:str, group_name:str, sid:str) -> Awaitable[Tuple[bool, Union[Exception, None]]]:
 		try:
 			async with samrrpc_from_smb(self.connection, auth_level=self.force_rpc_auth) as samrpc:
@@ -372,6 +464,13 @@ async def list_services(self) -> AsyncGenerator[Tuple[SMBService, Union[Exceptio
 		except Exception as e:
 			yield None, e
 
+	async def get_service_config(self, service_name:str) -> Awaitable[Tuple[ServiceStatus, Union[Exception, None]]]:
+		try:
+			async with remsvcrpc_from_smb(self.connection, auth_level=self.force_rpc_auth) as rpc:
+				return await rpc.get_config(service_name)
+		except Exception as e:
+			return None, e
+
 	async def enable_service(self, service_name:str) -> Awaitable[Tuple[bool, Union[Exception, None]]]:
 		try:
 			async with remsvcrpc_from_smb(self.connection, auth_level=self.force_rpc_auth) as rpc:

aiosmb/commons/interfaces/share.py

@@ -49,6 +49,27 @@ async def connect(self, connection):
 		except Exception as e:
 			return None, e
 
+	async def snapshots(self, connection, path = '\\'):
+		"""
+		Lists all snapshots of the share
+		"""
+		try:
+			if self.tree_id is None:
+				_, err = await self.connect(connection)
+				if err is not None:
+					raise err
+				tree_id = self.tree_id
+			if tree_id is not None:
+				file_id, err = await connection.create(tree_id, path, FileAccessMask.READ_CONTROL, 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(tree_id, file_id, CtlCode.FSCTL_SRV_ENUMERATE_SNAPSHOTS)
+			if err is not None:
+				return None, err
+			return snapshots, None
+		except Exception as e:
+			return None, e
+
 	async def get_security_descriptor(self, connection) -> Awaitable[Tuple[SECURITY_DESCRIPTOR, Union[Exception, None]]]:
 		if self.security_descriptor is None:
 			file_id = None

aiosmb/commons/utils/glob2re.py

@@ -38,4 +38,5 @@ def glob2re(pat):
                 res = '%s[%s]' % (res, stuff)
         else:
             res = res + re.escape(c)
-    return res + '\Z(?ms)'
+    # Anchor to end of string; no inline flags
+    return res + r'\Z'

aiosmb/connection.py

@@ -997,7 +997,7 @@ async def create(self, tree_id, file_path, desired_access, share_mode, create_op
 				raise SMBConnectionTerminated()
 
 			if tree_id not in self.TreeConnectTable_id:
-				raise Exception('Unknown Tree ID!')
+				raise Exception('Unknown Tree ID! %s' % tree_id)
 
 			command = CREATE_REQ()
 			command.RequestedOplockLevel  = oplock_level
@@ -1058,9 +1058,9 @@ async def read(self, tree_id, file_id, offset = 0, length = 0):
 				raise SMBConnectionTerminated()
 
 			if tree_id not in self.TreeConnectTable_id:
-				raise Exception('Unknown Tree ID!')
+				raise Exception('Unknown Tree ID! %s' % tree_id)
 			if file_id not in self.FileHandleTable:
-				raise Exception('Unknown File ID!')
+				raise Exception('Unknown File ID! %s' % file_id)
 
 			async with self.FileHandleTable[file_id].oplock:
 				header = SMB2Header_SYNC()
@@ -1174,9 +1174,9 @@ async def query_info(self, tree_id, file_id, info_type = QueryInfoType.FILE, inf
 				raise SMBConnectionTerminated()
 
 			if tree_id not in self.TreeConnectTable_id:
-				raise Exception('Unknown Tree ID!')
+				raise Exception('Unknown Tree ID! %s' % tree_id)
 			if file_id not in self.FileHandleTable:
-				raise Exception('Unknown File ID!')
+				raise Exception('Unknown File ID! %s' % file_id)
 
 			command = QUERY_INFO_REQ()
 			command.InfoType = info_type
@@ -1239,9 +1239,9 @@ async def query_directory(self, tree_id, file_id, search_pattern = '*', resume_i
 				raise SMBConnectionTerminated()
 
 			if tree_id not in self.TreeConnectTable_id:
-				raise Exception('Unknown Tree ID!')
+				raise Exception('Unknown Tree ID! %s' % tree_id)
 			if file_id not in self.FileHandleTable:
-				raise Exception('Unknown File ID!')
+				raise Exception('Unknown File ID! %s' % file_id)
 
 
 			command = QUERY_DIRECTORY_REQ()
@@ -1281,7 +1281,7 @@ async def query_directory(self, tree_id, file_id, search_pattern = '*', resume_i
 		except Exception as e:
 			return None, e
 
-	async def ioctl(self, tree_id, file_id, ctlcode, data = None, flags = IOCTLREQFlags.IS_IOCTL):
+	async def ioctl(self, tree_id, file_id, ctlcode, data = None, flags = IOCTLREQFlags.IS_IOCTL, MaxOutputResponse=65535):
 		try:
 			if self.session_closed == True or self.status == SMBConnectionStatus.CLOSED:
 				raise SMBConnectionTerminated()
@@ -1291,6 +1291,7 @@ async def ioctl(self, tree_id, file_id, ctlcode, data = None, flags = IOCTLREQFl
 			command.FileId  = file_id
 			command.Flags = flags
 			command.Buffer = data
+			command.MaxOutputResponse = MaxOutputResponse
 
 			header = SMB2Header_SYNC()
 			header.Command  = SMB2Command.IOCTL

aiosmb/dcerpc/v5/interfaces/remoteregistry.py

@@ -258,7 +258,7 @@ async def DeleteValue(self, key, value):
 		try:
 			key = await self.__get_rawhandle(key)
 
-			_, err = await  rrp.hBaseRegDeleteValue(self.dce, key, value)
+			_, err = await rrp.hBaseRegDeleteValue(self.dce, key, value)
 			if err is not None:
 				raise err