RITSEC CTF 2022 - Crypto - Bad AES

badaes.py

"""
# RITSEC CTF 2022 - Crypto - Bad AES
## Custom AES implementation where Mix Columns and Shift Rows steps switch places

A secret government agency uses a 16-letter passphrase that is encrypted 
to create their passwords for their computers. An insider within the agency 
told me that everyday employees input their passphrase into this secret 
encryption scheme to receive their password for the day & the key used to 
encrypt their passphrase is changed by the agency daily.
(This is so their passwords change every day without the employee having 
to memorize a new  password daily).

When walking past a co-worker, my insider saw a sticky note with their 
passphrase reading,

I_h8t3_w0rK_h3rE

After some further snooping, I’ve discovered their encryption scheme is a 
slight variation on the AES algorithm. There is only one change made to the 
AES algorithm. Within the diffusion layer of the AES cipher, the Mix Column 
Layer is used before the Shift Rows Layer.

Today, my insider came to me with the key that was used to create the passwords 
for today.

S33cReT-aGecy174

I need you to reverse engineer their variation of the AES algorithm and encrypt 
the passphrase to have access to an inside computer without exposing my insider.
Submit the ciphertext in hex within the RS{} brackets.

Hint: Swap Mix Columns and Shift Rows, completely

----------------------------------------------------------------------------

Answer was based on these resources:
https://en.wikipedia.org/wiki/Advanced_Encryption_Standard
https://en.wikipedia.org/wiki/Rijndael_MixColumns
https://en.wikipedia.org/wiki/AES_key_schedule

Code based on / Credit goes to
https://medium.com/wearesinch/building-aes-128-from-the-ground-up-with-python-8122af44ebf9
https://www.samiam.org/key-schedule.html

"""

# Substitution Box
S = [ 0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5, 0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 
0xab, 0x76, 0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0, 0xad, 0xd4, 0xa2, 0xaf, 0x9c, 
0xa4, 0x72, 0xc0, 0xb7, 0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc, 0x34, 0xa5, 0xe5, 0xf1, 
0x71, 0xd8, 0x31, 0x15, 0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05, 0x9a, 0x07, 0x12, 0x80, 
0xe2, 0xeb, 0x27, 0xb2, 0x75, 0x09, 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0, 0x52, 0x3b, 
0xd6, 0xb3, 0x29, 0xe3, 0x2f, 0x84, 0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b, 0x6a, 
0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf, 0xd0, 0xef, 0xaa, 0xfb, 0x43, 0x4d, 0x33, 0x85, 
0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c, 0x9f, 0xa8, 0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 
0xf5, 0xbc, 0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2, 0xcd, 0x0c, 0x13, 0xec, 0x5f, 0x97, 
0x44, 0x17, 0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73, 0x60, 0x81, 0x4f, 0xdc, 0x22, 
0x2a, 0x90, 0x88, 0x46, 0xee, 0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb, 0xe0, 0x32, 0x3a, 0x0a, 
0x49, 0x06, 0x24, 0x5c, 0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79, 0xe7, 0xc8, 0x37, 
0x6d, 0x8d, 0xd5, 0x4e, 0xa9, 0x6c, 0x56, 0xf4, 0xea, 0x65, 0x7a, 0xae, 0x08, 0xba, 0x78, 
0x25, 0x2e, 0x1c, 0xa6, 0xb4, 0xc6, 0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a, 0x70, 
0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e, 0x61, 0x35, 0x57, 0xb9, 0x86, 0xc1, 0x1d, 0x9e, 
0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e, 0x94, 0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 
0xdf, 0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68, 0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 
0xbb, 0x16 ]

def str_to_int_array(str_in):
  return [ord(c) for c in str_in]

assert str_to_int_array('AB') == [65,66]

def sub(to_sub):
  return [S[c] for c in to_sub]

assert sub([0xa5]) == [6]

def rot(arr_4):
  assert isinstance(arr_4, list), 'arr_4 must be a list'
  assert len(arr_4) == 4, 'arr_4 must be of length 4'
  arr_4.append(arr_4.pop(0))
  return arr_4

def shift(to_shift):
  assert isinstance(to_shift, list), 'to_shift must be a list'
  assert len(to_shift) == 16, 'to_shift must be of length 16'
  return [to_shift[0], to_shift[5], to_shift[10], to_shift[15],
         to_shift[4], to_shift[9], to_shift[14], to_shift[3],
         to_shift[8], to_shift[13], to_shift[2], to_shift[7],
         to_shift[12], to_shift[1], to_shift[6], to_shift[11]]

assert shift(str_to_int_array("ABCDEFGHIJKLMNOP")) == str_to_int_array("AFKPEJODINCHMBGL")

