aes.py CryptoJS-ify patch

aes-cryptojs.patch

--- aes.py	2023-06-06 15:43:19.437326158 +0900
+++ aes-cryptojs.py	2023-06-07 11:07:27.063221638 +0900
@@ -11,6 +11,8 @@
 
 Although this is an exercise, the `encrypt` and `decrypt` functions should
 provide reasonable security to encrypted messages.
+
+Modified to do some "invalid operations" like CryptoJS.
 """
 
 
@@ -178,13 +180,11 @@
     This is a raw implementation of AES, without key stretching or IV
     management. Unless you need that, please use `encrypt` and `decrypt`.
     """
-    rounds_by_key_size = {16: 10, 24: 12, 32: 14}
     def __init__(self, master_key):
         """
         Initializes the object with a given key.
         """
-        assert len(master_key) in AES.rounds_by_key_size
-        self.n_rounds = AES.rounds_by_key_size[len(master_key)]
+        self.n_rounds = 6 + len(master_key) // 4
         self._key_matrices = self._expand_key(master_key)
 
     def _expand_key(self, master_key):
@@ -209,7 +209,7 @@
                 # XOR with first byte of R-CON, since the others bytes of R-CON are 0.
                 word[0] ^= r_con[i]
                 i += 1
-            elif len(master_key) == 32 and len(key_columns) % iteration_size == 4:
+            elif len(master_key) >= 32 and len(key_columns) % iteration_size == 4:
                 # Run word through S-box in the fourth iteration when using a
                 # 256-bit key.
                 word = [s_box[b] for b in word]