diff --git a/src/smbclient/_pool.py b/src/smbclient/_pool.py index c5c8cf1..3a7dc50 100644 --- a/src/smbclient/_pool.py +++ b/src/smbclient/_pool.py @@ -210,7 +210,7 @@ def dfs_request(tree: TreeConnect, path: str) -> DFSReferralResponse: return dfs_response -def delete_session(server, port=445, connection_cache=None): +def delete_session(server, port=445, connection_cache=None, timeout=60): """ Deletes the connection in the connection pool for the server specified. This will also close all sessions associated with the connection. @@ -226,7 +226,7 @@ def delete_session(server, port=445, connection_cache=None): connection = connection_cache.get(connection_key, None) if connection: del connection_cache[connection_key] - connection.disconnect(close=True) + connection.disconnect(close=True, timeout=timeout) def get_smb_tree( diff --git a/src/smbprotocol/connection.py b/src/smbprotocol/connection.py index be06a40..de7fc40 100644 --- a/src/smbprotocol/connection.py +++ b/src/smbprotocol/connection.py @@ -931,7 +931,7 @@ def connect(self, dialect=None, timeout=60, preferred_encryption_algos=None, pre elif context_type == NegotiateContextType.SMB2_SIGNING_CAPABILITIES: self.signing_algorithm_id = context["data"]["signing_algorithms"][0] - def disconnect(self, close=True): + def disconnect(self, close=True, timeout=None): """ Closes the connection as well as logs off any of the Disconnects the TCP connection and shuts down the socket listener @@ -943,7 +943,7 @@ def disconnect(self, close=True): # We cannot close the session or tree if the socket has been closed. if close and self.transport.connected: for session in list(self.session_table.values()): - session.disconnect(True) + session.disconnect(True, timeout=timeout) log.info("Disconnecting transport connection") self.transport.close() diff --git a/src/smbprotocol/session.py b/src/smbprotocol/session.py index f521173..d9a3481 100644 --- a/src/smbprotocol/session.py +++ b/src/smbprotocol/session.py @@ -394,7 +394,7 @@ def connect(self): log.info("Verifying the SMB Setup Session signature as auth is successful") self.connection.verify_signature(response, self.session_id, force=True) - def disconnect(self, close=True): + def disconnect(self, close=True, timeout=None): """ Logs off the session @@ -418,7 +418,7 @@ def disconnect(self, close=True): request = self.connection.send(logoff, sid=self.session_id) log.info("Session: %s - Receiving Logoff response", self.username) - res = self.connection.receive(request) + res = self.connection.receive(request, timeout=timeout) res_logoff = SMB2Logoff() res_logoff.unpack(res["data"].get_value()) log.debug(res_logoff)