* Copyright (c) 2017 Stephen Heumann
*
* Permission to use, copy, modify, and/or distribute this software for any
* purpose with or without fee is hereby granted, provided that the above
* copyright notice and this permission notice appear in all copies.
*
* THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
* WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
* MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
* ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
* WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
* ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
* OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

* This makes a function wrapper that is callable from C,
* taking a pointer to the context structure as its argument.
	macro
	CFunction &fn
	phb
	plx
	ply
	tdc
	pld
	plb
	plb
	phy
	phx
	plb
	pha
	jsl	&fn
	pld
	rtl
	mend

*The 'core' function applied to some words when computing the AES key schedule
	macro
	ExpandKeyCore &xorback,&rconoffset

	ShortRegs
	
	ldy	rk-3,x
	lda	Sbox,y
	eor	Rcon-&xorback+&rconoffset,x
	eor	rk-&xorback,x
	sta	rk,x
	
	ldy	rk-2,x
	lda	Sbox,y
	eor	rk+1-&xorback,x
	sta	rk+1,x
	
	ldy	rk-1,x
	lda	Sbox,y
	eor	rk+2-&xorback,x
	sta	rk+2,x
	
	ldy	rk-4,x
	lda	Sbox,y
	eor	rk+3-&xorback,x
	sta	rk+3,x

	LongRegs

	mend

* Secondary substitution step used when expanding AES-256 keys
	macro
	ExpandKeySubst &xorback,&rconoffset

	ShortRegs
	
	ldy	rk-4,x
	lda	Sbox,y
	eor	rk-&xorback,x
	sta	rk,x
	
	ldy	rk+1-4,x
	lda	Sbox,y
	eor	rk+1-&xorback,x
	sta	rk+1,x
	
	ldy	rk+2-4,x
	lda	Sbox,y
	eor	rk+2-&xorback,x
	sta	rk+2,x
	
	ldy	rk+3-4,x
	lda	Sbox,y
	eor	rk+3-&xorback,x
	sta	rk+3,x

	LongRegs

	mend

* Generate consecutive words of key schedule that don't use above functions
	macro
	ExpandKeyIter &xorback,&nwords
	lcla	&i
	lda	rk,x
.loop1
	eor	rk+&i+4-&xorback,x
	sta	rk+&i+4,x
&i	seta	&i+4
	aif	&i/4<&nwords,.loop1

&i	seta	2
	lda	rk+2,x
.loop2
	eor	rk+&i+4-&xorback,x
	sta	rk+&i+4,x
&i	seta	&i+4
	aif	&i/4<&nwords,.loop2
	mend

* Do an initial AddRoundKey step on the starting state (for encryption)
	macro
	AddInitialRoundKey
	lcla	&i
.top
	lda	state1+&i
	eor	rk+&i
	sta	state1+&i
&i	seta	&i+2
	aif	&i<16,.top
	mend

* Do a full normal round, including (in effect) SubBytes through AddRoundKey
	macro
&lbl	NormalRound &round
&lbl	anop
	aif	&round/2*2=&round,.evenround
	MixColumn 0,0,5,10,15,state1,state2
	MixColumn 4,4,9,14,3,state1,state2
	MixColumn 8,8,13,2,7,state1,state2
	MixColumn 12,12,1,6,11,state1,state2
	ago	.done
.evenround
	MixColumn 0,0,5,10,15,state2,state1
	MixColumn 4,4,9,14,3,state2,state1
	MixColumn 8,8,13,2,7,state2,state1
	MixColumn 12,12,1,6,11,state2,state1	
.done
	mend

* Do the operations on one column for a normal round.
	macro
	MixColumn &i,&A,&B,&C,&D,&state,&out

	aif	(&i=0).AND.(&round<>1),.skip
	ldy	&state+&D
.skip
	lda	Sbox,Y
	pha
	ldx	&state+&A
	eor	Xtime2Sbox,X
	ldy	&state+&B
	eor	Xtime3Sbox,Y
	ldy	&state+&C
	eor	Sbox,Y
	eor	rk+&round*16+&i
	sta	&out+&i

	pla
	eor	Xtime3Sbox,Y	
	eor	Sbox,X
	ldy	&state+&B
	eor	Xtime2Sbox,Y
	eor	rk+&round*16+&i+1
	sta	&out+&i+1

	lda	Sbox,Y
	pha
	ldy	&state+&D
	eor	Xtime3Sbox,Y
	eor	Sbox,X
	ldy	&state+&C
	eor	Xtime2Sbox,Y
	eor	rk+&round*16+&i+2
	sta	&out+&i+2

	pla
	eor	Sbox,Y
	eor	Xtime3Sbox,X
	ldy	&state+&D
	eor	Xtime2Sbox,Y
	aif	(rk+&round*16+&i+3)>255,.bigindex
	eor	rk+&round*16+&i+3
	ago	.cont
.bigindex
	ldx	#&round*16+&i+3
	eor	rk,X
.cont
	sta	&out+&i+3
	aif	&i<>12,.skip2
	tay
.skip2
	mend

