Improve key detection

DeDRM_plugin/__init__.py

@@ -317,7 +317,41 @@ class DeDRM(FileTypePlugin):
         import calibre_plugins.dedrm.ineptepub as ineptepub
 
         if ineptepub.adeptBook(inf.name):
-            print("{0} v{1}: {2} is a secure Adobe Adept ePub".format(PLUGIN_NAME, PLUGIN_VERSION, os.path.basename(path_to_ebook)))
+            book_uuid = None
+            try: 
+                # This tries to figure out which Adobe account UUID the book is licensed for. 
+                # If we know that we can directly use the correct key instead of having to
+                # try them all.
+                book_uuid = ineptepub.adeptGetUserUUID(inf.name)
+            except: 
+                pass
+
+            if book_uuid is None: 
+                print("{0} v{1}: {2} is a secure Adobe Adept ePub".format(PLUGIN_NAME, PLUGIN_VERSION, os.path.basename(path_to_ebook)))
+            else: 
+                print("{0} v{1}: {2} is a secure Adobe Adept ePub for UUID {3}".format(PLUGIN_NAME, PLUGIN_VERSION, os.path.basename(path_to_ebook), book_uuid))
+
+
+            if book_uuid is not None: 
+                # Check if we have a key with that UUID in its name: 
+                for keyname, userkeyhex in dedrmprefs['adeptkeys'].items():
+                    if not book_uuid.lower() in keyname.lower(): 
+                        continue
+
+                    # Found matching key
+                    userkey = codecs.decode(userkeyhex, 'hex')
+                    print("{0} v{1}: Trying UUID-matched encryption key {2:s}".format(PLUGIN_NAME, PLUGIN_VERSION, keyname))
+                    of = self.temporary_file(".epub")
+                    try: 
+                        result = ineptepub.decryptBook(userkey, inf.name, of.name)
+                        of.close()
+                        if result == 0:
+                            print("{0} v{1}: Decrypted with key {2:s} after {3:.1f} seconds".format(PLUGIN_NAME, PLUGIN_VERSION,keyname,time.time()-self.starttime))
+                            return of.name
+                    except:
+                        print("{0} v{1}: Exception when decrypting after {2:.1f} seconds - trying other keys".format(PLUGIN_NAME, PLUGIN_VERSION, time.time()-self.starttime))
+                        traceback.print_exc()
+
 
             # Attempt to decrypt epub with each encryption key (generated or provided).
             for keyname, userkeyhex in dedrmprefs['adeptkeys'].items():
@@ -363,7 +397,10 @@ class DeDRM(FileTypePlugin):
                     scriptpath = os.path.join(self.alfdir,"adobekey.py")
                     defaultkeys, defaultnames = WineGetKeys(scriptpath, ".der",dedrmprefs['adobewineprefix'])
 
-                self.default_key = defaultkeys[0]
+                try: 
+                    self.default_key = defaultkeys[0]
+                except: 
+                    print("{0} v{1}: No ADE key found".format(PLUGIN_NAME, PLUGIN_VERSION))
             except:
                 print("{0} v{1}: Exception when getting default Adobe Key after {2:.1f} seconds".format(PLUGIN_NAME, PLUGIN_VERSION, time.time()-self.starttime))
                 traceback.print_exc()
@@ -470,7 +507,10 @@ class DeDRM(FileTypePlugin):
                 scriptpath = os.path.join(self.alfdir,"adobekey.py")
                 defaultkeys, defaultnames = WineGetKeys(scriptpath, ".der",dedrmprefs['adobewineprefix'])
 
-            self.default_key = defaultkeys[0]
+            try:
+                self.default_key = defaultkeys[0]
+            except: 
+                print("{0} v{1}: No ADE key found".format(PLUGIN_NAME, PLUGIN_VERSION))
         except:
             print("{0} v{1}: Exception when getting default Adobe Key after {2:.1f} seconds".format(PLUGIN_NAME, PLUGIN_VERSION, time.time()-self.starttime))
             traceback.print_exc()

DeDRM_plugin/ineptepub.py

@@ -393,6 +393,42 @@ def adeptBook(inpath):
             return True
     return False
 
