code: def key_expand(key, Nk=4):
      
      assert Nk in {4, 6, 8}
      Nr = Nk + 6
  
      key = key.reshape(Nk, 32)
      rkey = exprzeros(4*(Nr+1), 32)
  
      for i in range(Nk):
          rkey[i] = key[i]
  
      for i in range(Nk, 4*(Nr+1)):
          if i % Nk == 0:
              rkey[i] = rkey[i-Nk] ^ subword(rotword(rkey[i-1])) ^ RCON[i//Nk].zext(32-8)
          elif Nk > 6 and i % Nk == 4:
              rkey[i] = rkey[i-Nk] ^ subword(rkey[i-1])
          else:
              rkey[i] = rkey[i-Nk] ^ rkey[i-1]
  
      return rkey
docs: Expand the key into the round key.