Improve argument parsing

1 files changed, 45 insertions, 26 deletions

diff --git a/xorcat.c b/xorcat.c
index 715fc44..5664ca3 100644
--- a/xorcat.c
+++ b/xorcat.c
@@ -1,18 +1,12 @@
-#include <stdlib.h>
-#include <stdio.h>
-#include <string.h>
-#include <stdint.h>
 #include <unistd.h>
 #include <err.h>
+#include <stdbool.h>
+#include <string.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
 
-static const char usage[] = "USAGE: xorcat KEYMASK";
-
-void
-die(const char *fmtstr)
-{
-	fprintf(stderr, "%s\n", fmtstr);
-	exit(1);
-}
+static const char usage[] = "USAGE: xorcat [-h] [-m] KEYMASK";
 
 void
 xorcat(uint8_t *indata, size_t insize, uint8_t *key, size_t keylen)
@@ -59,18 +53,41 @@ collect(size_t *size)
 int
 main(int argc, const char **argv)
 {
-	uint8_t *mask, *key;
+	uint8_t *mask, *key, keyb;
+	const char *maskarg;
 	size_t keylen;
 	uint8_t *indata;
 	size_t insize;
-	unsigned maskb;
+	bool markers;
 	ssize_t i;
 
-	if (argc != 2) die(usage);
+	if (argc <= 1) {
+		printf("%s\n", usage);
+		exit(0);
+	}
 
-	if (strlen(argv[1]) % 2) die(usage);
+	maskarg = NULL;
+	markers = false;
+	for (i = 1; i < argc; i++) {
+		if (!strcmp(argv[i], "-h")) {
+			printf("%s\n", usage);
+			exit(0);
+		} else if (!strcmp(argv[i], "-m")) {
+			markers = true;
+		} else if (maskarg) {
+			errx(1, "Multiple mask arguments");
+		} else {
+			maskarg = argv[i];
+		}
+	}
+
+	if (!maskarg)
+		errx(1, "Missing mask");
 
-	keylen = strlen(argv[1]) / 2;
+	if (strlen(maskarg) % 2)
+		errx(1, "Bad mask length");
+
+	keylen = strlen(maskarg) / 2;
 	mask = calloc(keylen, 1);
 	key = calloc(keylen, 1);
 	if (!mask || !key) err(1, "malloc");
@@ -79,20 +96,22 @@ main(int argc, const char **argv)
 		if (!strncmp(argv[1] + 2 * i, "??", 2)) {
 			mask[i] = 1;
 		} else {
-			if (!sscanf(argv[1] + 2 * i, "%02X", &maskb))
-				die("Bad mask");
-			mask[i] = (uint8_t) maskb;
+			if (!sscanf(maskarg + 2 * i, "%02hhX", &keyb))
+				errx(1, "Bad mask");
+			key[i] = (uint8_t) keyb;
 		}
 	}
 
 	indata = collect(&insize);
 	while (1) {
-		printf("\n");
-		printf(">>>>>>>>>>>>>>>>>>>>>>>>>>>>> ");
-		for (i = 0; i < keylen; i++)
-			printf("%02X", key[i]);
-		printf("\n");
-		fflush(stdout);
+		if (markers) {
+			printf("\n");
+			printf(">>>>>>>>>>>>>>>>>>>>>>>>>>>>> ");
+			for (i = 0; i < keylen; i++)
+				printf("%02X", key[i]);
+			printf("\n");
+			fflush(stdout);
+		}
 
 		xorcat(indata, insize, key, keylen);