def mult_by_2(v):
    s = v << 1
    s &= 0xff
    if (v & 128) != 0:
        s = s ^ 0x1b
    return s

def mult_by_3(v):
    return mult_by_2(v) ^ v

def mix_column(column):
    r = [
        mult_by_2(column[0]) ^ mult_by_3(column[1]) ^ column[2] ^ column[3],
        mult_by_2(column[1]) ^ mult_by_3(column[2]) ^ column[3] ^ column[0],
        mult_by_2(column[2]) ^ mult_by_3(column[3]) ^ column[0] ^ column[1],
        mult_by_2(column[3]) ^ mult_by_3(column[0]) ^ column[1] ^ column[2],
    ]
    return r

def mix(to_mix):
  mixed = []
  for i in range(0,16,4):
    mixed.extend( mix_column(to_mix[i:i+4]) )
  return mixed

assert mix([219,19,83,69,242,10,34,92,1,1,1,1,198,198,198,198]) == [142,77,161,188,159,220,88,157,1,1,1,1,198,198,198,198]

def rcon(i):
  precomputed = [0,1,2,4,8,0x10,0x20,0x40,0x80,0x1b,0x36]
  return precomputed[i]

def schedule_core(key_in, i):
  key_in = sub(rot(key_in))
  key_in[0] = key_in[0] ^ rcon(i)
  return key_in

def expand_key(key_in):
  # c is 16 because the first sub-key is the user-supplied key
  c = 16
  i = 1
  t = [0,0,0,0]
  key_in.extend([0] * 10*16 )
  # We need 11 sets of sixteen bytes each for 128-bit mode
  while c < 176:
    # Copy the temporary variable over from the last 4-byte block
    for a in range(4):
      t[a] = key_in[a + c - 4]
    # Every four blocks (of four bytes) do a complex calculation
    if (c % 16 == 0):
      t = schedule_core(t,i)
      i = i + 1
    for a in range(4):
      key_in[c] = key_in[c - 16] ^ t[a]
      c = c + 1
  return key_in

expanded_example = [0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,
                    0x62,0x63,0x63,0x63,0x62,0x63,0x63,0x63,0x62,0x63,0x63,0x63,0x62,0x63,0x63,0x63,
                    0x9b,0x98,0x98,0xc9,0xf9,0xfb,0xfb,0xaa,0x9b,0x98,0x98,0xc9,0xf9,0xfb,0xfb,0xaa,
                    0x90,0x97,0x34,0x50,0x69,0x6c,0xcf,0xfa,0xf2,0xf4,0x57,0x33,0x0b,0x0f,0xac,0x99,
                    0xee,0x06,0xda,0x7b,0x87,0x6a,0x15,0x81,0x75,0x9e,0x42,0xb2,0x7e,0x91,0xee,0x2b,
                    0x7f,0x2e,0x2b,0x88,0xf8,0x44,0x3e,0x09,0x8d,0xda,0x7c,0xbb,0xf3,0x4b,0x92,0x90,
                    0xec,0x61,0x4b,0x85,0x14,0x25,0x75,0x8c,0x99,0xff,0x09,0x37,0x6a,0xb4,0x9b,0xa7,
                    0x21,0x75,0x17,0x87,0x35,0x50,0x62,0x0b,0xac,0xaf,0x6b,0x3c,0xc6,0x1b,0xf0,0x9b,
                    0x0e,0xf9,0x03,0x33,0x3b,0xa9,0x61,0x38,0x97,0x06,0x0a,0x04,0x51,0x1d,0xfa,0x9f,
                    0xb1,0xd4,0xd8,0xe2,0x8a,0x7d,0xb9,0xda,0x1d,0x7b,0xb3,0xde,0x4c,0x66,0x49,0x41,
                    0xb4,0xef,0x5b,0xcb,0x3e,0x92,0xe2,0x11,0x23,0xe9,0x51,0xcf,0x6f,0x8f,0x18,0x8e]

