Skip to content

Commit 1f6d8e7

Browse files
pranjalg1331pranjalg1331
committed
file analyzers test+ base class improved
Signed-off-by: pranjalg1331 <[email protected]>
1 parent 2e14055 commit 1f6d8e7

34 files changed

+1481
-33
lines changed

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

Lines changed: 62 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import hashlib
2+
import logging
23
import os
34
from contextlib import ExitStack
45
from types import SimpleNamespace
@@ -7,6 +8,8 @@
78
from api_app.analyzers_manager.models import AnalyzerConfig
89
from tests.mock_utils import MockUpResponse
910

11+
logger = logging.getLogger(__name__)
12+
1013

1114
class BaseFileAnalyzerTest(TestCase):
1215
analyzer_class = None
@@ -59,71 +62,83 @@ def get_all_supported_mimetypes(cls) -> set:
5962
return set(cls.MIMETYPE_TO_FILENAME.keys())
6063

6164
@classmethod
62-
def get_extra_config(self) -> dict:
65+
def get_extra_config(cls) -> dict:
6366
"""
6467
Subclasses can override this to provide additional runtime configuration
6568
specific to their analyzer (e.g., API keys, URLs, retry counts, etc.).
66-
67-
Returns:
68-
dict: Extra configuration parameters for the analyzer
6969
"""
7070
return {}
7171

7272
def get_mocked_response(self):
7373
"""
7474
Subclasses override this to define expected mocked output.
75-
76-
Can return:
77-
1. A single patch object: patch('module.function')
78-
2. A list of patch objects: [patch('module.func1'), patch('module.func2')]
79-
3. A context manager: patch.multiple() or ExitStack()
8075
"""
81-
raise NotImplementedError
76+
raise NotImplementedError("Subclasses must implement get_mocked_response()")
8277

8378
@classmethod
84-
def _apply_patches(self, patches):
79+
def _apply_patches(cls, patches):
8580
"""Helper method to apply single or multiple patches"""
8681
if patches is None:
8782
return ExitStack() # No-op context manager
8883

89-
# If it's already a context manager, return as-is
9084
if hasattr(patches, "__enter__") and hasattr(patches, "__exit__"):
9185
return patches
9286

93-
# If it's a list of patches, use ExitStack to manage them
9487
if isinstance(patches, (list, tuple)):
9588
stack = ExitStack()
9689
for patch_obj in patches:
9790
stack.enter_context(patch_obj)
9891
return stack
9992

100-
# Single patch object
10193
return patches
10294

95+
def setUp(self):
96+
super().setUp()
97+
logger.info("Setting up test environment for file analyzer")
98+
if self.analyzer_class:
99+
analyzer_module = self.analyzer_class.__module__
100+
logging.getLogger(analyzer_module).setLevel(logging.CRITICAL)
101+
logging.getLogger("api_app.analyzers_manager").setLevel(logging.WARNING)
102+
103+
def tearDown(self):
104+
super().tearDown()
105+
logger.info("Tearing down test environment for file analyzer")
106+
if self.analyzer_class:
107+
analyzer_module = self.analyzer_class.__module__
108+
logging.getLogger(analyzer_module).setLevel(logging.NOTSET)
109+
logging.getLogger("api_app.analyzers_manager").setLevel(logging.NOTSET)
110+
103111
def test_analyzer_on_supported_filetypes(self):
104112
if self.analyzer_class is None:
105113
self.skipTest("analyzer_class is not set")
106-
config = AnalyzerConfig.objects.get(
107-
python_module=self.analyzer_class.python_module
108-
)
109-
print(config)
110114

111-
# If supported_filetypes is None or empty, use all available mimetypes
112-
if config.supported_filetypes:
113-
supported_types = config.supported_filetypes
114-
else:
115-
supported_types = self.get_all_supported_mimetypes()
115+
logger.info("Starting file analyzer test for: %s", self.analyzer_class.__name__)
116+
117+
try:
118+
config = AnalyzerConfig.objects.get(
119+
python_module=self.analyzer_class.python_module
120+
)
121+
except AnalyzerConfig.DoesNotExist:
122+
self.fail(
123+
f"No AnalyzerConfig found for {self.analyzer_class.python_module}"
124+
)
125+
126+
logger.debug("Loaded analyzer config: %s", config)
127+
128+
supported_types = (
129+
config.supported_filetypes or self.get_all_supported_mimetypes()
130+
)
116131

117132
for mimetype in supported_types:
118133
with self.subTest(mimetype=mimetype):
134+
logger.info("Testing mimetype: %s", mimetype)
119135

120136
try:
121137
file_bytes = self.get_sample_file_bytes(mimetype)
122-
except (ValueError, OSError):
123-
print(f"SKIPPING {mimetype}")
138+
except (ValueError, OSError) as e:
139+
logger.warning("Skipping %s due to error: %s", mimetype, str(e))
124140
continue
125141

