added option for using custom values in bluh function and lenght fixes in bluh function

1 files changed, 21 insertions, 10 deletions

diff --git a/src/onetimebluh.c b/src/onetimebluh.c
index 3726489..905e8bb 100644
--- a/src/onetimebluh.c
+++ b/src/onetimebluh.c
@@ -9,7 +9,7 @@
 
 void xor(int ed, char* mess, char* keyf, char* outp);		/* operates the bitwise XOR between mess and keyf and puts the output to outp */
 void keyrand(int nb, char* outp);				/* generates random numbers using RAND_bytes from openssl and puts them into outp */
-void bluh(int c, char* mess, char* outp);			/* performs the binary dump of the input file and prints that to outp */
+void bluh(int c, char* mess, char* outp, char* ch);		/* performs the binary dump of the input file and prints that to outp */
 void help(char* av[]);						/* prints the help message */
 
 int main(int argc, char* argv[]) {
@@ -18,6 +18,7 @@ int main(int argc, char* argv[]) {
 	int command = 0;
 	int comm = 0;
 	int tear = 0;
+	char* chars = "01";
 	char* message = NULL;
 	char* keyfile = NULL;
 	char* output = NULL;
@@ -30,6 +31,7 @@ int main(int argc, char* argv[]) {
 		int option_index = 0;
 		static struct option options[] = {
 			{"bluh",	required_argument,	0,	'b'},
+			{"char",	required_argument,	0,	'c'},
 			{"decrypt",	required_argument,	0,	'd'},
 			{"encrypt",	required_argument,	0,	'e'},
 			{"key-file",	required_argument,	0,	'k'},
@@ -42,7 +44,7 @@ int main(int argc, char* argv[]) {
 			{0,		0,			0,	0},
 		};
 
-		if ((opt = getopt_long(argc, argv, "b:d:e:ghk:n:o:tu:", options, &option_index)) == -1)
+		if ((opt = getopt_long(argc, argv, "b:c:d:e:ghk:n:o:tu:", options, &option_index)) == -1)
 			break;
 
 		switch (opt) {
@@ -54,6 +56,13 @@ int main(int argc, char* argv[]) {
 			command ++;
 			comm = 'b';
 			break;
+		case 'c':
+			if (strlen(argv[optind - 1]) != 2){
+				printf("Too much or too many characters inserted!");
+				exit(EXIT_FAILURE);
+			}
+			chars = argv[optind - 1];
+			break;
 		case 'd':
 			message = argv[optind-1];
 			command ++;
@@ -135,7 +144,7 @@ int main(int argc, char* argv[]) {
 	} else if (comm == 'g') {
 		keyrand(nbytes, output);
 	} else if (comm == 'b' || comm == 'u') {
-		bluh(comm, message, output);
+		bluh(comm, message, output, chars);
 	}
 
 	exit(EXIT_SUCCESS);
@@ -250,10 +259,12 @@ void keyrand(int nb, char* outp) {
 	return;
 }
 
-void bluh(int c, char* mess, char* outp) {
+void bluh(int c, char* mess, char* outp, char* ch) {
 
 	char* defbluh = "bluhed";
 	char* defunbluh = "unbluhed";
+	char a = ch[0];
+	char b = ch[1];
 
 	/* In absence of input by users next block sets the default values */
 
@@ -286,7 +297,7 @@ void bluh(int c, char* mess, char* outp) {
 
 		for (i = 0; i < len; i ++) {
 			for (j = 0; j < 8; j ++) {
-				bits[8 * i + j] = (bytes[i] & (1 << (7 -j))) ? '1' : '0';
+				bits[8 * i + j] = (bytes[i] & (1 << (7 -j))) ? b : a;
 				printf("%c", bits[8 * i + j]);
 			}
 
@@ -303,17 +314,17 @@ void bluh(int c, char* mess, char* outp) {
 		printf("Message successfully bluhed!\n");
 
 	} else if (c == 'u') {
-		char* comp = malloc(len + 1);
-		char* bits = malloc(8 * len);
+		char* comp = malloc((len / 8) + 1);
+		char* bits = malloc(len);
 
-		fread(bits, sizeof(char), 8 * len, mex);
-		memset(comp, 0, (len + 1));
+		fread(bits, sizeof(char), len, mex);
+		memset(comp, 0, ((len / 8) + 1));
 
 		int i, j;
 
 		for (i = 0; i < len; i ++) {
 			for (j = 0; j < 8; j ++) {
-				if (bits[8 * i + j] == '1') {
+				if (bits[8 * i + j] == b) {
 					comp[i] |= 1 << (7 - j);
 				}
 			}