expanded_example2 = [0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,
                     0xe8,0xe9,0xe9,0xe9,0x17,0x16,0x16,0x16,0xe8,0xe9,0xe9,0xe9,0x17,0x16,0x16,0x16,
                     0xad,0xae,0xae,0x19,0xba,0xb8,0xb8,0x0f,0x52,0x51,0x51,0xe6,0x45,0x47,0x47,0xf0,
                     0x09,0x0e,0x22,0x77,0xb3,0xb6,0x9a,0x78,0xe1,0xe7,0xcb,0x9e,0xa4,0xa0,0x8c,0x6e,
                     0xe1,0x6a,0xbd,0x3e,0x52,0xdc,0x27,0x46,0xb3,0x3b,0xec,0xd8,0x17,0x9b,0x60,0xb6,
                     0xe5,0xba,0xf3,0xce,0xb7,0x66,0xd4,0x88,0x04,0x5d,0x38,0x50,0x13,0xc6,0x58,0xe6,
                     0x71,0xd0,0x7d,0xb3,0xc6,0xb6,0xa9,0x3b,0xc2,0xeb,0x91,0x6b,0xd1,0x2d,0xc9,0x8d,
                     0xe9,0x0d,0x20,0x8d,0x2f,0xbb,0x89,0xb6,0xed,0x50,0x18,0xdd,0x3c,0x7d,0xd1,0x50,
                     0x96,0x33,0x73,0x66,0xb9,0x88,0xfa,0xd0,0x54,0xd8,0xe2,0x0d,0x68,0xa5,0x33,0x5d,
                     0x8b,0xf0,0x3f,0x23,0x32,0x78,0xc5,0xf3,0x66,0xa0,0x27,0xfe,0x0e,0x05,0x14,0xa3,
                     0xd6,0x0a,0x35,0x88,0xe4,0x72,0xf0,0x7b,0x82,0xd2,0xd7,0x85,0x8c,0xd7,0xc3,0x26]

assert expand_key( [0] * 16 ) == expanded_example, f"Key expansion failed.\n\nResult \n{expand_key( [0] * 16 )}\n\nExpected \n{expanded_example}"
assert expand_key( [0xff] * 16 ) == expanded_example2, f"Key expansion failed.\n\nResult \n{expand_key( [0xff] * 16 )}\n\nExpected \n{expanded_example2}"


def key_for_round(key, round):
  return key[round*16:round*16+16]

def add_key(data, key):
  return [ c ^ k for c, k in zip(data, key) ]

def aes_enc(data, key):
    # Asumes data comes padded to 16 chars

    # From now on it's all integer arrays
    data = str_to_int_array(data)
    key = str_to_int_array(key)

    expanded_key = expand_key(key)

    # Apply the original key to the blocks before start the rounds
    round_key = key_for_round(expanded_key, 0)
    temp_data = []
    for i in range(0,len(data),16):
      block = data[i:i+16]
      temp_data.extend( add_key(block, round_key) )
    data = temp_data

    # Now we can move to the main part of the algorithm
    for r in range(1, 10):
      temp_data = []
      round_key = key_for_round(expanded_key, r)

      for i in range(0,len(data),16):
        block = data[i:i+16]
        sub_step = sub(block)
        shift_step = shift(sub_step)
        mix_step = mix(shift_step)
        add_step = add_key(mix_step,round_key)
        temp_data.extend(add_step)
      
      data = temp_data
  

    # A final round without the mix columns
    temp_data = []
    round_key = key_for_round(expanded_key, 10)

    for i in range(0,len(data),16):
      block = data[i:i+16]
      temp_data.extend(
        add_key(
          shift(
            sub(block)
          ),
          round_key
        )
      )

    data = temp_data

    return bytes(data)

plaintext = 'I_h8t3_w0rK_h3rE'
key = 'S33cReT-aGecy174'

ciphertext = aes_enc(plaintext, key)

assert ciphertext.hex() == '9b6791616f69abdce161b400dc213953'


def bad_aes_enc(data, key):
    # Asumes data comes padded to 16 chars

    # From now on it's all integer arrays
    data = str_to_int_array(data)
    key = str_to_int_array(key)

    expanded_key = expand_key(key)

    # Apply the original key to the blocks before start the rounds
    round_key = key_for_round(expanded_key, 0)
    temp_data = []
    for i in range(0,len(data),16):
      block = data[i:i+16]
      temp_data.extend( add_key(block, round_key) )
    data = temp_data

    # Now we can move to the main part of the algorithm
    for r in range(1, 10):
      temp_data = []
      round_key = key_for_round(expanded_key, r)

      for i in range(0,len(data),16):
        block = data[i:i+16]
        sub_step = sub(block)
        mix_step = mix(sub_step)
        shift_step = shift(mix_step)
        add_step = add_key(shift_step,round_key)
        temp_data.extend(add_step)
      
      data = temp_data
  

    # A final round without the mix columns
    temp_data = []
    round_key = key_for_round(expanded_key, 10)

    for i in range(0,len(data),16):
      block = data[i:i+16]
      temp_data.extend(
        add_key(
          mix(
            sub(block)
          ),
          round_key
        )
      )

    data = temp_data

    return bytes(data)

ciphertext = bad_aes_enc(plaintext, key)

# :D 
assert "RS{" + ciphertext.hex() + "}" == "RS{f0df90ef84d71bf2dfc72f3ba1713cec}"