Complete AES implementation for python3

pwman/util/crypto/AES.py

@@ -308,16 +308,16 @@ def new(key,mode=MODE_ECB,IV=None,counter=None,segment_size=None):
 
 block_size = 16
 
+
 class AES(BlockCipher):
-    
+
     key_error_message = ("Key should be 128, 192 or 256 bits")
     block_size = 16
-    
-    def __init__(self,key,mode,IV,counter,segment_size):
+
+    def __init__(self, key, mode, IV, counter, segment_size):
         cipher_module = rijndael
         self.blocksize = 16
         args = {'block_size':16}
-        
         BlockCipher.__init__(self,key,mode,IV,counter,cipher_module,segment_size,args)
 
     def keylen_valid(self,key):

pwman/util/crypto/blockcipher.py

@@ -23,6 +23,7 @@
 from . import util
 from . import padding
 import collections
+import sys
 
 MODE_ECB = 1
 MODE_CBC = 2
@@ -55,7 +56,6 @@ class BlockCipher():
             self.IV = '\x00'*self.blocksize
         else:
             self.IV = IV
-
         if mode != MODE_XTS:
             self.cipher = cipher_module(self.key,**args)
         if mode == MODE_ECB:
@@ -227,7 +227,10 @@ class ECB:
     """ECB chaining mode
     """
     def __init__(self, codebook, blocksize):
-        self.cache = ''
+        if sys.version_info.major > 2:
+            self.cache = b''
+        else:
+            self.cache = ''
         self.codebook = codebook
         self.blocksize = blocksize
 
@@ -257,7 +260,10 @@ class ECB:
             else:
                 output_blocks.append(self.codebook.decrypt( self.cache[i:i + self.blocksize] ))
         self.cache = self.cache[i+self.blocksize:]
