From 46fcb3badca644e167749a709d4f28d6a747bdfc Mon Sep 17 00:00:00 2001
From: Stephen Heumann <stephenheumann@gmail.com>
Date: Tue, 27 Jun 2017 18:35:04 -0500
Subject: [PATCH] Add support for AES-256 decryption.

---
 aes.asm    | 18 ++++++++++++++----
 aes.h      |  1 +
 aes.macros | 12 ++++++++++++
 aestest.c  |  4 ++++
 4 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/aes.asm b/aes.asm
index ea4ecbb..02095bd 100644
--- a/aes.asm
+++ b/aes.asm
@@ -345,17 +345,27 @@ aes192_decrypt start
 	CFunction AES192_DECRYPT
 	end
 
-AES192_DECRYPT start
+aes256_decrypt start
+	CFunction AES256_DECRYPT
+	end
+
+AES256_DECRYPT start
 	using	tables
-	
+	InvFinalRound 14
+	InvNormalRound 13
+	InvNormalRound 12
+	jmp	cont1
+
+AES192_DECRYPT entry
 	InvFinalRound 12
+cont1	anop
 	InvNormalRound 11
 	InvNormalRound 10
-	jmp	cont
+	jmp	cont2
 	
 AES128_DECRYPT entry
 	InvFinalRound 10
-cont	anop
+cont2	anop
 	InvNormalRound 9
 	InvNormalRound 8
 	InvNormalRound 7
diff --git a/aes.h b/aes.h
index 4e171e2..374d1d0 100644
--- a/aes.h
+++ b/aes.h
@@ -14,3 +14,4 @@ void aes_encrypt(struct aes_state *state);
 
 void aes128_decrypt(struct aes_state *state);
 void aes192_decrypt(struct aes_state *state);
+void aes256_decrypt(struct aes_state *state);
diff --git a/aes.macros b/aes.macros
index bf41309..59d6af4 100644
--- a/aes.macros
+++ b/aes.macros
@@ -246,7 +246,13 @@
 	
 .loop
 	lda	&state+&i
+	aif	(rk+&round*16+&i)>255,.bigindex
 	eor	rk+&round*16+&i
+	ago	.cont
+.bigindex
+	ldx	#&round*16+&i
+	eor	rk,X
+.cont
 	sta	&state+&i
 &i	seta	&i+2
 	aif	&i<16,.loop
@@ -319,7 +325,13 @@
 	InvFinalRoundStep &to,&from
 	
 	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
 	sta	state2+&to
diff --git a/aestest.c b/aestest.c
index c4172c6..f7deef5 100644
--- a/aestest.c
+++ b/aestest.c
@@ -98,6 +98,10 @@ void aes256_test(void) {
     aes_encrypt(&aes_state);
     
     print_hexbytes("Output:       ", aes_state.data, 16);
+
+    aes256_decrypt(&aes_state);
+    
+    print_hexbytes("Decrypted:    ", aes_state.data, 16);
 }
 
 unsigned long aes128_time_test(unsigned int iters) {