126-
# Apply patches using the improved system
127142
patches = self.get_mocked_response()
128143
with self._apply_patches(patches):
129144
md5 = hashlib.md5(file_bytes).hexdigest()
@@ -133,23 +148,37 @@ def test_analyzer_on_supported_filetypes(self):
133148
analyzer.filename = f"test_file_{mimetype}"
134149
analyzer.md5 = md5
135150
analyzer.read_file_bytes = lambda: file_bytes
151+
136152
analyzer._job = SimpleNamespace()
137153
analyzer._job.analyzable = SimpleNamespace()
138154
analyzer._job.analyzable.name = analyzer.filename
139155
analyzer._job.analyzable.mimetype = mimetype
156+
analyzer._job.analyzable.sha256 = hashlib.sha256(
157+
file_bytes
158+
).hexdigest()
159+
analyzer._job_id = ""
140160

141-
test_file_path = self.get_sample_file_path(mimetype)
142-
analyzer._FileAnalyzer__filepath = test_file_path
161+
analyzer._FileAnalyzer__filepath = self.get_sample_file_path(
162+
mimetype
163+
)
143164

144-
extra_config = self.get_extra_config()
145-
for key, value in extra_config.items():
165+
for key, value in self.get_extra_config().items():
146166
setattr(analyzer, key, value)
147167

148-
response = analyzer.run()
168+
try:
169+
response = analyzer.run()
170+
logger.info("Analyzer ran successfully for %s", mimetype)
171+
except Exception as e:
172+
logger.exception(
173+
"Analyzer raised an exception for %s", mimetype
174+
)
175+
self.fail(
176+
f"Analyzer run failed for {mimetype}: {type(e).__name__}: {e}"
177+
)
178+
149179
self.assertIsInstance(
150180
response,
151181
(dict, MockUpResponse),
152182
f"Expected dict or MockUpResponse, got {type(response)} with value: {response}",
153183
)
154-
155-
print(f"SUCCESS {mimetype}")
184+
logger.debug("Successful result for %s: %s", mimetype, response)

‎tests/api_app/analyzers_manager/unit_tests/file_analyzers/test_elf_info.py‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
# from unittest.mock import patch, mock_open, MagicMock
2+
# from api_app.analyzers_manager.file_analyzers.elf_info import ELFInfo
3+
# from .base_test_class import BaseFileAnalyzerTest
4+
# from elftools.construct import Container
5+
6+
7+
# class TestELFInfo(BaseFileAnalyzerTest):
8+
# analyzer_class = ELFInfo
9+
10+
# def get_mocked_response(self):
11+
# # Mock ELF file content and structure
12+
# mock_elf = MagicMock()
13+
14+
# # Mock ELF header as Container object
15+
# mock_e_ident = Container()
16+
# mock_e_ident.EI_MAG = [0x7f, 0x45, 0x4c, 0x46] # ELF magic number
17+
# mock_e_ident.EI_CLASS = 2 # 64-bit
18+
# mock_e_ident.EI_DATA = 1 # Little endian
19+
# mock_e_ident.EI_VERSION = 1
20+
# mock_e_ident.EI_OSABI = 0
21+
# mock_e_ident.EI_ABIVERSION = 0
22+
23+
# mock_header = Container()
24+
# mock_header.e_ident = mock_e_ident
25+
# mock_header.e_type = 2 # ET_EXEC
26+
# mock_header.e_machine = 62 # EM_X86_64
27+
# mock_header.e_version = 1
28+
# mock_header.e_entry = 0x400000
29+
# mock_header.e_phoff = 64
30+
# mock_header.e_shoff = 4096
31+
# mock_header.e_flags = 0
32+
# mock_header.e_ehsize = 64
33+
# mock_header.e_phentsize = 56
34+
# mock_header.e_phnum = 8
35+
# mock_header.e_shentsize = 64
36+
# mock_header.e_shnum = 29
37+
# mock_header.e_shstrndx = 28
38+
39+
# # Configure mock ELF object
40+
# mock_elf.header = mock_header
41+
# mock_elf.elfclass = 64
42+
# mock_elf.little_endian = True
43+
44+
# # Mock telfhash result
45+
# mock_telfhash_result = ({
46+
# 'telfhash': 'T1234567890ABCDEF1234567890ABCDEF12345678',
47+
# 'header': {
48+
# 'machine': 'x86_64',
49+
# 'class': 'ELF64',
50+
# 'data': 'little_endian'
51+
# },
52+
# 'sections': [
53+
# {'name': '.text', 'size': 1024},
54+
# {'name': '.data', 'size': 512}
55+
# ],
56+
# 'symbols': 42,
57+
# 'imports': 15,
58+
# 'exports': 8
59+
# },)
60+
61+
# # Create patches for all the dependencies
62+
# patches = [
63+
# patch('builtins.open', mock_open()),
64+
# patch('elftools.elf.elffile.ELFFile', return_value=mock_elf),
65+
# patch('telfhash.telfhash', return_value=mock_telfhash_result)
66+
# ]
67+
68+
# return patches

