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.

diff --git a/DeDRM_calibre_plugin/DeDRM_plugin/kfxdedrm.py b/DeDRM_calibre_plugin/DeDRM_plugin/kfxdedrm.py
index 8885e097e12f7e0a8d7136a4599534fbdffcebcc..66317adbc184db6c6feb4c5533d05bc0a44c5f43 100644 (file)
--- 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')