-        return ''.join(output_blocks)
+        if sys.version_info.major < 3:
+            return ''.join(output_blocks)
+        else:
+            return output_blocks[0]
 
 class CBC:
     """CBC chaining mode

pwman/util/crypto/rijndael.py

@@ -53,6 +53,7 @@ If any strings are of the wrong length a ValueError is thrown
 # deleting all the comments and renaming all the variables
 
 import copy
+import sys
 
 
 #-----------------------
@@ -252,9 +253,19 @@ class rijndael:
 
         # copy user material bytes into temporary ints
         tk = []
-        for i in range(0, KC):
-            tk.append((ord(key[i * 4]) << 24) | (ord(key[i * 4 + 1]) << 16) |
-                (ord(key[i * 4 + 2]) << 8) | ord(key[i * 4 + 3]))
+
+        if sys.version_info.major < 3:
+            for i in range(0, KC):
+                tk.append((ord(key[i * 4]) << 24) |
+                          (ord(key[i * 4 + 1]) << 16) |
+                          (ord(key[i * 4 + 2]) << 8) |
+                          ord(key[i * 4 + 3]))
+        else:
+            for i in range(0, KC):
+                tk.append((key[i * 4] << 24) |
+                          (key[i * 4 + 1] << 16) |
+                          (key[i * 4 + 2] << 8) |
+                          (key[i * 4 + 3]))
 
         # copy values into round key arrays
         t = 0
@@ -326,11 +337,18 @@ class rijndael:
         # temporary work array
         t = []
         # plaintext to ints + key
-        for i in range(BC):
-            t.append((ord(plaintext[i * 4    ]) << 24 |
-                      ord(plaintext[i * 4 + 1]) << 16 |
-                      ord(plaintext[i * 4 + 2]) <<  8 |
-                      ord(plaintext[i * 4 + 3])        ) ^ Ke[0][i])
+        if sys.version_info.major < 3:
+            for i in range(BC):
+                t.append((ord(plaintext[i * 4    ]) << 24 |
+                          ord(plaintext[i * 4 + 1]) << 16 |
+                          ord(plaintext[i * 4 + 2]) <<  8 |
+                          ord(plaintext[i * 4 + 3])        ) ^ Ke[0][i])
+        else:
+            for i in range(BC):
+                t.append((plaintext[i * 4    ] << 24 |
+                          plaintext[i * 4 + 1] << 16 |
+                          plaintext[i * 4 + 2] <<  8 |
+                          plaintext[i * 4 + 3]) ^ Ke[0][i])
         # apply round transforms
         for r in range(1, ROUNDS):
             for i in range(BC):
@@ -346,8 +364,11 @@ class rijndael:
             result.append((S[(t[ i           ] >> 24) & 0xFF] ^ (tt >> 24)) & 0xFF)
             result.append((S[(t[(i + s1) % BC] >> 16) & 0xFF] ^ (tt >> 16)) & 0xFF)
             result.append((S[(t[(i + s2) % BC] >>  8) & 0xFF] ^ (tt >>  8)) & 0xFF)
-            result.append((S[ t[(i + s3) % BC]        & 0xFF] ^  tt       ) & 0xFF)
-        return "".join(list(map(chr, result)))
+            result.append((S[ t[(i + s3) % BC]  & 0xFF] ^  tt       ) & 0xFF)
+        if sys.version_info.major < 3:
+            return "".join(list(map(chr, result)))
+        else:
+            return bytearray(result)
 
     def decrypt(self, ciphertext):
         if len(ciphertext) != self.block_size:
@@ -369,11 +390,18 @@ class rijndael:
         # temporary work array
         t = [0] * BC
         # ciphertext to ints + key
-        for i in range(BC):
-            t[i] = (ord(ciphertext[i * 4    ]) << 24 |
-                    ord(ciphertext[i * 4 + 1]) << 16 |
-                    ord(ciphertext[i * 4 + 2]) <<  8 |
-                    ord(ciphertext[i * 4 + 3])        ) ^ Kd[0][i]
+        if sys.version_info.major < 3:
+            for i in range(BC):
+                t[i] = (ord(ciphertext[i * 4    ]) << 24 |
+                        ord(ciphertext[i * 4 + 1]) << 16 |
+                        ord(ciphertext[i * 4 + 2]) <<  8 |
+                        ord(ciphertext[i * 4 + 3])        ) ^ Kd[0][i]
+        else:
+            for i in range(BC):
+                t[i] = (ciphertext[i * 4    ] << 24 |
+                        ciphertext[i * 4 + 1] << 16 |
+                        ciphertext[i * 4 + 2] <<  8 |
+                        ciphertext[i * 4 + 3]        ) ^ Kd[0][i]
         # apply round transforms
         for r in range(1, ROUNDS):
             for i in range(BC):
@@ -390,14 +418,18 @@ class rijndael:
             result.append((Si[(t[(i + s1) % BC] >> 16) & 0xFF] ^ (tt >> 16)) & 0xFF)
             result.append((Si[(t[(i + s2) % BC] >>  8) & 0xFF] ^ (tt >>  8)) & 0xFF)
             result.append((Si[ t[(i + s3) % BC]        & 0xFF] ^  tt       ) & 0xFF)
-        return "".join(list(map(chr, result)))
 
+        if sys.version_info.major < 3:
+            return "".join(list(map(chr, result)))
+        else:
+            return bytearray(result)
 
 def encrypt(key, block):
     return rijndael(key, len(block)).encrypt(block)
 
 
 def decrypt(key, block):
+    import pdb; pdb.set_trace()
     return rijndael(key, len(block)).decrypt(block)
 
 

pwman/util/crypto_engine.py

@@ -30,8 +30,8 @@ import time
 try:
     from Crypto.Cipher import AES
     from Crypto.Protocol.KDF import PBKDF2
-except ImportError:
-    # PyCryptop not found, we use a compatible implementation
+except ImportError as E:
+    # PyCrypto not found, we use a compatible implementation
     # in pure Python.
     # This is good for Windows where software installation suck
     # or embeded devices where compilation is a bit harder