Browse Source

Ensure "secure" files are encrypted

pull/128/head
Adam Johnson 6 years ago
parent
commit
1bee0b1736
Signed by: Hoikas
GPG Key ID: 0B6515D6FF6F271E
  1. 44
      korman/exporter/outfile.py

44
korman/exporter/outfile.py

@ -27,6 +27,7 @@ import time
import weakref
import zipfile
_CHUNK_SIZE = 0xA00000
_encoding = locale.getpreferredencoding(False)
def _hashfile(filename, hasher, block=0xFFFF):
@ -92,6 +93,42 @@ class _OutputFile:
if self.file_data is None:
self.file_data = self.id_data.as_string()
# Last chance for encryption...
enc = kwargs.get("enc", None)
if enc is not None:
self._encrypt(enc)
def _encrypt(self, enc):
backing_stream = hsRAMStream()
with plEncryptedStream().open(backing_stream, fmCreate, enc) as enc_stream:
if self.file_path:
if plEncryptedStream.IsFileEncrypted(self.file_path):
with plEncryptedStream().open(self.file_path, fmRead, plEncryptedStream.kEncAuto) as dec_stream:
self._enc_spin_wash(enc_stream, dec_stream)
else:
with hsFileStream().open(self.file_path, fmRead) as dec_stream:
self._enc_spin_wash(enc_stream, dec_stream)
elif self.file_data:
if isinstance(self.file_data, str):
enc_stream.write(self.file_data.encode(_encoding))
else:
enc_stream.write(self.file_data)
else:
raise RuntimeError()
self.file_data = backing_stream.buffer
# do NOT copy over an unencrypted file...
self.file_path = None
def _enc_spin_wash(self, enc_stream, dec_stream):
while True:
size_rem = dec_stream.size - dec_stream.pos
readsz = min(size_rem, _CHUNK_SIZE)
if readsz == 0:
break
data = dec_stream.read(readsz)
enc_stream.write(data)
def __eq__(self, rhs):
return str(self) == str(rhs)
@ -102,10 +139,10 @@ class _OutputFile:
if self.file_path:
with open(self.file_path, "rb") as handle:
h = md5()
data = handle.read(0xFFFF)
data = handle.read(_CHUNK_SIZE)
while data:
h.update(data)
data = handle.read(0xFFFF)
data = handle.read(_CHUNK_SIZE)
return h.digest()
elif self.file_data is not None:
if isinstance(self.file_data, str):
@ -196,6 +233,9 @@ class OutputFiles:
if not stream is backing_stream:
backing_stream.close()
# Not passing enc as a keyword argument to the output file definition. It makes more
# sense to yield an encrypted stream from this context manager and encrypt as we go
# instead of doing lots of buffer copying to encrypt as a post step.
dirname = kwargs.get("dirname", "dat")
kwargs = {
"file_type": _FileType.generated_dat if dirname == "dat" else

Loading…
Cancel
Save