+# Checks the license file and returns the UUID the book is licensed for. 
+# This is used so that the Calibre plugin can pick the correct decryption key
+# first try without having to loop through all possible keys.
+def adeptGetUserUUID(inpath): 
+    with closing(ZipFile(open(inpath, 'rb'))) as inf:
+        try:
+            rights = etree.fromstring(inf.read('META-INF/rights.xml'))
+            adept = lambda tag: '{%s}%s' % (NSMAP['adept'], tag)
+            expr = './/%s' % (adept('user'),)
+            user_uuid = ''.join(rights.findtext(expr))
+            if user_uuid[:9] != "urn:uuid:": 
+                return None
+            return user_uuid[9:]
+        except:
+            return None
+
+def verify_book_key(bookkey):
+    if bookkey[-17] != '\x00' and bookkey[-17] != 0:
+        # Byte not null, invalid result
+        return False
+
+    if ((bookkey[0] != '\x02' and bookkey[0] != 2) and
+        ((bookkey[0] != '\x00' and bookkey[0] != 0) or 
+        (bookkey[1] != '\x02' and bookkey[1] != 2))):
+        # Key not starting with "00 02" or "02" -> error
+        return False
+
+    keylen = len(bookkey) - 17
+    for i in range(1, keylen):
+        if bookkey[i] == 0 or bookkey[i] == '\x00':
+            # Padding data contains a space - that's not allowed. 
+            # Probably bad decryption.
+            return False
+
+    return True
+
 def decryptBook(userkey, inpath, outpath):
     if AES is None:
         raise ADEPTError("PyCrypto or OpenSSL must be installed.")
@@ -416,7 +452,7 @@ def decryptBook(userkey, inpath, outpath):
             bookkey = rsa.decrypt(codecs.decode(bookkey.encode('ascii'), 'base64'))
             # Padded as per RSAES-PKCS1-v1_5
             if len(bookkey) > 16:
-                if bookkey[-17] == '\x00' or bookkey[-17] == 0:
+                if verify_book_key(bookkey):
                     bookkey = bookkey[-16:]
                 else:
                     print("Could not decrypt {0:s}. Wrong key".format(os.path.basename(inpath)))

DeDRM_plugin/ineptpdf.py

@@ -1587,6 +1587,26 @@ class PDFDocument(object):
         self.ready = True
         return
 
+    def verify_book_key(self, bookkey):
+        if bookkey[-17] != '\x00' and bookkey[-17] != 0:
+            # Byte not null, invalid result
+            return False
+
+        if ((bookkey[0] != '\x02' and bookkey[0] != 2) and
+            ((bookkey[0] != '\x00' and bookkey[0] != 0) or 
+            (bookkey[1] != '\x02' and bookkey[1] != 2))):
+            # Key not starting with "00 02" or "02" -> error
+            return False
+
+        keylen = len(bookkey) - 17
+        for i in range(1, keylen):
+            if bookkey[i] == 0 or bookkey[i] == '\x00':
+                # Padding data contains a space - that's not allowed. 
+                # Probably bad decryption.
+                return False
+
+        return True
+
     def initialize_ebx(self, password, docid, param):
         self.is_printable = self.is_modifiable = self.is_extractable = True
         rsa = RSA(password)
@@ -1597,12 +1617,14 @@ class PDFDocument(object):
         expr = './/{http://ns.adobe.com/adept}encryptedKey'
         bookkey = codecs.decode(''.join(rights.findtext(expr)).encode('utf-8'),'base64')
         bookkey = rsa.decrypt(bookkey)
-        #if bookkey[0] != 2:
-        #    raise ADEPTError('error decrypting book session key')
+
         if len(bookkey) > 16:
-            if bookkey[-17] == '\x00' or bookkey[-17] == 0:
+            if (self.verify_book_key(bookkey)):
                 bookkey = bookkey[-16:]
                 length = 16
+            else:
+                raise ADEPTError('error decrypting book session key')
+
         ebx_V = int_value(param.get('V', 4))
         ebx_type = int_value(param.get('EBX_ENCRYPTIONTYPE', 6))
         # added because of improper booktype / decryption book session key errors