From 3c15e009afbbaeea8af0c0cf77118c65301039e3 Mon Sep 17 00:00:00 2001
From: Alex Romero <ntalexio2@gmail.com>
Date: Thu, 11 Jul 2024 02:21:03 +0330
Subject: [PATCH] Add new test cases for test_smb.py (#1760)

* remove spaces

* add remained functions to test_getData() in test_smb.py

* add new test case test_queryInfo() in test_smb.py
---
 tests/SMB_RPC/test_smb.py | 27 +++++++++++++++++++++------
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/tests/SMB_RPC/test_smb.py b/tests/SMB_RPC/test_smb.py
index 52b2b380b8..b73f195445 100644
--- a/tests/SMB_RPC/test_smb.py
+++ b/tests/SMB_RPC/test_smb.py
@@ -87,7 +87,7 @@ def test_connection(self):
         self.assertEqual(credentials, (self.username, self.password, self.domain, '', '', '', None, None))
         smb.logoff()
         del(smb)
-        
+
     def test_close_connection(self):
         smb = self.create_connection()
         smb.login(self.username, self.password, self.domain)
@@ -156,7 +156,7 @@ def test_createFile(self):
         smb.deleteFile(self.share, self.file + '.bak')
         smb.disconnectTree(tid)
         smb.logoff()
-        
+
     def test_readwriteFile(self):
         smb = self.create_connection()
         smb.login(self.username, self.password, self.domain)
@@ -176,9 +176,8 @@ def test_readwriteFile(self):
         smb.closeFile(tid, fid)
         smb.deleteFile(self.share, self.file)
         smb.disconnectTree(tid)
-        
         smb.logoff()
-         
+
     def test_createdeleteDirectory(self):
         smb = self.create_connection()
         smb.login(self.username, self.password, self.domain)
@@ -194,7 +193,7 @@ def test_createdeleteDirectory(self):
                 smb.deleteDirectory(self.share, nested_dir)
                 smb.deleteDirectory(self.share, self.directory)
         smb.logoff()
- 
+
     def test_getData(self):
         smb = self.create_connection()
         smb.login(self.username, self.password, self.domain)
@@ -205,6 +204,8 @@ def test_getData(self):
         smb.getServerOS()
         smb.doesSupportNTLMv2()
         smb.isLoginRequired()
+        smb.isSigningRequired()
+        smb.getIOCapabilities()
         smb.logoff()
 
     def test_getServerName(self):
@@ -266,7 +267,21 @@ def test_getSessionKey(self):
         smb.login(self.username, self.password, self.domain)
         smb.getSessionKey()
         smb.logoff()
-        
+
+    def test_queryInfo(self):
+        smb = self.create_connection()
+        smb.login(self.username, self.password, self.domain)
+        tid = smb.connectTree(self.share)
+        fid = smb.createFile(tid, self.file)
+        file_info = smb.queryInfo(tid, fid)
+        self.assertEqual(file_info["AllocationSize"], 0)
+        self.assertEqual(file_info["EndOfFile"], 0)
+        self.assertEqual(file_info["Directory"], 0)
+        smb.closeFile(tid,fid)
+        smb.deleteFile(self.share, self.file)
+        smb.disconnectTree(tid)
+        smb.logoff()
+
     def __is_socket_opened(self, s):
         # We assume that if socket is selectable, it's open; and if it were not, it's closed.
         # Note: this method is accurate as long as the file descriptor used for the socket is not re-used