* Do final round, including (in effect) SubBytes, ShiftRows, and AddRoundKey.
	macro
	FinalRound &round

	FinalRoundStep 3,15,1
	FinalRoundStep 15,11
	FinalRoundStep 11,7
	FinalRoundStep 7,3

	FinalRoundStep 0,0
	FinalRoundStep 4,4
	FinalRoundStep 8,8
	FinalRoundStep 12,12

	FinalRoundStep 13,1
	FinalRoundStep 1,5
	FinalRoundStep 5,9
	FinalRoundStep 9,13

	FinalRoundStep 10,2
	FinalRoundStep 2,10
	FinalRoundStep 14,6
	FinalRoundStep 6,14
	
	mend

* Do the final round operations for one byte.
	macro
	FinalRoundStep &to,&from,&skipldy

	aif	C:&skipldy,.skip
	ldy	state2+&from
.skip
	lda	Sbox,Y
	aif	(rk+&round*16+&to)>255,.bigindex
	eor	rk+&round*16+&to
	ago	.cont
.bigindex
	ldx	#&round*16+&to
	eor	rk,X
.cont
	sta	state1+&to

	mend

* Perform an inverse normal round (for decryption)
	macro
	InvNormalRound &round,&state
	lcla	&i
	lclc	&state
	lclc	&out

	aif	&round/2*2=&round,.evenround
&state	setc	state2
&out	setc	state1
	ago	.cont
.evenround
&state	setc	state1
&out	setc	state2
.cont

	InvMixColumn 12,1,6,11,12,1
	InvMixColumn 0,5,10,15,0
	InvMixColumn 8,13,2,7,8
	InvMixColumn 4,9,14,3,4,dotax=1
	mend

* Perform the operations for one column in an inverse normal round
	macro
	InvMixColumn &A,&B,&C,&D,&i,&skipldx,&dotax

	aif	C:&skipldx,.skip
	ldx	&state+&i+2
.skip
	lda	Xtime9,X
	ldy	&state+&i+0
	eor	XtimeB,Y
	ldy	&state+&i+1
	eor	XtimeD,Y
	ldy	&state+&i+3
	eor	XtimeE,Y
	tay
	lda	InvSbox,Y
	eor	rk+(&round-1)*16+&D
	sta	&out+&D
	
	ldy	&state+&i+0
	lda	XtimeE,Y
	ldy	&state+&i+1
	eor	XtimeB,Y
	eor	XtimeD,X
	ldy	&state+&i+3
	eor	Xtime9,Y
	tay
	lda	InvSbox,Y
	eor	rk+(&round-1)*16+&A
	sta	&out+&A
	
	ldy	&state+&i+0
	lda	Xtime9,Y
	ldy	&state+&i+1
	eor	XtimeE,Y
	eor	XtimeB,X
	ldy	&state+&i+3
	eor	XtimeD,Y
	tay
	lda	InvSbox,Y
	eor	rk+(&round-1)*16+&B
	sta	&out+&B
	
	ldy	&state+&i+0
	lda	XtimeD,Y
	ldy	&state+&i+1
	eor	Xtime9,Y
	eor	XtimeE,X
	ldy	&state+&i+3
	eor	XtimeB,Y
	tay
	lda	InvSbox,Y
	eor	rk+(&round-1)*16+&C
	aif	(C:&dotax).AND.(&round<>1),.dotax
	sta	&out+&C
	ago	.done
.dotax
	tax
.done
	mend

* Do the inverse final round steps for one byte.
	macro
	InvFinalRoundStep &to,&from,&dotax
	
	lda	state1+&from
	aif	(rk+&round*16+&from)>255,.bigindex
	eor	rk+&round*16+&from
	ago	.cont
.bigindex
	ldx	#&round*16+&from
	eor	rk,X
.cont
	tay
	lda	InvSbox,Y

	aif	(rk+(&round-1)*16+&to)>255,.bigindex2
	eor	rk+(&round-1)*16+&to
	ago	.cont2
.bigindex2
	ldx	#(&round-1)*16+&to
	eor	rk,X
.cont2
	aif	C:&dotax,.dotax
	sta	state2+&to
	ago	.done
.dotax
	tax
.done
	mend

* Do the inverse of the final round (which comes first for decryption).
	macro
	InvFinalRound &round
	
	InvFinalRoundStep 0,0
	InvFinalRoundStep 4,4
	InvFinalRoundStep 8,8
	InvFinalRoundStep 12,12
	
	InvFinalRoundStep 1,13
	InvFinalRoundStep 13,9
	InvFinalRoundStep 9,5
	InvFinalRoundStep 5,1
	
	InvFinalRoundStep 15,3
	InvFinalRoundStep 3,7
	InvFinalRoundStep 7,11
	InvFinalRoundStep 11,15

	InvFinalRoundStep 2,10
	InvFinalRoundStep 10,2
	InvFinalRoundStep 6,14
	InvFinalRoundStep 14,6,1
	
	mend
	
* Set registers to 8 bits
	macro
	ShortRegs
	sep	#$30
	longa	off
	longi	off
	mend

* Set registers to 16 bits
	macro
	LongRegs
	rep	#$30
	longa	on
	longi	on
	mend