Cleaned up unit tests

Signed-off-by: Andy Ragusa <[email protected]>
Cisco-Talos · Jun 17, 2024 · 74de119 · 74de119
1 parent de7eeba
commit 74de119
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 71 deletions.
diff --git a/unit_tests/clamscan/ole2_encryption_test.py b/unit_tests/clamscan/ole2_encryption_test.py
@@ -7,6 +7,7 @@
 import sys
 import os
 import re
+import shutil
 
 sys.path.append('../unit_tests')
 import testcase
@@ -21,19 +22,6 @@ def setUpClass(cls):
     def tearDownClass(cls):
         super(TC, cls).tearDownClass()
 
-    # Find the metadata.json file and verify its contents.
-    def verify_metadata_json(self, tempdir, expected=[], unexpected=[]):
-        for parent, dirs, files in os.walk(tempdir):
-            for f in files:
-                if "metadata.json" == f:
-                    with open(os.path.join(parent, f)) as handle:
-                        metadata_json = handle.read()
-                        self.verify_output(metadata_json, expected=expected, unexpected=unexpected)
-
-                    # There is only one metadata.json per scan.
-                    # We found it, so we can break out of the loop.
-                    break
-
     def setUp(self):
         super(TC, self).setUp()
 
@@ -82,10 +70,8 @@ def test_FAT_doc_metadata(self):
 
         assert output.ec == 0  # clean
 
-        expected_strings = [ 
-            '"Encrypted":1',
-        ]
-        self.assertStrings(tempdir, neededStrings)
+        expected_strings = [ '"Encrypted":"ENCRYPTION_TYPE_UNKNOWN"' ]
+        self.verify_metadata_json(tempdir, expected_strings)
 
     def test_ministream_doc(self):
         self.step_name('Test ministream doc')
@@ -123,10 +109,8 @@ def test_ministream_doc_metadata(self):
 
         assert output.ec == 0  # clean
 
-        neededStrings = [ '"Encrypted":1'
-                , '"EncryptedWithVelvetSweatshop":0'
-                ]
-        self.assertStrings(tempdir, neededStrings)
+        expected_strings = [ '"Encrypted":"ENCRYPTION_TYPE_UNKNOWN"' ]
+        self.verify_metadata_json(tempdir, expected_strings)
 
 
     def test_FAT_docx(self):
@@ -165,10 +149,8 @@ def test_FAT_docx_metadata(self):
 
         assert output.ec == 0  # clean
 
-        neededStrings = [ '"Encrypted":1'
-                , '"EncryptedWithVelvetSweatshop":0'
-                ]
-        self.assertStrings(tempdir, neededStrings)
+        expected_strings = [ '"Encrypted":"ENCRYPTION_TYPE_UNKNOWN"' ]
+        self.verify_metadata_json(tempdir, expected_strings)
 
     def test_ministream_docx(self):
         self.step_name('Test ministream docx')
@@ -206,10 +188,8 @@ def test_ministream_docx_metadata(self):
 
         assert output.ec == 0  # clean
 
-        neededStrings = [ '"Encrypted":1'
-                , '"EncryptedWithVelvetSweatshop":0'
-                ]
-        self.assertStrings(tempdir, neededStrings)
+        expected_strings = [ '"Encrypted":"ENCRYPTION_TYPE_UNKNOWN"' ]
+        self.verify_metadata_json(tempdir, expected_strings)
 
     def test_FAT_dot(self):
         self.step_name('Test FAT dot')
@@ -247,10 +227,8 @@ def test_FAT_dot_metadata(self):
 
         assert output.ec == 0  # clean
 
-        neededStrings = [ '"Encrypted":1'
-                , '"EncryptedWithVelvetSweatshop":0'
-                ]
-        self.assertStrings(tempdir, neededStrings)
+        expected_strings = [ '"Encrypted":"ENCRYPTION_TYPE_UNKNOWN"' ]
+        self.verify_metadata_json(tempdir, expected_strings)
 
     def test_ministream_dot(self):
         self.step_name('Test ministream dot')
@@ -288,10 +266,8 @@ def test_ministream_dot_metadata(self):
 
         assert output.ec == 0  # clean
 
-        neededStrings = [ '"Encrypted":1'
-                , '"EncryptedWithVelvetSweatshop":0'
-                ]
-        self.assertStrings(tempdir, neededStrings)
+        expected_strings = [ '"Encrypted":"ENCRYPTION_TYPE_UNKNOWN"' ]
+        self.verify_metadata_json(tempdir, expected_strings)
 
     def test_FAT_ppsx(self):
         self.step_name('Test FAT ppsx')
@@ -329,10 +305,8 @@ def test_FAT_ppsx_metadata(self):
 
         assert output.ec == 0  # clean
 
-        neededStrings = [ '"Encrypted":1'
-                , '"EncryptedWithVelvetSweatshop":0'
-                ]
-        self.assertStrings(tempdir, neededStrings)
+        expected_strings = [ '"Encrypted":"ENCRYPTION_TYPE_UNKNOWN"' ]
+        self.verify_metadata_json(tempdir, expected_strings)
 
     def test_ministream_ppsx(self):
         self.step_name('Test ministream ppsx')
