improve base file class

Signed-off-by: pranjalg1331 <[email protected]>

Author: pranjalg1331

tests/api_app/analyzers_manager/unit_tests/file_analyzers/base_test_class.py

@@ -1,4 +1,5 @@
 import hashlib
+import logging
 import os
 from contextlib import ExitStack
 from types import SimpleNamespace
@@ -7,6 +8,8 @@
 from api_app.analyzers_manager.models import AnalyzerConfig
 from tests.mock_utils import MockUpResponse
 
+logger = logging.getLogger(__name__)
+
 
 class BaseFileAnalyzerTest(TestCase):
     analyzer_class = None
@@ -59,71 +62,83 @@ def get_all_supported_mimetypes(cls) -> set:
         return set(cls.MIMETYPE_TO_FILENAME.keys())
 
     @classmethod
-    def get_extra_config(self) -> dict:
+    def get_extra_config(cls) -> dict:
         """
         Subclasses can override this to provide additional runtime configuration
         specific to their analyzer (e.g., API keys, URLs, retry counts, etc.).
-
-        Returns:
-            dict: Extra configuration parameters for the analyzer
         """
         return {}
 
     def get_mocked_response(self):
         """
         Subclasses override this to define expected mocked output.
-
-        Can return:
-        1. A single patch object: patch('module.function')
-        2. A list of patch objects: [patch('module.func1'), patch('module.func2')]
-        3. A context manager: patch.multiple() or ExitStack()
         """
-        raise NotImplementedError
+        raise NotImplementedError("Subclasses must implement get_mocked_response()")
 
     @classmethod
-    def _apply_patches(self, patches):
+    def _apply_patches(cls, patches):
         """Helper method to apply single or multiple patches"""
         if patches is None:
             return ExitStack()  # No-op context manager
 
-        # If it's already a context manager, return as-is
         if hasattr(patches, "__enter__") and hasattr(patches, "__exit__"):
             return patches
 
-        # If it's a list of patches, use ExitStack to manage them
         if isinstance(patches, (list, tuple)):
             stack = ExitStack()
             for patch_obj in patches:
                 stack.enter_context(patch_obj)
             return stack
 
-        # Single patch object
         return patches
 
+    def setUp(self):
+        super().setUp()
+        logger.info("Setting up test environment for file analyzer")
+        if self.analyzer_class:
+            analyzer_module = self.analyzer_class.__module__
+            logging.getLogger(analyzer_module).setLevel(logging.CRITICAL)
+            logging.getLogger("api_app.analyzers_manager").setLevel(logging.WARNING)
+
+    def tearDown(self):
+        super().tearDown()
+        logger.info("Tearing down test environment for file analyzer")
+        if self.analyzer_class:
+            analyzer_module = self.analyzer_class.__module__
+            logging.getLogger(analyzer_module).setLevel(logging.NOTSET)
+            logging.getLogger("api_app.analyzers_manager").setLevel(logging.NOTSET)
+
     def test_analyzer_on_supported_filetypes(self):
         if self.analyzer_class is None:
             self.skipTest("analyzer_class is not set")
-        config = AnalyzerConfig.objects.get(
-            python_module=self.analyzer_class.python_module
-        )
-        print(config)
 
-        # If supported_filetypes is None or empty, use all available mimetypes
-        if config.supported_filetypes:
-            supported_types = config.supported_filetypes
-        else:
-            supported_types = self.get_all_supported_mimetypes()
+        logger.info("Starting file analyzer test for: %s", self.analyzer_class.__name__)
+
+        try:
+            config = AnalyzerConfig.objects.get(
+                python_module=self.analyzer_class.python_module
+            )
+        except AnalyzerConfig.DoesNotExist:
+            self.fail(
+                f"No AnalyzerConfig found for {self.analyzer_class.python_module}"
+            )
+
+        logger.debug("Loaded analyzer config: %s", config)
+
+        supported_types = (
+            config.supported_filetypes or self.get_all_supported_mimetypes()
+        )
 
         for mimetype in supported_types:
             with self.subTest(mimetype=mimetype):
+                logger.info("Testing mimetype: %s", mimetype)
 
                 try:
                     file_bytes = self.get_sample_file_bytes(mimetype)
-                except (ValueError, OSError):
-                    print(f"SKIPPING {mimetype}")
+                except (ValueError, OSError) as e:
+                    logger.warning("Skipping %s due to error: %s", mimetype, str(e))
                     continue
 
-                # Apply patches using the improved system
                 patches = self.get_mocked_response()
                 with self._apply_patches(patches):
                     md5 = hashlib.md5(file_bytes).hexdigest()
@@ -133,26 +148,37 @@ def test_analyzer_on_supported_filetypes(self):
                     analyzer.filename = f"test_file_{mimetype}"
                     analyzer.md5 = md5
                     analyzer.read_file_bytes = lambda: file_bytes
