Implemented option parsing

Signed-off-by: Gon Namprasertkul <foresterblox@gmail.com>
Signed-off-by: Ian Moffett <ian@osmora.org>

1 files changed, 26 insertions, 10 deletions

diff --git a/src/main.c b/src/main.c
index f969b6b..e8e736f 100644
--- a/src/main.c
+++ b/src/main.c
@@ -35,6 +35,8 @@
 #include <unistd.h>
 #include <info.h>
 
+#define DEFAULT_FILENAME "fob"
+
 #if __BYTE_ORDER__ != __ORDER_LITTLE_ENDIAN__
 #error "Big endian machines not supported yet"
 #endif
@@ -59,7 +61,6 @@ read_file(const char *fname, size_t *size_out)
     size_t bufsize;
 
     if (access(fname, F_OK) != 0) {
-        fprintf(stderr, "%s does not exist!\n", fname);
         return NULL;
     }
 
@@ -216,29 +217,44 @@ obfuscate(const struct cpu_info *info, char *buf, size_t buf_size)
 }
 
 int
-main(int argc, const char **argv)
+main(int argc, char **argv)
 {
     size_t buf_size;
     char *buf;
     struct cpu_info info = { 0 };
+    int8_t opt;
+    char *output = DEFAULT_FILENAME;
+    char *input = NULL;
+
+    while ((opt = getopt(argc, (char *const *)argv, "o:")) != -1)
+    {
+        if (opt == 'o') {
+            output = optarg;
+        }
+    }
+
+    input = argv[optind];
 
-    if (argc < 2) {
+    if (input == NULL) {
         fprintf(stderr, "Usage: %s <file>\n", argv[0]);
         return 1;
     }
 
-#if defined(__x86_64__)
-    amd64_cpu_tests(&info);
-#endif  /* __x86_64__ */
-
-    buf = read_file(argv[1], &buf_size);
+    buf = read_file(input, &buf_size);
     if (buf == NULL) {
+        fprintf(stderr, "Failed to read %s. Does it exist?\n", input);
         return 1;
     }
 
+#if defined(__x86_64__)
+    amd64_cpu_tests(&info);
+#endif  /* __x86_64__ */
+
     obfuscate(&info, buf, buf_size);
-    writeback_file(argv[1], buf, buf_size);
-    free(buf);
+    writeback_file(output, buf, buf_size);
+
+    fprintf(stdout, "Written to %s\n", output);
 
+    free(buf);
     return 0;
 }