From 4400d8d1d463a38d013ffffb5b6a403445dfc2dd Mon Sep 17 00:00:00 2001 From: Conrad Meyer Date: Thu, 5 Apr 2018 10:32:59 -0700 Subject: [PATCH] kfxdrm: Traipse through the kfx-zip more efficiently We only need to read the magic bytes headers to identify files of the correct type. Avoid slurping the entire contents out of the zip if it's the wrong file. --- DeDRM_calibre_plugin/DeDRM_plugin/kfxdedrm.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/DeDRM_calibre_plugin/DeDRM_plugin/kfxdedrm.py b/DeDRM_calibre_plugin/DeDRM_plugin/kfxdedrm.py index 8885e09..66317ad 100644 --- a/DeDRM_calibre_plugin/DeDRM_plugin/kfxdedrm.py +++ b/DeDRM_calibre_plugin/DeDRM_plugin/kfxdedrm.py @@ -36,8 +36,11 @@ class KFXZipBook: def processBook(self, totalpids): with zipfile.ZipFile(self.infile, 'r') as zf: for filename in zf.namelist(): - data = zf.read(filename) - if data.startswith('\xeaDRMION\xee'): + with zf.open(filename) as fh: + data = fh.read(8) + if data != '\xeaDRMION\xee': + continue + data += fh.read() if self.voucher is None: self.decrypt_voucher(totalpids) print u'Decrypting KFX DRMION: {0}'.format(filename) @@ -51,9 +54,13 @@ class KFXZipBook: def decrypt_voucher(self, totalpids): with zipfile.ZipFile(self.infile, 'r') as zf: for info in zf.infolist(): - if info.file_size < 0x10000: - data = zf.read(info.filename) - if data.startswith('\xe0\x01\x00\xea') and 'ProtectedData' in data: + with zf.open(info.filename) as fh: + data = fh.read(4) + if data != '\xe0\x01\x00\xea': + continue + + data += fh.read() + if 'ProtectedData' in data: break # found DRM voucher else: raise Exception(u'The .kfx-zip archive contains an encrypted DRMION file without a DRM voucher')