+
                     analyzer._job = SimpleNamespace()
                     analyzer._job.analyzable = SimpleNamespace()
                     analyzer._job.analyzable.name = analyzer.filename
                     analyzer._job.analyzable.mimetype = mimetype
                     analyzer._job.analyzable.sha256 = hashlib.sha256(
                         file_bytes
                     ).hexdigest()
+                    analyzer._job_id = ""
 
-                    test_file_path = self.get_sample_file_path(mimetype)
-                    analyzer._FileAnalyzer__filepath = test_file_path
+                    analyzer._FileAnalyzer__filepath = self.get_sample_file_path(
+                        mimetype
+                    )
 
-                    extra_config = self.get_extra_config()
-                    for key, value in extra_config.items():
+                    for key, value in self.get_extra_config().items():
                         setattr(analyzer, key, value)
 
-                    response = analyzer.run()
+                    try:
+                        response = analyzer.run()
+                        logger.info("Analyzer ran successfully for %s", mimetype)
+                    except Exception as e:
+                        logger.exception(
+                            "Analyzer raised an exception for %s", mimetype
+                        )
+                        self.fail(
+                            f"Analyzer run failed for {mimetype}: {type(e).__name__}: {e}"
+                        )
+
                     self.assertIsInstance(
                         response,
                         (dict, MockUpResponse),
                         f"Expected dict or MockUpResponse, got {type(response)} with value: {response}",
                     )
-
-                    print(f"SUCCESS {mimetype}")
+                    logger.debug("Successful result for %s: %s", mimetype, response)

tests/api_app/analyzers_manager/unit_tests/file_analyzers/test_quark_engine.py

@@ -0,0 +1,45 @@
+from unittest.mock import MagicMock, patch
+
+from api_app.analyzers_manager.file_analyzers.quark_engine import QuarkEngine
+
+from .base_test_class import BaseFileAnalyzerTest
+
+
+class TestQuarkEngine(BaseFileAnalyzerTest):
+    analyzer_class = QuarkEngine
+
+    def get_extra_config(self):
+        return {}
+
+    def get_mocked_response(self):
+        # Mock the Report class and its methods
+        mock_report = MagicMock()
+        mock_json_report = {
+            "crimes": [
+                {
+                    "crime": "Send SMS",
+                    "weight": 50,
+                    "confidence": "80%",
+                    "permissions": ["android.permission.SEND_SMS"],
+                    "api": [
+                        {
+                            "class": "Landroid/telephony/SmsManager",
+                            "method": "sendTextMessage",
+                        }
+                    ],
+                    "register": [
+                        {"class": "Lcom/example/MainActivity", "method": "onCreate"}
+                    ],
+                }
+            ],
+            "total_score": 50,
+            "summary": {"threat_level": "Medium", "total_crimes": 1},
+        }
+        mock_report.get_report.return_value = mock_json_report
+
+        return [
+            # Mock the Report class from where it's actually imported
+            patch("quark.report.Report", return_value=mock_report),
+            # Mock the DIR_PATH from where it's actually imported
+            patch("quark.config.DIR_PATH", "/mock/rules/path"),
+        ]

tests/api_app/analyzers_manager/unit_tests/file_analyzers/test_rlt_info.py

@@ -0,0 +1,49 @@
+# RTFInfo Test Class
+from unittest.mock import MagicMock, patch
+
+from api_app.analyzers_manager.file_analyzers.rtf_info import RTFInfo
+
+from .base_test_class import BaseFileAnalyzerTest
+
+
+class TestRTFInfo(BaseFileAnalyzerTest):
+    analyzer_class = RTFInfo
+
+    def get_extra_config(self):
+        return {}
+
+    def get_mocked_response(self):
+        # Create mock RTF objects
+        mock_rtf_obj1 = MagicMock()
+        mock_rtf_obj1.is_ole = True
+        mock_rtf_obj1.class_name = b"OLE2Link"
+        mock_rtf_obj1.format_id = 2
+        mock_rtf_obj1.oledata_size = 1024
+        mock_rtf_obj1.is_package = False
+        mock_rtf_obj1.oledata_md5 = "abc123def456"
+        mock_rtf_obj1.clsid = "00000000-0000-0000-C000-000000000046"
+        mock_rtf_obj1.clsid_desc = "OLE Link Object"
+
+        mock_rtf_obj2 = MagicMock()
+        mock_rtf_obj2.is_ole = True
+        mock_rtf_obj2.class_name = b"Equation.3"
+        mock_rtf_obj2.format_id = 1
+        mock_rtf_obj2.oledata_size = 2048
+        mock_rtf_obj2.is_package = True
+        mock_rtf_obj2.filename = "malicious.exe"
+        mock_rtf_obj2.src_path = "C:\\temp\\malicious.exe"
+        mock_rtf_obj2.temp_path = "C:\\temp\\temp123.tmp"
+        mock_rtf_obj2.olepkgdata_md5 = "def456ghi789"
+        mock_rtf_obj2.clsid = None
+
+        # Create mock parser instance
+        mock_parser_instance = MagicMock()
+        mock_parser_instance.objects = [mock_rtf_obj1, mock_rtf_obj2]
+        mock_parser_instance.parse.return_value = None
+
+        return [
+            patch(
+                "api_app.analyzers_manager.file_analyzers.rtf_info.RtfObjParser",
+                return_value=mock_parser_instance,
+            ),
+        ]

