// File: permutationlock[P]aes_math/lib/aes.cpp
// Function: expandKey(std::vector<unsigned char, std::allocator<unsigned char>>, int)
vector<QSMatrix<GaloisPolynomial>> expandKey(vector<unsigned char> key, int rounds){
    vector<GaloisPolynomial> keyBytes;
    keyBytes.reserve(16*(rounds+1));    // Reserve space so inserts are faster
    
    // Just copy in key for first 16 bytes
    for(int i=0; i<16; i++){
        keyBytes.push_back(extractPoly(key));
    }
    
    int iteration = 1;
    
    // Continue until we have a 128 bit key for each round + an initial key
    while(keyBytes.size()<16*(rounds+1)){
        for(int w=0; w<4; w++){
            // Copy last 4 bytes
            vector<GaloisPolynomial> word;
            for(int i=0; i<4; i++){
                word.push_back(keyBytes[keyBytes.size()-5+i]);
            }
            
            // Apply core operations if first word
            if(w==0) keyExpandCore(word, iteration);
            
            // Add each byte with that 16 bytes back
            for(int i=0; i<4; i++){
                word[i] += keyBytes[keyBytes.size()-16+i];
            }
            
            // Append
            keyBytes.insert(keyBytes.end(), word.begin(), word.end());
        }
        
        iteration += 1;
    }
    
    // Convert to vector of matrices
    vector<QSMatrix<GaloisPolynomial>> keyMatrices;
    
    for(int i=0; i<rounds+1; i++){
        vector<GaloisPolynomial> v;
        for(int b=0; b<16; b++){
            v.push_back(keyBytes[i*16+b]);
        }
        keyMatrices.push_back(QSMatrix<GaloisPolynomial>(4, 4, v));
    }
    
    return keyMatrices;
}

