agent: Add unit test for RansomwareBitflipEncryptor

This commit is contained in:
Mike Salvatore 2021-06-23 09:19:46 -04:00
parent 707b40608a
commit f1a365def2
1 changed files with 30 additions and 0 deletions

View File

@ -0,0 +1,30 @@
from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import (
ALL_ZEROS_PDF,
ALL_ZEROS_PDF_CLEARTEXT_SHA256,
ALL_ZEROS_PDF_ENCRYPTED_SHA256,
TEST_KEYBOARD_TXT,
TEST_KEYBOARD_TXT_CLEARTEXT_SHA256,
TEST_KEYBOARD_TXT_ENCRYPTED_SHA256,
)
from tests.utils import hash_file
from infection_monkey.ransomware.ransomware_bitflip_encryptor import RansomwareBitflipEncryptor
EXTENSION = ".new"
def with_extension(filename):
return f"{filename}{EXTENSION}"
def test_listed_files_encrypted(ransomware_target):
file_list = [ransomware_target / ALL_ZEROS_PDF, ransomware_target / TEST_KEYBOARD_TXT]
assert hash_file(file_list[0]) == ALL_ZEROS_PDF_CLEARTEXT_SHA256
assert hash_file(file_list[1]) == TEST_KEYBOARD_TXT_CLEARTEXT_SHA256
encryptor = RansomwareBitflipEncryptor(".new")
encryptor.encrypt_files(file_list)
assert hash_file(with_extension(file_list[0])) == ALL_ZEROS_PDF_ENCRYPTED_SHA256
assert hash_file(with_extension(file_list[1])) == TEST_KEYBOARD_TXT_ENCRYPTED_SHA256