tests/api_app/analyzers_manager/unit_tests/file_analyzers/test_signature_info.py

@@ -0,0 +1,32 @@
+from unittest.mock import MagicMock, patch
+
+from api_app.analyzers_manager.file_analyzers.signature_info import SignatureInfo
+
+from .base_test_class import BaseFileAnalyzerTest
+
+
+class TestSignatureInfo(BaseFileAnalyzerTest):
+    analyzer_class = SignatureInfo
+
+    def get_extra_config(self):
+        return {}
+
+    def get_mocked_response(self):
+        # Create mock process
+        mock_process = MagicMock()
+        mock_process.returncode = 0
+        mock_process.communicate.return_value = (
+            b"Signature verification: ok\nCertificate info:\n  Subject: CN=Test Certificate\n  Issuer: CN=Test CA\n",
+            b"",
+        )
+
+        return [
+            patch(
+                "api_app.analyzers_manager.file_analyzers.signature_info.Popen",
+                return_value=mock_process,
+            ),
+            patch(
+                "api_app.analyzers_manager.file_analyzers.signature_info.settings.PROJECT_LOCATION",
+                "/mock/project",
+            ),
+        ]

tests/api_app/analyzers_manager/unit_tests/file_analyzers/test_strings_info.py

@@ -0,0 +1,53 @@
+from unittest.mock import patch
+
+from api_app.analyzers_manager.file_analyzers.strings_info import StringsInfo
+
+from .base_test_class import BaseFileAnalyzerTest
+
+
+class TestStringsInfo(BaseFileAnalyzerTest):
+    analyzer_class = StringsInfo
+
+    def get_extra_config(self):
+        return {
+            "max_number_of_strings": 100,
+            "max_characters_for_string": 200,
+            "rank_strings": 1,  # Enable string ranking
+        }
+
+    def get_mocked_response(self):
+        # Mock strings extracted from binary
+        mock_strings_first_call = [
+            "http://malicious-site.com/payload",
+            "CreateFileA",
+            "WriteFile",
+            "RegSetValueA",
+            "https://c2server.evil/beacon",
+            "malware.exe",
+            "C:\\Windows\\System32\\cmd.exe",
+            "ftp://badsite.com/data",
+            "Some random string",
+            "Another test string",
+        ]
+
+        # Mock ranked strings (second call when rank_strings is enabled)
+        mock_ranked_strings = [
+            "http://malicious-site.com/payload",
+            "https://c2server.evil/beacon",
+            "CreateFileA",
+            "ftp://badsite.com/data",
+            "WriteFile",
+            "RegSetValueA",
+            "malware.exe",
+            "C:\\Windows\\System32\\cmd.exe",
+        ]
+
+        return [
+            patch(
+                "api_app.analyzers_manager.file_analyzers.strings_info.StringsInfo._docker_run",
+                side_effect=[
+                    mock_strings_first_call,  # First call for flarestrings
+                    mock_ranked_strings,  # Second call for rank_strings (if enabled)
+                ],
+            ),
+        ]

tests/api_app/analyzers_manager/unit_tests/file_analyzers/test_suricate.py

@@ -0,0 +1,39 @@
+from unittest.mock import patch
+
+from api_app.analyzers_manager.file_analyzers.suricata import Suricata
+
+from .base_test_class import BaseFileAnalyzerTest
+
+
+class TestSuricata(BaseFileAnalyzerTest):
+    analyzer_class = Suricata
+
+    def get_extra_config(self):
+        return {"reload_rules": True, "extended_logs": False, "signatures": {}}
+
+    def get_mocked_response(self):
+        # Mock Suricata analysis results with network alerts
+        mock_suricata_report = {
+            "data": [],
+            "stats": {
+                "capture": {"kernel_packets": 150, "kernel_drops": 0},
+                "decoder": {
+                    "pkts": 150,
+                    "bytes": 45600,
+                    "invalid": 0,
+                    "ipv4": 148,
+                    "ipv6": 2,
+                    "ethernet": 150,
+                    "tcp": 100,
+                    "udp": 48,
+                    "icmpv4": 2,
+                },
+            },
+        }
+
+        return [
+            patch(
+                "api_app.analyzers_manager.file_analyzers.suricata.Suricata._docker_run",
+                return_value=mock_suricata_report,
+            ),
+        ]