57 lines
1.9 KiB
Python
57 lines
1.9 KiB
Python
import unittest
|
|
import pickle
|
|
import tempfile
|
|
from detection.pickle_detection import pickleScanner, pickleDataDetection
|
|
from unittest.mock import patch
|
|
|
|
|
|
class TestPickleScanner(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
# Create temporary files with valid and malicious data
|
|
self.valid_data = {"key": "value"}
|
|
self.malicious_data = b"\x80\x03csubprocess\ncheck_output\nq\x00X\x05\x00\x00\x00echo 1q\x01\x85q\x02Rq\x03."
|
|
|
|
self.valid_file = tempfile.NamedTemporaryFile(delete=False)
|
|
self.valid_file.write(pickle.dumps(self.valid_data))
|
|
self.valid_file.close()
|
|
|
|
self.malicious_file = tempfile.NamedTemporaryFile(delete=False)
|
|
self.malicious_file.write(self.malicious_data)
|
|
self.malicious_file.close()
|
|
|
|
def tearDown(self):
|
|
# Clean up temporary files
|
|
import os
|
|
|
|
os.remove(self.valid_file.name)
|
|
os.remove(self.malicious_file.name)
|
|
|
|
def test_valid_pickle(self):
|
|
with open(self.valid_file.name, "rb") as file:
|
|
scanner = pickleScanner(file)
|
|
print(scanner.maliciousModule)
|
|
scanner.load()
|
|
output = scanner.output()
|
|
self.assertEqual(output["ReduceCount"], 0)
|
|
self.assertEqual(output["maliciousModule"], [])
|
|
|
|
def test_malicious_pickle(self):
|
|
with open(self.malicious_file.name, "rb") as file:
|
|
scanner = pickleScanner(file)
|
|
scanner.load()
|
|
output = scanner.output()
|
|
self.assertEqual(output["ReduceCount"], 1)
|
|
self.assertIn(("subprocess", "check_output"), output["maliciousModule"])
|
|
|
|
@patch("builtins.print")
|
|
def test_pickleDataDetection_no_output_file(self, mock_print):
|
|
# test output to stdout if filename is not given
|
|
with patch("builtins.print") as mock_print:
|
|
pickleDataDetection(self.valid_file.name)
|
|
mock_print.assert_called_once()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|