// Assembly:
subq $0x158, %rsp            # imm = 0x158
movq %rsi, 0x38(%rsp)
movq %rdi, 0x40(%rsp)
movq %rdi, %rax
movq %rax, 0x48(%rsp)
movq %rdi, 0x150(%rsp)
movq %rsi, 0x148(%rsp)
movl %edx, 0x144(%rsp)
leaq 0x128(%rsp), %rdi
movq %rdi, 0x50(%rsp)
callq 0xc420
movq 0x50(%rsp), %rdi
movl 0x144(%rsp), %eax
shll $0x4, %eax
addl $0x10, %eax
movslq %eax, %rsi
callq 0xcac0
jmp 0xb06e
movl $0x0, 0x118(%rsp)
cmpl $0x10, 0x118(%rsp)
jge 0xb111
movq 0x38(%rsp), %rsi
leaq 0xf8(%rsp), %rdi
callq 0xaca0
jmp 0xb09b
leaq 0x128(%rsp), %rdi
leaq 0xf8(%rsp), %rsi
callq 0xcc00
jmp 0xb0b2
leaq 0xf8(%rsp), %rdi
callq 0x6f00
movl 0x118(%rsp), %eax
addl $0x1, %eax
movl %eax, 0x118(%rsp)
jmp 0xb079
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x120(%rsp)
movl %eax, 0x11c(%rsp)
jmp 0xb4d9
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x120(%rsp)
movl %eax, 0x11c(%rsp)
leaq 0xf8(%rsp), %rdi
callq 0x6f00
jmp 0xb4d9
movl $0x1, 0xf4(%rsp)
leaq 0x128(%rsp), %rdi
callq 0xcc30
movl 0x144(%rsp), %ecx
addl $0x1, %ecx
shll $0x4, %ecx
movslq %ecx, %rcx
cmpq %rcx, %rax
jae 0xb35b
movl $0x0, 0xf0(%rsp)
cmpl $0x4, 0xf0(%rsp)
jge 0xb345
leaq 0xd8(%rsp), %rdi
callq 0xc420
movl $0x0, 0xd4(%rsp)
cmpl $0x4, 0xd4(%rsp)
jge 0xb1f3
leaq 0x128(%rsp), %rdi
movq %rdi, 0x30(%rsp)
callq 0xcc30
movq 0x30(%rsp), %rdi
movslq 0xd4(%rsp), %rcx
leaq -0x5(%rax,%rcx), %rsi
callq 0xc4a0
movq %rax, %rsi
leaq 0xd8(%rsp), %rdi
callq 0xc430
jmp 0xb1b8
jmp 0xb1ba
movl 0xd4(%rsp), %eax
addl $0x1, %eax
movl %eax, 0xd4(%rsp)
jmp 0xb173
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x120(%rsp)
movl %eax, 0x11c(%rsp)
leaq 0xd8(%rsp), %rdi
callq 0xc2f0
jmp 0xb4d9
cmpl $0x0, 0xf0(%rsp)
jne 0xb215
movl 0xf4(%rsp), %esi
leaq 0xd8(%rsp), %rdi
callq 0xad60
jmp 0xb213
jmp 0xb215
movl $0x0, 0xd0(%rsp)
cmpl $0x4, 0xd0(%rsp)
jge 0xb291
leaq 0x128(%rsp), %rdi
movq %rdi, 0x20(%rsp)
callq 0xcc30
movq 0x20(%rsp), %rdi
movslq 0xd0(%rsp), %rcx
leaq -0x10(%rax,%rcx), %rsi
callq 0xc4a0
movq %rax, 0x28(%rsp)
movslq 0xd0(%rsp), %rsi
leaq 0xd8(%rsp), %rdi
callq 0xc4a0
movq 0x28(%rsp), %rsi
movq %rax, %rdi
callq 0x5d40
jmp 0xb27c
jmp 0xb27e
movl 0xd0(%rsp), %eax
addl $0x1, %eax
movl %eax, 0xd0(%rsp)
jmp 0xb220
leaq 0x128(%rsp), %rdi
movq %rdi, 0x10(%rsp)
callq 0xcd10
movq %rax, 0xc0(%rsp)
leaq 0xc8(%rsp), %rdi
leaq 0xc0(%rsp), %rsi
callq 0xcd40
leaq 0xd8(%rsp), %rdi
movq %rdi, 0x8(%rsp)
callq 0xcd80
movq 0x8(%rsp), %rdi
movq %rax, 0xb8(%rsp)
callq 0xcd10
movq 0x10(%rsp), %rdi
movq %rax, 0xb0(%rsp)
movq 0xc8(%rsp), %rsi
movq 0xb8(%rsp), %rdx
movq 0xb0(%rsp), %rcx
callq 0xcc50
movq %rax, 0x18(%rsp)
jmp 0xb315
movq 0x18(%rsp), %rax
movq %rax, 0xa8(%rsp)
leaq 0xd8(%rsp), %rdi
callq 0xc2f0
movl 0xf0(%rsp), %eax
addl $0x1, %eax
movl %eax, 0xf0(%rsp)
jmp 0xb14d
movl 0xf4(%rsp), %eax
addl $0x1, %eax
movl %eax, 0xf4(%rsp)
jmp 0xb11c
movq 0x40(%rsp), %rdi
movb $0x0, 0xa7(%rsp)
callq 0xcdb0
movl $0x0, 0xa0(%rsp)
movl 0xa0(%rsp), %eax
movl 0x144(%rsp), %ecx
addl $0x1, %ecx
cmpl %ecx, %eax
jge 0xb4a3
leaq 0x88(%rsp), %rdi
callq 0xc420
movl $0x0, 0x84(%rsp)
cmpl $0x10, 0x84(%rsp)
jge 0xb413
movl 0xa0(%rsp), %eax
shll $0x4, %eax
movl 0x84(%rsp), %ecx
addl %ecx, %eax
movslq %eax, %rsi
leaq 0x128(%rsp), %rdi
callq 0xc4a0
movq %rax, %rsi
leaq 0x88(%rsp), %rdi
callq 0xc430
jmp 0xb3e8
jmp 0xb3ea
movl 0x84(%rsp), %eax
addl $0x1, %eax
movl %eax, 0x84(%rsp)
jmp 0xb3a9
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x120(%rsp)
movl %eax, 0x11c(%rsp)
jmp 0xb48a
leaq 0x58(%rsp), %rdi
movl $0x4, %edx
leaq 0x88(%rsp), %rcx
movl %edx, %esi
callq 0xc160
jmp 0xb42e
movq 0x40(%rsp), %rdi
leaq 0x58(%rsp), %rsi
callq 0xcdc0
jmp 0xb43f
leaq 0x58(%rsp), %rdi
callq 0xc360
leaq 0x88(%rsp), %rdi
callq 0xc2f0
movl 0xa0(%rsp), %eax
addl $0x1, %eax
movl %eax, 0xa0(%rsp)
jmp 0xb378
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x120(%rsp)
movl %eax, 0x11c(%rsp)
leaq 0x58(%rsp), %rdi
callq 0xc360
leaq 0x88(%rsp), %rdi
callq 0xc2f0
movq 0x40(%rsp), %rdi
callq 0xcdf0
jmp 0xb4d9
movb $0x1, 0xa7(%rsp)
testb $0x1, 0xa7(%rsp)
jne 0xb4bf
movq 0x40(%rsp), %rdi
callq 0xcdf0
leaq 0x128(%rsp), %rdi
callq 0xc2f0
movq 0x48(%rsp), %rax
addq $0x158, %rsp            # imm = 0x158
retq
leaq 0x128(%rsp), %rdi
callq 0xc2f0
movq 0x120(%rsp), %rdi
callq 0x22b0
nopw %cs:(%rax,%rax)