@@ -370,10 +344,8 @@ def test_ministream_ppsx_metadata(self):
 
         assert output.ec == 0  # clean
 
-        neededStrings = [ '"Encrypted":1'
-                , '"EncryptedWithVelvetSweatshop":0'
-                ]
-        self.assertStrings(tempdir, neededStrings)
+        expected_strings = [ '"Encrypted":"ENCRYPTION_TYPE_UNKNOWN"' ]
+        self.verify_metadata_json(tempdir, expected_strings)
 
     def test_FAT_pptx(self):
         self.step_name('Test FAT pptx')
@@ -411,10 +383,8 @@ def test_FAT_pptx_metadata(self):
 
         assert output.ec == 0  # clean
 
-        neededStrings = [ '"Encrypted":1'
-                , '"EncryptedWithVelvetSweatshop":0'
-                ]
-        self.assertStrings(tempdir, neededStrings)
+        expected_strings = [ '"Encrypted":"ENCRYPTION_TYPE_UNKNOWN"' ]
+        self.verify_metadata_json(tempdir, expected_strings)
 
     def test_ministream_pptx(self):
         self.step_name('Test ministream pptx')
@@ -452,10 +422,8 @@ def test_ministream_pptx_metadata(self):
 
         assert output.ec == 0  # clean
 
-        neededStrings = [ '"Encrypted":1'
-                , '"EncryptedWithVelvetSweatshop":0'
-                ]
-        self.assertStrings(tempdir, neededStrings)
+        expected_strings = [ '"Encrypted":"ENCRYPTION_TYPE_UNKNOWN"' ]
+        self.verify_metadata_json(tempdir, expected_strings)
 
     def test_FAT_xls(self):
         self.step_name('Test FAT xls')
@@ -493,11 +461,8 @@ def test_FAT_xls_metadata(self):
 
         assert output.ec == 0  # clean
 
-        neededStrings = [ '"Encrypted":1'
-                , '"EncryptedWithVelvetSweatshop":0'
-                , '"RC4Encryption":1'
-                ]
-        self.assertStrings(tempdir, neededStrings)
+        expected_strings = [ '"Encrypted":"RC4"' ]
+        self.verify_metadata_json(tempdir, expected_strings)
 
     def test_ministream_xls(self):
         self.step_name('Test ministream xls')
@@ -535,11 +500,8 @@ def test_ministream_xls_metadata(self):
 
         assert output.ec == 0  # clean
 
-        neededStrings = [ '"Encrypted":1'
-                , '"EncryptedWithVelvetSweatshop":0'
-                , '"RC4Encryption":1'
-                ]
-        self.assertStrings(tempdir, neededStrings)
+        expected_strings = [ '"Encrypted":"RC4"' ]
+        self.verify_metadata_json(tempdir, expected_strings)
 
     def test_FAT_xlsx(self):
         self.step_name('Test FAT xlsx')
@@ -577,10 +539,8 @@ def test_FAT_xlsx_metadata(self):
 
         assert output.ec == 0  # clean
 
-        neededStrings = [ '"Encrypted":1'
-                , '"EncryptedWithVelvetSweatshop":0'
-                ]
-        self.assertStrings(tempdir, neededStrings)
+        expected_strings = [ '"Encrypted":"ENCRYPTION_TYPE_UNKNOWN"' ]
+        self.verify_metadata_json(tempdir, expected_strings)
 
     def test_ministream_xlsx(self):
         self.step_name('Test ministream xlsx')
@@ -618,10 +578,8 @@ def test_ministream_xlsx_metadata(self):
 
         assert output.ec == 0  # clean
 
-        neededStrings = [ '"Encrypted":1'
-                , '"EncryptedWithVelvetSweatshop":0'
-                ]
-        self.assertStrings(tempdir, neededStrings)
+        expected_strings = [ '"Encrypted":"ENCRYPTION_TYPE_UNKNOWN"' ]
+        self.verify_metadata_json(tempdir, expected_strings)
 
 
 

diff --git a/unit_tests/testcase.py b/unit_tests/testcase.py
@@ -508,6 +508,20 @@ def execute_command(self, cmd, **kwargs):
         """
         return self.execute(cmd, **kwargs)
 
+    # Find the metadata.json file and verify its contents.
+    def verify_metadata_json(self, tempdir, expected=[], unexpected=[]):
+        for parent, dirs, files in os.walk(tempdir):
+            for f in files:
+                if "metadata.json" == f:
+                    with open(os.path.join(parent, f)) as handle:
+                        metadata_json = handle.read()
+                        self.verify_output(metadata_json, expected=expected, unexpected=unexpected)
+
+                    # There is only one metadata.json per scan.
+                    # We found it, so we can break out of the loop.
+                    break
+
+
 
 class Logger(object):