1 files changed, 31 insertions, 15 deletions

diff --git a/lib/lib2.c b/lib/lib2.c
index d84c781..f9904b6 100644
--- a/lib/lib2.c
+++ b/lib/lib2.c
@@ -1,3 +1,4 @@
+#include "lib3.h"
 #include "lib2.h"
 #include "lib.h"
 
@@ -422,7 +423,7 @@ void send_user(char *encrypted_user, int length)
 }
 
 
-int challenge16_encrypt(char *input, char **encrypted)
+int challenge16_encrypt(char *input, char **encrypted, int cbc_mode)
 {
 	char *prepend = "comment1=cooking\%20MCs;userdata=";
 	char *append = ";comment2=\%20like\%20a\%20pound\%20of\%20bacon";
@@ -465,24 +466,39 @@ int challenge16_encrypt(char *input, char **encrypted)
 	memcpy(&res[strlen(prepend)+strlen(input)+2*quote_char], append, strlen(append));
 	res[strlen(prepend)+strlen(input)+strlen(append)+2*quote_char+1] = '\0';	
 	// padding
-	unencrypted = pkcs7_padding(res, strlen(res), 16);
-	*encrypted = malloc(strlen(unencrypted));
-	aes_cbc(unencrypted, strlen(unencrypted), *encrypted, key, iv , 1);
-	return strlen(unencrypted);
+	if(cbc_mode) {
+		unencrypted = pkcs7_padding(res, strlen(res), 16);
+		*encrypted = malloc(strlen(unencrypted));
+		aes_cbc(unencrypted, strlen(unencrypted), *encrypted, key, iv , 1);
+		return strlen(unencrypted);
+	}
+	// otherwise is CTR mode
+	else {
+		*encrypted = malloc(strlen(res));
+		aes_ctr(res, strlen(res), *encrypted, key, iv);
+		return strlen(res);
+	}
 }
 
-void challenge16_decrypt(char *encrypted, int length)
+void challenge16_decrypt(char *encrypted, int length, int cbc_mode)
 {
 	char *unencrypted = malloc(length);
 	char *unpadd= malloc(length);
-	
-	aes_cbc(encrypted, length, unencrypted, key, iv, 0);
-	// unpadd
-	int ret = valid_pkcs7_padding(unencrypted, length, unpadd, 16);
-	if(!ret) {
-		printf("no valid padding!\n");
-		return;
+
+	if(cbc_mode) {
+		aes_cbc(encrypted, length, unencrypted, key, iv, 0);
+		// unpadd
+		int ret = valid_pkcs7_padding(unencrypted, length, unpadd, 16);
+		if(!ret) {
+			printf("no valid padding!\n");
+			return;
+		}
+		// look for string ;admin=true;
+		printf("unencrpyted string: %s\n", unpadd);
+	}
+	// we are in ctr mode
+	else {
+		aes_ctr(encrypted, length, unencrypted, key, iv);
+		printf("unencrpted string: %s\n", unencrypted);
 	}
-	// look for string ;admin=true;
-	printf("unencrpyted string: %s\n", unpadd);
 }