More Python2 bugfixes

7 files changed, 75 insertions, 19 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b5537cd..c10dc19 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -101,4 +101,5 @@ This is v10.0.9, a release candidate for v10.1.0. I don't expect there to be maj
 
 - Fix a bug where decrypting a 40-bit RC4 pdf with R=2 didn't work.
 - Fix a bug where decrypting a 256-bit AES pdf with V=5 didn't work.
-- Fix bugs in kgenpids.py and kindlekey.py that caused it to fail on Python 2 (#380).
+- Fix bugs in kgenpids.py, alfcrypto.py, mobidedrm.py and kindlekey.py that caused it to fail on Python 2 (#380).
+- Fix some bugs (Python 2 and Python 3) in erdr2pml.py (untested).
diff --git a/DeDRM_plugin/__init__.py b/DeDRM_plugin/__init__.py
index 9098817..983ecc8 100644
--- a/DeDRM_plugin/__init__.py
+++ b/DeDRM_plugin/__init__.py
@@ -96,7 +96,10 @@ import traceback
 #@@CALIBRE_COMPAT_CODE@@
 
 try: 
-    import __version
+    try: 
+        from . import __version
+    except:
+        import __version
 except: 
     print("#############################")
     print("Failed to load the DeDRM plugin")
@@ -134,8 +137,10 @@ try:
 except:
     config_dir = ""
 
-
-import utilities
+try: 
+    from . import utilities
+except: 
+    import utilities
 
 
 PLUGIN_NAME = __version.PLUGIN_NAME
@@ -915,6 +920,9 @@ class DeDRM(FileTypePlugin):
             # perhaps we need to get a new default Kindle for Mac/PC key
             defaultkeys = []
             print("{0} v{1}: Failed to decrypt with error: {2}".format(PLUGIN_NAME, PLUGIN_VERSION,e.args[0]))
+
+            traceback.print_exc()
+
             print("{0} v{1}: Looking for new default Kindle Key after {2:.1f} seconds".format(PLUGIN_NAME, PLUGIN_VERSION, time.time()-self.starttime))
 
             try:
diff --git a/DeDRM_plugin/alfcrypto.py b/DeDRM_plugin/alfcrypto.py
index ecb7916..3b9f319 100644
--- a/DeDRM_plugin/alfcrypto.py
+++ b/DeDRM_plugin/alfcrypto.py
@@ -8,6 +8,7 @@
 # pbkdf2.py Copyright © 2009 Daniel Holth <[email protected]>
 # pbkdf2.py This code may be freely used and modified for any purpose.
 
+import sys
 import hmac
 from struct import pack
 import hashlib
@@ -25,7 +26,10 @@ class Pukall_Cipher(object):
             raise Exception("PC1: Bad key length")
         wkey = []
         for i in range(8):
-            wkey.append(key[i*2]<<8 | key[i*2+1])
+            if sys.version_info[0] == 2:
+                wkey.append(ord(key[i*2])<<8 | ord(key[i*2+1]))
+            else: 
+                wkey.append(key[i*2]<<8 | key[i*2+1])
         dst = bytearray(len(src))
         for i in range(len(src)):
             temp1 = 0;
@@ -37,7 +41,12 @@ class Pukall_Cipher(object):
                 sum2  = (sum2+sum1)&0xFFFF
                 temp1 = (temp1*20021+1)&0xFFFF
                 byteXorVal ^= temp1 ^ sum2
-            curByte = src[i]
+
+            if sys.version_info[0] == 2:
+                curByte = ord(src[i])
+            else:
+                curByte = src[i]
+
             if not decryption:
                 keyXorVal = curByte * 257;
             curByte = ((curByte ^ (byteXorVal >> 8)) ^ byteXorVal) & 0xFF
@@ -45,7 +54,12 @@ class Pukall_Cipher(object):
                 keyXorVal = curByte * 257;
             for j in range(8):
                 wkey[j] ^= keyXorVal;
-            dst[i] = curByte
+            
+            if sys.version_info[0] == 2:
+                dst[i] = chr(curByte)
+            else: 
+                dst[i] = curByte
+                
         return bytes(dst)
 
 class Topaz_Cipher(object):
@@ -103,7 +117,7 @@ class KeyIVGen(object):
         def xorbytes( a, b ):
             if len(a) != len(b):
                 raise Exception("xorbytes(): lengths differ")
-            return bytes([x ^ y for x, y in zip(a, b)])
+            return bytes(bytearray([x ^ y for x, y in zip(a, b)]))
 
         def prf( h, data ):
             hm = h.copy()
diff --git a/DeDRM_plugin/erdr2pml.py b/DeDRM_plugin/erdr2pml.py
index 1ec9993..7d87aa7 100755
--- a/DeDRM_plugin/erdr2pml.py
+++ b/DeDRM_plugin/erdr2pml.py
@@ -79,12 +79,16 @@ except ImportError:
 
 #@@CALIBRE_COMPAT_CODE@@
 
-from utilities import SafeUnbuffered
+try: 
+    from utilities import SafeUnbuffered
+    from argv_utils import unicode_argv
+except: 
+    from . import utilities, argv_utils
 
 iswindows = sys.platform.startswith('win')
 isosx = sys.platform.startswith('darwin')
 
-from argv_utils import unicode_argv
+
 
 import cgi
 import logging
@@ -141,14 +145,20 @@ def sanitizeFileName(name):
 
 def fixKey(key):
     def fixByte(b):
+        if sys.version_info[0] == 2:
+            b = ord(b)
+
         return b ^ ((b ^ (b<<1) ^ (b<<2) ^ (b<<3) ^ (b<<4) ^ (b<<5) ^ (b<<6) ^ (b<<7) ^ 0x80) & 0x80)
-    return bytes([fixByte(a) for a in key])
+    return bytes(bytearray([fixByte(a) for a in key]))
 
 def deXOR(text, sp, table):
-    r=''
+    r=b''
     j = sp
     for i in range(len(text)):
-        r += chr(ord(table[j]) ^ ord(text[i]))
+        if sys.version_info[0] == 2:
+            r += chr(ord(table[j]) ^ ord(text[i]))
+        else: 
+            r += bytes(bytearray([table[j] ^ text[i]]))
         j = j + 1
         if j == len(table):
             j = 0
diff --git a/DeDRM_plugin/kgenpids.py b/DeDRM_plugin/kgenpids.py
index d11c3da..670ed93 100644
--- a/DeDRM_plugin/kgenpids.py
+++ b/DeDRM_plugin/kgenpids.py
@@ -190,6 +190,10 @@ def getKindlePids(rec209, token, serialnum):
     if isinstance(serialnum,str):
         serialnum = serialnum.encode('utf-8')
 
+    if sys.version_info[0] == 2:
+        if isinstance(serialnum,unicode):
+            serialnum = serialnum.encode('utf-8')
+
     if rec209 is None:
         return [serialnum]
 
diff --git a/DeDRM_plugin/kindlekey.py b/DeDRM_plugin/kindlekey.py
index e8a0e65..6297959 100644
--- a/DeDRM_plugin/kindlekey.py
+++ b/DeDRM_plugin/kindlekey.py
@@ -62,7 +62,12 @@ except NameError:
 
 # Routines common to Mac and PC
 
-from utilities import SafeUnbuffered
+try: 
+    from utilities import SafeUnbuffered
+    from argv_utils import unicode_argv
+except:
+    from . import utilities, argv_utils
+    
 
 try:
     from calibre.constants import iswindows, isosx
@@ -70,7 +75,7 @@ except:
     iswindows = sys.platform.startswith('win')
     isosx = sys.platform.startswith('darwin')
 
-from argv_utils import unicode_argv
+
 
 class DrmException(Exception):
     pass
diff --git a/DeDRM_plugin/mobidedrm.py b/DeDRM_plugin/mobidedrm.py
index 0374836..36d30c4 100755
--- a/DeDRM_plugin/mobidedrm.py
+++ b/DeDRM_plugin/mobidedrm.py
@@ -124,7 +124,11 @@ def getSizeOfTrailingDataEntries(ptr, size, flags):
         if size <= 0:
             return result
         while True:
-            v = ptr[size-1]
+            if sys.version_info[0] == 2:
+                v = ord(ptr[size-1])
+            else:
+                v = ptr[size-1]
+
             result |= (v & 0x7F) << bitpos
             bitpos += 7
             size -= 1
@@ -140,7 +144,10 @@ def getSizeOfTrailingDataEntries(ptr, size, flags):
     # if multibyte data is included in the encryped data, we'll
     # have already cleared this flag.
     if flags & 1:
-        num += (ptr[size - num - 1] & 0x3) + 1
+        if sys.version_info[0] == 2:
+            num += (ord(ptr[size - num - 1]) & 0x3) + 1
+        else: 
+            num += (ptr[size - num - 1] & 0x3) + 1
     return num
 
 
@@ -299,7 +306,10 @@ class MobiBook:
         for pid in pidlist:
             bigpid = pid.encode('utf-8').ljust(16,b'\0')
             temp_key = PC1(keyvec1, bigpid, False)
-            temp_key_sum = sum(temp_key) & 0xff
+            if sys.version_info[0] == 2:
+                temp_key_sum = sum(map(ord,temp_key)) & 0xff
+            else:
+                temp_key_sum = sum(temp_key) & 0xff
             found_key = None
             for i in range(count):
                 verification, size, type, cksum, cookie = struct.unpack('>LLLBxxx32s', data[i*0x30:i*0x30+0x30])
@@ -315,7 +325,11 @@ class MobiBook:
             # Then try the default encoding that doesn't require a PID
             pid = '00000000'
             temp_key = keyvec1
-            temp_key_sum = sum(temp_key) & 0xff
+            if sys.version_info[0] == 2:
+                temp_key_sum = sum(map(ord,temp_key)) & 0xff
+            else:
+                temp_key_sum = sum(temp_key) & 0xff
+            
             for i in range(count):
                 verification, size, type, cksum, cookie = struct.unpack('>LLLBxxx32s', data[i*0x30:i*0x30+0x30])
                 if cksum == temp_key_sum: