Fix python2 issues in kgenpids and kindlekey

CHANGELOG.md

@@ -99,4 +99,6 @@ This is v10.0.9, a release candidate for v10.1.0. I don't expect there to be maj
 
 ## Fixes on master (not yet released):
 
-- (none)
+- 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).

DeDRM_plugin/ineptpdf.py

@@ -1366,14 +1366,14 @@ class PDFDocument(object):
 
     def process_with_aes(self, key, encrypt, data, repetitions = 1, iv = None):
         if iv is None:
-            keylen = len(key)
-            iv = bytes([0x00]*keylen)
+            iv = bytes(bytearray(16))
+
+        aes = AES.new(key, AES.MODE_CBC, iv)
 
         if not encrypt:
-            plaintext = AES.new(key,AES.MODE_CBC,iv, True).decrypt(data)
+            plaintext = aes.decrypt(data)
             return plaintext
         else:
-            aes = AES.new(key, AES.MODE_CBC, iv, False)
             new_data = bytes(data * repetitions)
             crypt = aes.encrypt(new_data)
             return crypt
@@ -1394,10 +1394,18 @@ class PDFDocument(object):
                     raise Exception("K1 < 32 ...")
                 #def process_with_aes(self, key: bytes, encrypt: bool, data: bytes, repetitions: int = 1, iv: bytes = None):
                 E = self.process_with_aes(K[:16], True, K1, 64, K[16:32])
-                K = (hashlib.sha256, hashlib.sha384, hashlib.sha512)[sum(E) % 3](E).digest()
+                E = bytearray(E)
+
+                E_mod_3 = 0
+                for i in range(16):
+                    E_mod_3 += E[i]
+
+                E_mod_3 %= 3
+
+                K = (hashlib.sha256, hashlib.sha384, hashlib.sha512)[E_mod_3](E).digest()
 
                 if round_number >= 64:
-                    ch = int.from_bytes(E[-1:], "big", signed=False)
+                    ch = E[-1:][0]  # get last byte
                     if ch <= round_number - 32:
                         done = True
 
@@ -1478,14 +1486,23 @@ class PDFDocument(object):
             EncMetadata = b'True'
         if (EncMetadata == ('False' or 'false') or V < 4) and R >= 4:
             hash.update(codecs.decode(b'ffffffff','hex'))
+
+        # Finish hash:
+        hash = hash.digest()
+        
         if R >= 3:
             # 8
             for _ in range(50):
-                hash = hashlib.md5(hash.digest()[:length//8])
-        key = hash.digest()[:length//8]
+                hash = hashlib.md5(hash[:length//8]).digest()
+        if R == 2:
+            # R=2 only uses first five bytes.
+            key = hash[:5]
+        else:
+            key = hash[:length//8]
+
         if R == 2:
             # Algorithm 3.4
-            u1 = ARC4.new(key).decrypt(password)
+            u1 = ARC4.new(key).decrypt(self.PASSWORD_PADDING)
         elif R >= 3:
             # Algorithm 3.5
             hash = hashlib.md5(self.PASSWORD_PADDING) # 2
@@ -1498,6 +1515,7 @@ class PDFDocument(object):
                     k = b''.join(bytes([c ^ i]) for c in key )
                 x = ARC4.new(k).decrypt(x)
             u1 = x+x # 32bytes total
+        
         if R == 2:
             is_authenticated = (u1 == U)
         else:

DeDRM_plugin/kgenpids.py

@@ -53,11 +53,17 @@ def SHA1(message):
 def encode(data, map):
     result = b''
     for char in data:
-        value = char
+        if sys.version_info[0] == 2:
+            value = ord(char)
+        else:
+            value = char
+
         Q = (value ^ 0x80) // len(map)
         R = value % len(map)
-        result += bytes([map[Q]])
-        result += bytes([map[R]])
+
+        result += bytes(bytearray([map[Q]]))
+        result += bytes(bytearray([map[R]]))
+
     return result
 
 # Hash the bytes in data and then encode the digest with the characters in map
@@ -84,8 +90,11 @@ def decode(data,map):
 def getTwoBitsFromBitField(bitField,offset):
     byteNumber = offset // 4
     bitPosition = 6 - 2*(offset % 4)
-    return bitField[byteNumber] >> bitPosition & 3
-
+    if sys.version_info[0] == 2:
+        return ord(bitField[byteNumber]) >> bitPosition & 3
+    else:
+        return bitField[byteNumber] >> bitPosition & 3
+    
 # Returns the six bits at offset from a bit field
 def getSixBitsFromBitField(bitField,offset):
     offset *= 3
@@ -97,7 +106,8 @@ def encodePID(hash):
     global charMap3
     PID = b''
     for position in range (0,8):
-        PID += bytes([charMap3[getSixBitsFromBitField(hash,position)]])
+        PID += bytes(bytearray([charMap3[getSixBitsFromBitField(hash,position)]]))
+        
     return PID
 
 # Encryption table used to generate the device PID
@@ -134,7 +144,7 @@ def generateDevicePID(table,dsn,nbRoll):
         index = (index+1) %8
     for counter in range (0,8):
         index = ((((pid[counter] >>5) & 3) ^ pid[counter]) & 0x1f) + (pid[counter] >> 7)
-        pidAscii += bytes([charMap4[index]])
+        pidAscii += bytes(bytearray([charMap4[index]]))
     return pidAscii
 
 def crc32(s):
@@ -150,7 +160,7 @@ def checksumPid(s):
     for i in (0,1):
         b = crc & 0xff
         pos = (b // l) ^ (b % l)
-        res += bytes([charMap4[pos%l]])
+        res += bytes(bytearray([charMap4[pos%l]]))
         crc >>= 8
     return res
 
@@ -161,14 +171,17 @@ def pidFromSerial(s, l):
     crc = crc32(s)
     arr1 = [0]*l
     for i in range(len(s)):
-        arr1[i%l] ^= s[i]
+        if sys.version_info[0] == 2:
+            arr1[i%l] ^= ord(s[i])
+        else: 
+            arr1[i%l] ^= s[i]
     crc_bytes = [crc >> 24 & 0xff, crc >> 16 & 0xff, crc >> 8 & 0xff, crc & 0xff]
     for i in range(l):
         arr1[i] ^= crc_bytes[i&3]
     pid = b""
     for i in range(l):
         b = arr1[i] & 0xff
-        pid += bytes([charMap4[(b >> 7) + ((b >> 5 & 3) ^ (b & 0x1f))]])
+        pid += bytes(bytearray([charMap4[(b >> 7) + ((b >> 5 & 3) ^ (b & 0x1f))]]))
     return pid
 
 

DeDRM_plugin/kindlekey.py

@@ -115,11 +115,17 @@ def primes(n):
 def encode(data, map):
     result = b''
     for char in data:
-        value = char
+        if sys.version_info[0] == 2:
+            value = ord(char)
+        else:
+            value = char
+
         Q = (value ^ 0x80) // len(map)
         R = value % len(map)
-        result += bytes([map[Q]])
-        result += bytes([map[R]])
+
+        result += bytes(bytearray([map[Q]]))
+        result += bytes(bytearray([map[R]]))
+
     return result
 
 # Hash the bytes in data and then encode the digest with the characters in map