tests: Move hash_file() into tests/utils.py

This commit is contained in:
Mike Salvatore 2021-06-23 09:01:42 -04:00
parent ab40518881
commit 45ba743418
2 changed files with 12 additions and 11 deletions

View File

@ -1,10 +1,9 @@
import hashlib
import os import os
import shutil import shutil
from pathlib import Path from pathlib import Path
import pytest import pytest
from tests.utils import is_user_admin from tests.utils import hash_file, is_user_admin
from infection_monkey.ransomware.ransomware_payload import EXTENSION, RansomewarePayload from infection_monkey.ransomware.ransomware_payload import EXTENSION, RansomewarePayload
@ -33,15 +32,6 @@ TEST_KEYBOARD_TXT_ENCRYPTED_SHA256 = (
) )
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 with_extension(filename): def with_extension(filename):
return f"{filename}{EXTENSION}" return f"{filename}{EXTENSION}"

View File

@ -1,5 +1,7 @@
import ctypes import ctypes
import hashlib
import os import os
from pathlib import Path
def is_user_admin(): def is_user_admin():
@ -7,3 +9,12 @@ def is_user_admin():
return os.getuid() == 0 return os.getuid() == 0
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()