From 7738b97d0a9c77292af74499b45a20bfab599552 Mon Sep 17 00:00:00 2001
From: x1phosura <x1phosura@x1phosura.zone>
Date: Sun, 8 May 2022 23:25:28 -0700
Subject: [PATCH] Add block-aligned AES ECB encryption (mbedtls)

---
 set1/Makefile      | 23 ++++++++------
 set1/README.md     |  6 ++--
 set1/src/aes-ecb.c | 76 ++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 94 insertions(+), 11 deletions(-)
 create mode 100644 set1/src/aes-ecb.c

diff --git a/set1/Makefile b/set1/Makefile
index e01c43e..971b29e 100644
--- a/set1/Makefile
+++ b/set1/Makefile
@@ -8,18 +8,23 @@ SRC_DIR = src
 OUT_DIR = bin
 TARGETS = base64 fixed-xor char-freq-analyze repeating-key-xor aes-ecb
 
+AES_LIB = mbedtls
+#AES_LIB = tiny-AES-c
+
+# choose which AES library to use
+ifeq (mbedtls,$(AES_LIB))
+    EXT_INCLUDES = -Iaes-libs/mbedtls/include/
+    EXT_LIBS     = -Laes-libs/mbedtls/library
+    LDLIBS       = -lmbedcrypto
+else
+    #EXT_INCLUDES = -Iaes-libs/tiny-AES-c
+    EXT_LIBS     = -Laes-libs/tiny-AES-c
+    #LDLIBS       = -laes
+endif
+
 
 all: $(TARGETS)
 
-# TODO toggle mbedtls or tiny-AES-c
-#EXT_INCLUDES = -Iaes-libs/mbedtls/include/
-#EXT_LIBS     = -Laes-libs/mbedtls/library
-#LDLIBS       = -lmbedcrypto
-#
-EXT_INCLUDES = -Iaes-libs/tiny-AES-c
-#EXT_LIBS     = -Laes-libs/tiny-AES-c
-#LDLIBS       = -laes
-
 %: $(SRC_DIR)/%.c
 	$(CC) $(CFLAGS) -o $(OUT_DIR)/$@ $<
 
diff --git a/set1/README.md b/set1/README.md
index debfe6b..8ccb17c 100644
--- a/set1/README.md
+++ b/set1/README.md
@@ -9,7 +9,9 @@ Provided by base64.c
 Provided by fixed-key-xor.c
 
 ### Challenge 3: Single-byte XOR cipher
-Provided by ch3-brute-force.c and char-freq-analyze.c
+Provided by:
+- ch3-brute-force.c
+- char-freq-analyze.c (almost complete)
 
 ### Challenge 4: Detect single-character XOR
 
@@ -21,7 +23,7 @@ Provided by repeating-key-xor.c
 
 
 ### Challenge 7: AES in ECB mode
-
+Provided by aes-ecb.c (almost complete)
 
 ### Challenge 8: Detect AES in ECB mode
 
diff --git a/set1/src/aes-ecb.c b/set1/src/aes-ecb.c
new file mode 100644
index 0000000..819abcf
--- /dev/null
+++ b/set1/src/aes-ecb.c
@@ -0,0 +1,76 @@
+#include <stdbool.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#define MBEDTLS
+
+#ifdef MBEDTLS
+
+#include <mbedtls/aes.h>    // might need to be in double-quotes
+#include <mbedtls/config.h>
+
+#else
+
+#include <tiny-AES-c/aes.h>
+
+#endif
+
+void print_buffer(uint8_t *buffer, size_t buff_len)
+{
+	size_t i;
+	for (i = 0; i < buff_len; ++i) {
+		putchar(buffer[i]);
+	}
+}
+
+// TODO: write generic AES function that uses mbedtls or tiny-AES-c
+bool aes_128_enc_ecb(uint8_t key[16], uint8_t *plaintext, size_t plaintext_len,
+                                    uint8_t *ciphertext, size_t ciphertext_len)
+{
+	size_t i;
+	mbedtls_aes_context aes_ctx;
+
+	if (plaintext_len > ciphertext_len) {
+		fprintf(stderr, "Plaintext longer than ciphertext buffer\n");
+		return false;
+	}
+
+	if (mbedtls_aes_setkey_enc(&aes_ctx, key, 128) != 0) {
+		fprintf(stderr, "Setting AES encryption key\n");
+		return false;
+	}
+
+	for (i = 0; i < plaintext_len; i = i + 16) {
+		// TODO: handle ciphertext not block size-aligned
+		if (mbedtls_aes_crypt_ecb(&aes_ctx, MBEDTLS_AES_ENCRYPT,
+					  plaintext + i, ciphertext + i) != 0) {
+			fprintf(stderr, "Performing AES encryption\n");
+			return false;
+		}
+	}
+
+	return true;
+}
+
+
+int main()
+{
+	// for testing (TODO: take key, plaintext as args)
+	uint8_t key[16];
+	uint8_t message[32];
+	uint8_t ciph_out[48];
+	memset(key, 'B', sizeof(key));
+	memset(message, 'A', sizeof(message));
+	memset(ciph_out, 0, sizeof(ciph_out));
+
+	if (!aes_128_enc_ecb(key, message, sizeof(message),
+	                           ciph_out, sizeof(ciph_out))) {
+		return EXIT_FAILURE;
+	}
+
+	print_buffer(ciph_out, sizeof(ciph_out));
+
+	return EXIT_SUCCESS;                              
+}
+