‎tests/api_app/analyzers_manager/unit_tests/file_analyzers/test_file_info.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
from unittest.mock import patch
2+
3+
from api_app.analyzers_manager.file_analyzers.file_info import FileInfo
4+
5+
from .base_test_class import BaseFileAnalyzerTest
6+
7+
8+
class TestFileInfo(BaseFileAnalyzerTest):
9+
analyzer_class = FileInfo
10+
11+
def get_mocked_response(self):
12+
return [
13+
# Mock file type detection
14+
patch(
15+
"magic.from_file",
16+
side_effect=lambda path, mime=False: (
17+
"application/pdf" if mime else "PDF document, version 1.4"
18+
),
19+
),
20+
# Mock hash functions
21+
patch(
22+
"api_app.helpers.calculate_md5",
23+
return_value="d41d8cd98f00b204e9800998ecf8427e",
24+
),
25+
patch(
26+
"api_app.helpers.calculate_sha1",
27+
return_value="da39a3ee5e6b4b0d3255bfef95601890afd80709",
28+
),
29+
patch(
30+
"api_app.helpers.calculate_sha256",
31+
return_value="e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
32+
),
33+
patch("pydeep.hash_file", return_value=b"3:AOn4:An"),
34+
patch("tlsh.hash", return_value="T1234567890ABCDEF"),
35+
# Disable exiftool to avoid subprocess issues
36+
patch.object(FileInfo, "exiftool_path", None),
37+
]

‎tests/api_app/analyzers_manager/unit_tests/file_analyzers/test_filescan.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
from unittest.mock import patch
2+
3+
from api_app.analyzers_manager.file_analyzers.filescan import FileScanUpload
4+
from tests.mock_utils import MockUpResponse
5+
6+
from .base_test_class import BaseFileAnalyzerTest
7+
8+
9+
class TestFileScanUpload(BaseFileAnalyzerTest):
10+
analyzer_class = FileScanUpload
11+
12+
def get_extra_config(self):
13+
return {"_api_key": "sample_key"}
14+
15+
def get_mocked_response(self):
16+
return [
17+
patch(
18+
"requests.post",
19+
return_value=MockUpResponse({"flow_id": 1}, 200),
20+
),
21+
patch(
22+
"requests.get",
23+
return_value=MockUpResponse(
24+
{
25+
"allFinished": True,
26+
"general": {
27+
"verdict": "clean",
28+
"file_type": "exe",
29+
"file_name": "test_file.exe",
30+
},
31+
"finalVerdict": "no threat detected",
32+
},
33+
200,
34+
),
35+
),
36+
]

‎tests/api_app/analyzers_manager/unit_tests/file_analyzers/test_floss.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
from unittest.mock import patch
2+
3+
from api_app.analyzers_manager.file_analyzers.floss import Floss
4+
5+
from .base_test_class import BaseFileAnalyzerTest
6+
7+
8+
class TestFloss(BaseFileAnalyzerTest):
9+
analyzer_class = Floss
10+
11+
def get_extra_config(self):
12+
return {
13+
"max_no_of_strings": {"decoded": 10, "stack": 5},
14+
"rank_strings": {"decoded": True, "stack": False},
15+
}
16+
17+
def get_mocked_response(self):
18+
return [
19+
patch(
20+
"api_app.analyzers_manager.file_analyzers.floss.Floss._docker_run",
21+
side_effect=[
22+
{
23+
"strings": {
24+
"decoded": ["de_string1", "de_string2"],
25+
"stack": ["st_string1", "st_string2"],
26+
}
27+
},
28+
# second call for ranking decoded strings only
29+
["de_string1", "de_string2"], # simulate ranked strings
30+
],
31+
),
32+
]

‎tests/api_app/analyzers_manager/unit_tests/file_analyzers/test_goresym.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
from unittest.mock import patch
2+
3+
from api_app.analyzers_manager.file_analyzers.goresym import GoReSym
4+
5+
from .base_test_class import BaseFileAnalyzerTest
6+
7+
8+
class TestGoReSym(BaseFileAnalyzerTest):
9+
analyzer_class = GoReSym
10+
11+
def get_extra_config(self):
12+
return {
13+
"default": True,
14+
"paths": False,
15+
"types": False,
16+
"manual": "",
17+
"version": "",
18+
}
19+
20+
def get_mocked_response(self):
21+
return [
22+
patch(
23+
"api_app.analyzers_manager.file_analyzers.goresym.GoReSym._docker_run",
24+
return_value={
25+
"report": {
26+
"Version": "1.22.3",
27+
"OS": "linux",
28+
"BuildInfo": {
29+
"GoVersion": "go1.22.3",
30+
"Path": "github.com/example/project",
31+
},
32+
}
33+
},
34+
)
35+
]

0 commit comments

Comments
 (0)