; void aes128_key_schedule(unsigned char key[176])
; void aes128_encrypt(unsigned char buffer[16], const unsigned char key[176])

PUBLIC aes128_key_schedule, aes128_key_inverse, aes128_encrypt

_TEXT SEGMENT ALIGN(8) EXECUTE READ ALIAS('.text') 'CODE'

aes128_key_schedule:
    movdqu          xmm0, XMMWORD PTR [rcx]    ; load key
    aeskeygenassist xmm1, xmm0, 1              ; prepare  1st round key
    call            aes128_key_expansion       ; prepare  1st round key
    aeskeygenassist xmm1, xmm0, 2              ; prepare  2nd round key
    call            aes128_key_expansion       ; prepare  2nd round key
    aeskeygenassist xmm1, xmm0, 4              ; prepare  3rd round key
    call            aes128_key_expansion       ; prepare  3rd round key
    aeskeygenassist xmm1, xmm0, 8              ; prepare  4th round key
    call            aes128_key_expansion       ; prepare  4th round key
    aeskeygenassist xmm1, xmm0, 16             ; prepare  5th round key
    call            aes128_key_expansion       ; prepare  5th round key
    aeskeygenassist xmm1, xmm0, 32             ; prepare  6th round key
    call            aes128_key_expansion       ; prepare  6th round key
    aeskeygenassist xmm1, xmm0, 64             ; prepare  7th round key
    call            aes128_key_expansion       ; prepare  7th round key
    aeskeygenassist xmm1, xmm0, 128            ; prepare  8th round key
    call            aes128_key_expansion       ; prepare  8th round key
    aeskeygenassist xmm1, xmm0, 27             ; prepare  9th round key
    call            aes128_key_expansion       ; prepare  9th round key
    aeskeygenassist xmm1, xmm0, 54             ; prepare 10th round key
    ; call          aes128_key_expansion       ; prepare 10th round key
    ; ret
aes128_key_expansion:
    add             rcx, 16                    ; address of next round key
    pshufd          xmm1, xmm1, 255
    pxor            xmm1, xmm0
    pslldq          xmm0, 4
    pxor            xmm1, xmm0
    pslldq          xmm0, 4
    pxor            xmm1, xmm0
    pslldq          xmm0, 4
    pxor            xmm0, xmm1
    movdqu          XMMWORD PTR [rcx], xmm0    ; store round key
    ret

aes128_encrypt:
    movdqu          xmm0, XMMWORD PTR [rcx]    ; load plaintext
    pxor            xmm0, XMMWORD PTR [rdx]    ; initial xor
    mov             al, 9                      ; number of main rounds
aes128_encrypt_loop:
    add             rdx, 16                    ; next round key
    aesenc          xmm0, XMMWORD PTR [rdx]    ; encrypt round
    dec             al                         ; loop control variable
    jne             aes128_encrypt_loop
    aesenclast      xmm0, XMMWORD PTR [rdx+16] ; last encrypt round
    movdqu          XMMWORD PTR [rcx], xmm0    ; store ciphertext
    ret

_TEXT ENDS

END