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

View File

@ -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)