Tests: Remove hash_file() and use get_file_sha256_hash() instead

This commit is contained in:
Mike Salvatore 2021-07-16 10:21:07 -04:00
parent 8879dae276
commit 1d7476637d
3 changed files with 7 additions and 18 deletions

View File

@ -9,8 +9,8 @@ from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import
TEST_KEYBOARD_TXT_CLEARTEXT_SHA256, TEST_KEYBOARD_TXT_CLEARTEXT_SHA256,
TEST_KEYBOARD_TXT_ENCRYPTED_SHA256, TEST_KEYBOARD_TXT_ENCRYPTED_SHA256,
) )
from tests.utils import hash_file
from common.utils.file_utils import get_file_sha256_hash
from infection_monkey.ransomware.in_place_file_encryptor import InPlaceFileEncryptor from infection_monkey.ransomware.in_place_file_encryptor import InPlaceFileEncryptor
from infection_monkey.utils.bit_manipulators import flip_bits from infection_monkey.utils.bit_manipulators import flip_bits
@ -44,11 +44,11 @@ def test_file_encrypted(
): ):
test_keyboard = ransomware_target / file_name test_keyboard = ransomware_target / file_name
assert hash_file(test_keyboard) == cleartext_hash assert get_file_sha256_hash(test_keyboard) == cleartext_hash
in_place_bitflip_file_encryptor(test_keyboard) in_place_bitflip_file_encryptor(test_keyboard)
assert hash_file(test_keyboard) == encrypted_hash assert get_file_sha256_hash(test_keyboard) == encrypted_hash
def test_file_encrypted_in_place(in_place_bitflip_file_encryptor, ransomware_target): def test_file_encrypted_in_place(in_place_bitflip_file_encryptor, ransomware_target):
@ -70,4 +70,4 @@ def test_encrypted_file_has_new_extension(ransomware_target):
assert not test_keyboard.exists() assert not test_keyboard.exists()
assert encrypted_test_keyboard.exists() assert encrypted_test_keyboard.exists()
assert hash_file(encrypted_test_keyboard) == TEST_KEYBOARD_TXT_ENCRYPTED_SHA256 assert get_file_sha256_hash(encrypted_test_keyboard) == TEST_KEYBOARD_TXT_ENCRYPTED_SHA256

View File

@ -1,6 +1,6 @@
import pytest import pytest
from tests.utils import hash_file
from common.utils.file_utils import get_file_sha256_hash
from infection_monkey.ransomware.readme_dropper import leave_readme from infection_monkey.ransomware.readme_dropper import leave_readme
DEST_FILE = "README.TXT" DEST_FILE = "README.TXT"
@ -23,10 +23,10 @@ def test_readme_already_exists(src_readme, dest_readme):
leave_readme(src_readme, dest_readme) leave_readme(src_readme, dest_readme)
assert hash_file(dest_readme) == EMPTY_FILE_HASH assert get_file_sha256_hash(dest_readme) == EMPTY_FILE_HASH
def test_leave_readme(src_readme, dest_readme): def test_leave_readme(src_readme, dest_readme):
leave_readme(src_readme, dest_readme) leave_readme(src_readme, dest_readme)
assert hash_file(dest_readme) == README_HASH assert get_file_sha256_hash(dest_readme) == README_HASH

View File

@ -1,7 +1,5 @@
import ctypes import ctypes
import hashlib
import os import os
from pathlib import Path
def is_user_admin(): def is_user_admin():
@ -11,14 +9,5 @@ def is_user_admin():
return ctypes.windll.shell32.IsUserAnAdmin() return ctypes.windll.shell32.IsUserAnAdmin()
def hash_file(filepath: Path):
sha256 = hashlib.sha256()
with open(filepath, "rb") as f:
for block in iter(lambda: f.read(65536), b""):
sha256.update(block)
return sha256.hexdigest()
def raise_(ex): def raise_(ex):
raise ex raise ex