From 79db4fa8975cd445958f88a046c75243300b7b69 Mon Sep 17 00:00:00 2001
From: jakob <jakob@memeware.net>
Date: Fri, 3 Mar 2017 21:04:36 -0500
Subject: Implemented error checking to prevent against segfaulting when a
 corrupted archive is encountered.

---
 src/cli.c  | 16 +++++++++-------
 src/cli.h  |  2 +-
 src/main.c | 39 ++++++++++++++++++++++++++++++++-------
 3 files changed, 42 insertions(+), 15 deletions(-)

(limited to 'src')

diff --git a/src/cli.c b/src/cli.c
index 9f4a6bd..41ed8a1 100644
--- a/src/cli.c
+++ b/src/cli.c
@@ -39,8 +39,8 @@ struct params parse_args(int argc, char **argv) {
     int cur = 0, opt_index = 0, count = 0;
     static struct option long_opts[] = {
         {"help", no_argument, NULL, 'h'},
-        {"version", no_argument, NULL, 'v'},
-        {"quiet", no_argument, NULL, 'q'},
+        {"version", no_argument, NULL, 'V'},
+        {"verbose", no_argument, NULL, 'v'},
         {"extract", no_argument, NULL, 'e'},
         {"list", no_argument, NULL, 'l'},
         {"output", no_argument, NULL, 'o'},
@@ -50,24 +50,26 @@ struct params parse_args(int argc, char **argv) {
 
     do {
         count++;
-        cur = getopt_long(argc, argv, "hvelqo:", long_opts, &opt_index);
+        cur = getopt_long(argc, argv, "hVvelqo:g:", long_opts, &opt_index);
         switch (cur) {
             case 'h':
                 p.mode = HELP;
                 return p;
-            case 'v':
+            case 'V':
                 p.mode = VERSION;
                 return p;
-            case 'q':
-                p.quiet = true;
+            case 'v':
+                p.verbose = true;
                 break;
             case 'o':
                 count++;
                 p.out_len = strlen(optarg);
                 p.out = malloc(p.out_len + 2);
                 strcpy(p.out, optarg);
-                if (p.out[p.out_len - 1] != '/')
+                if (p.out[p.out_len - 1] != '/') {
                     p.out[p.out_len] = '/';
+                    p.out_len += 1;
+                }
                 break;
             case 'g':
                 count++;
diff --git a/src/cli.h b/src/cli.h
index 6ccfd86..662d955 100644
--- a/src/cli.h
+++ b/src/cli.h
@@ -35,7 +35,7 @@ enum {
 struct params {
     int    mode; /* Current mode of operation. */
     int    game; /* Which encryption keys to use. */
-    bool   quiet; /* Whether or not to suppress progress messages. */
+    bool   verbose; /* Whether or not to output progress messages. */
     int    vararg_index; /* Start index of paths in argv. */
     char   *out; /* Path to extract files to. */
     size_t out_len; /* Length of the output path string. */
diff --git a/src/main.c b/src/main.c
index da9a835..5363985 100644
--- a/src/main.c
+++ b/src/main.c
@@ -30,7 +30,10 @@
 #include "io.h"
 #include "table.h"
 
-#define VERSION_STR "2.0.0a1"
+#define EXIT_FAILURE 1
+#define EXIT_SUCCESS 0
+
+#define VERSION_STR "2.0.0a2"
 
 /* Pointer to a function to be mapped to entries in the table. */
 typedef void (*mapfn)(struct stream *archive, struct table_entry *e,
@@ -107,25 +110,47 @@ static char *get_path(struct params p, char *name) {
     size_t name_len = strlen(name);
     char *path = malloc(p.out_len + name_len + 2);
     strcpy(path, p.out);
-    strcpy(path + p.out_len + 1, name);
+    strcpy(path + p.out_len, name);
     return path;
 }
 
 
 /* Basic mapfn for printing the filename of each entry. */
 static void list(struct stream *s, struct table_entry *e, struct params p) {
+    if (e->filename == NULL) {
+        fprintf(stderr, "Found File entry without matching eliF. Archive may "
+                "be corrupted.\n");
+        return;
+    }
     printf("%s\n", e->filename);
 }
 
 
 /* Mapfn for extracting the contents of the archive. */
 static void extract(struct stream *s, struct table_entry *e, struct params p) {
+    if (e->filename == NULL) {
+        fprintf(stderr, "Found File entry without matching eliF. Archive may "
+                "be corrupted.\n");
+        return;
+    } else if (e->segments == NULL) {
+        fprintf(stderr, "Found eliF entry without matching File. Archive may "
+                "be corrupted.\n");
+        return;
+    }
+
     char *path = get_path(p, e->filename);
     make_dirs(path);
 
     FILE           *fp        = fopen(path, "wb+");
     struct stream  *segm_data;
-    struct segment *segm;    
+    struct segment *segm;
+
+    if (fp == NULL) {
+        perror(path);
+        exit(EXIT_FAILURE);
+    } else if (p.verbose) {
+        printf("Extracting %s to %s...\n", e->filename, path);
+    }
 
     for (uint64_t i = 0; i < e->segment_count; i++) {
         segm = e->segments[i];
@@ -137,11 +162,11 @@ static void extract(struct stream *s, struct table_entry *e, struct params p) {
             segm_data = stream_clone(s, segm->compressed_size);
         }
 
-        struct game_key k = get_key(NEKOPARA_VOLUME_1);
+        struct game_key k = get_key(p.game);
         uint8_t initial = derive_initial(k, e->key);
         uint8_t primary = derive_primary(k, e->key);
         stream_xor(segm_data, initial, primary);
-            
+
         stream_dump(fp, segm_data, segm_data->len);
         stream_free(segm_data);
     }
@@ -170,7 +195,7 @@ static void map_entries(char *path, struct params p, mapfn fn) {
 }
 
 
-int main(int argc, char **argv) {    
+int main(int argc, char **argv) {
     struct params p = parse_args(argc, argv);
     switch (p.mode) {
         case USAGE:
@@ -191,5 +216,5 @@ int main(int argc, char **argv) {
                 map_entries(argv[i], p, extract);
     }
     params_free(p);
-    return 0;
+    return EXIT_SUCCESS;
 }
-- 
cgit v1.3