Threaded.

1 files changed, 204 insertions, 81 deletions

diff --git a/main.c b/main.c
index faeb4c8..c64bf1b 100644
--- a/main.c
+++ b/main.c
@@ -1,9 +1,13 @@
-#include <tomcrypt.h>
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <stdbool.h>
 #include <assert.h>
+#include <threads.h>
+#include <stdatomic.h>
+#include <errno.h>
+
+#include <tomcrypt.h>
 
 const char *hexchars = "0123456789ABCDEF";
 
@@ -31,7 +35,11 @@ void new_key_search_ctx(
 
 /* search key in range [a, b), returns false if
  * no result yield from this call and searching is finished */
-bool yield_possible_key(key_search_ctx *ctx, uint32_t b) {
+bool yield_possible_key(
+        key_search_ctx *ctx,
+        uint32_t b,
+        atomic_bool *stop_signal
+) {
     if (ctx->finished) return false;
 
 //    const char[] hexchars = {0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37,
@@ -44,6 +52,11 @@ bool yield_possible_key(key_search_ctx *ctx, uint32_t b) {
     char key[8];
     symmetric_key skey;
     do {
+        if ((b != 0 && k >= b) || atomic_load(stop_signal)) {
+            // out of range, stop
+            ctx->finished = true;
+            return false;
+        }
         // convert key uint32 to char[]
         FILL_KEY(key, k, 0);
         FILL_KEY(key, k, 1);
@@ -66,14 +79,16 @@ bool yield_possible_key(key_search_ctx *ctx, uint32_t b) {
                 (const symmetric_key *) &skey
         ) != CRYPT_OK)
             continue; // failed to decrypt
-        // validate first 3 bytes of the cleartext
+        /* validate JPEG header (first 3 bytes) of the plaintext */
         if ((plaintext & 0xFFFFFFu) == 0xFFD8FFu) {
             ctx->yield = *(uint64_t *) key;
-            // if current `k` goes out of range, finished will be set to true
-            ctx->finished = k >= b;
+            /* if `next_possible_key` goes out of range,
+             * it means that we have searched all possible keys */
+            ctx->finished = ++ctx->next_possible_key >= b;
             return true;
         }
-    } while ((k = ++ctx->next_possible_key));
+        /* if b == 0 and k == 0, finish searching */
+    } while ((k = ++ctx->next_possible_key) || b != 0);
     ctx->finished = true;
     return false;
 }
@@ -92,17 +107,119 @@ int pkcs7_check_pad(const char *buf, size_t n) {
     return pad;
 }
 
+/* initialized in main function */
+typedef struct s_thread_param {
+    /* search range */
+    uint32_t a;
+    uint32_t b;
+    int worker_id;
+} thread_param;
+
+/* if a worker find this to be true, it will terminate */
+atomic_bool key_found;
+/* shared across workers */
+const char *ciphertext;
+/* should not be modified by workers */
+uint32_t ciphertext_len;
+/* the result generated by a lucky worker */
+volatile struct {
+    char *plaintext;
+    uint32_t len;
+} crack_result;
+
+int thread_worker(thread_param *param) {
+    key_search_ctx ctx;
+    const uint32_t b = param->b; /* search end */
+    uint32_t ciphertext_length = ciphertext_len;
+    new_key_search_ctx(&ctx, *(uint64_t *) ciphertext, param->a);
+    char *plaintext = malloc(ciphertext_length);
+    if (plaintext == NULL) {
+        perror("malloc");
+        return 1;
+    }
+    /* FOR DEBUGGING ONLY */
+//    assert(*(uint64_t *) ciphertext == 8022485120222247589ull);
+//    ctx.next_possible_key = 0xA0979B6Du;
+    while (yield_possible_key(&ctx, b, &key_found)) {
+        /* found a possible correct key */
+        /* validate it by calculating md5 hashsum of the plaintext */
+//        printf("[worker#%d] Possible key: %zu\n", param->worker_id, ctx.yield);
+//        fflush(stdout);
+
+        /* decrypt the whole ciphertext */
+        int err;
+        symmetric_key skey;
+        if ((err = des_setup((const unsigned char *) (&ctx.yield), 8, 0,
+                             &skey)) != CRYPT_OK) {
+            fprintf(stderr, "Err: setup: %s", error_to_string(err));
+            continue;
+        }
+
+        // TODO accelerate by checking the padding at first
+
+        uint_fast32_t blk_cnt = ciphertext_length >> 3;
+        for (uint_fast32_t blk = 0; blk < blk_cnt; ++blk) {
+            des_ecb_decrypt(
+                    (const unsigned char *) ((uint64_t *) ciphertext + blk),
+                    (unsigned char *) ((uint64_t *) plaintext + blk),
+                    (const symmetric_key *) &skey
+            );
+            /* error checking is unnecessary here */
+        }
+
+        int pad_length = pkcs7_check_pad(plaintext, ciphertext_length);
+        const unsigned int unpadded_length = ciphertext_length - pad_length;
+        assert(pad_length < ciphertext_length);
+        if (pad_length < 0) {
+            /* invalid pad, this key is incorrect, skip it */
+            fprintf(stderr, "Invalid pad.\n");
+            continue;
+        }
+
+        /* calculate md5 checksum of the decrypted plaintext */
+        char md5_out[16];
+        hash_state md;
+        md5_init(&md);
+        md5_process(&md, (const unsigned char *) plaintext, unpadded_length);
+        md5_done(&md, (unsigned char *) md5_out);
+
+        /* compare md5_out[0~3] with 8-byte ASCII hex string ctx.yield */
+        /* hex of first 4-byte of md5_out,
+         * 1 more byte to hold the '\0' terminator */
+        char md5_hex[8 + 1];
+        snprintf(md5_hex, 8 + 1, "%02X%02X%02X%02X",
+                 md5_out[0] & 0xFFu, md5_out[1] & 0xFFu,
+                 md5_out[2] & 0xFFu, md5_out[3] & 0xFFu);
+        if (!memcmp(md5_hex, (const char *) (&ctx.yield), 8)) {
+            atomic_store(&key_found, true);
+            printf("[+] FOUND KEY: %zu\n", ctx.yield);
+            crack_result.plaintext = plaintext;
+            crack_result.len = unpadded_length;
+            return 0;
+        }
+        /* otherwise the key is incorrect, continue searching */
+    }
+    /* either key is not found, or another worker has found the key */
+    return 0;
+}
+
 int main(int argc, char *argv[]) {
 //    uint64_t ciphertext = 8022485120222247589;
 //    unsigned char plaintext[8];
 //    const char *key = "A0979B6D";
 //    symmetric_key skey;
-    FILE *fp;
 
-    if (argc != 2 && argc != 3) {
-        printf("Usage: %s <fp_file> [<where_to_save_the_decrypted_file>]\n"
+    crack_result.plaintext = NULL;
+
+    if (argc == 1) {
+        USAGE:
+        printf("Usage: %s <fp_file> "
+               "[<where_to_save_the_decrypted_file>] "
+               "[-j <threads>]\n"
                "The decrypted image won't be saved if "
-               "save path is not specified.\n",
+               "save path is not specified.\n"
+               "threads: how many workers to run at the same time, "
+               "default: 1\n",
                argv[0]);
         return 0;
     }
@@ -110,13 +227,14 @@ int main(int argc, char *argv[]) {
     const char *plaintext_save_path = (argc == 3) ? (argv[2]) : NULL;
     const char *ciphertext_file_path = argv[1];
 
-    // open file
+    /* open file */
+    FILE *fp;
     if (!(fp = fopen(ciphertext_file_path, "rb"))) {
         perror("fopen");
         return 1;
     }
 
-    // test file header
+    /* validate file header */
     char header[8];
     if (fread(header, 1, 8, fp) != 8) {
         fprintf(stderr, "Cannot read first 8 bytes from file.\n");
@@ -127,7 +245,7 @@ int main(int argc, char *argv[]) {
         return 1;
     }
 
-    // read ciphertext into memory
+    /* read ciphertext into memory */
     fseek(fp, 0, SEEK_END);
     long file_length = ftell(fp);
     if (file_length <= 8) {
@@ -141,96 +259,101 @@ int main(int argc, char *argv[]) {
                 file_length);
         return 1;
     }
-    char *ciphertext = malloc(ciphertext_length);
+    char *ciphertext_buf = malloc(ciphertext_length);
     /* this buffer is for the future decryption usage,
      * storing padded plaintext (pkcs7) */
     char *plaintext = malloc(ciphertext_length);
-    if (ciphertext == NULL || plaintext == NULL) {
+    if (ciphertext_buf == NULL || plaintext == NULL) {
         perror("malloc");
         return 1;
     }
 
     fseek(fp, 8, SEEK_SET);
-    if (fread(ciphertext, 1, ciphertext_length, fp) != ciphertext_length) {
+    if (fread(ciphertext_buf, 1, ciphertext_length, fp) != ciphertext_length) {
         fprintf(stderr, "Cannot read the whole file.\n");
         return 1;
     }
 
-    /* start searching */
-    printf("Searching key...\n");
-    fflush(stdout);
+    ciphertext = ciphertext_buf;
+    ciphertext_len = ciphertext_length;
 
-    key_search_ctx ctx;
-    new_key_search_ctx(&ctx, *(uint64_t *) ciphertext);
-    /* FOR DEBUGGING ONLY */
-//    assert(*(uint64_t *) ciphertext == 8022485120222247589ull);
-//    ctx.next_possible_key = 0xA0979B6Du;
-    while (yield_possible_key(&ctx)) {
-        /* found a possible correct key */
-        /* validate it by calculating md5 hashsum of the plaintext */
-        printf("Possible key: %zu\n", ctx.yield);
-        fflush(stdout);
+    int threads = 1;
 
-        /* decrypt the whole ciphertext */
-        int err;
-        symmetric_key skey;
-        if ((err = des_setup((const unsigned char *) (&ctx.yield), 8, 0,
-                             &skey)) != CRYPT_OK) {
-            fprintf(stderr, "Err: setup: %s", error_to_string(err));
-            continue;
+    /* read thread count from argv */
+    for (int i = 1; i < argc; ++i) {
+        if (!strcmp(argv[i], "-j")) {
+            if (i == argc - 1) {
+                printf("-j requires an integer parameter.\n");
+                goto USAGE;
+            }
+            errno = 0;
+            threads = strtol(argv[i + 1], NULL, 10);
+            if (errno) {
+                printf("Invalid thread count number.\n");
+                goto USAGE; /* invalid integer */
+            }
+            break;
         }
+    }
 
-        uint_fast32_t blk_cnt = ciphertext_length >> 3;
-        for (uint_fast32_t blk = 0; blk < blk_cnt; ++blk) {
-            des_ecb_decrypt(
-                    (const unsigned char *) ((uint64_t *) ciphertext + blk),
-                    (unsigned char *) ((uint64_t *) plaintext + blk),
-                    (const symmetric_key *) &skey
-            );
-            /* error checking is unnecessary here */
-        }
+    /* start searching */
+    printf("Searching key (using %d workers)...\n", threads);
+    fflush(stdout);
 
-        int pad_length = pkcs7_check_pad(plaintext, ciphertext_length);
-        const unsigned int unpadded_length = ciphertext_length - pad_length;
-        assert(pad_length < ciphertext_length);
-        if (pad_length < 0) {
-            /* invalid pad, this key is incorrect, skip it */
-            fprintf(stderr, "Invalid pad.\n");
-            continue;
-        }
+    atomic_store(&key_found, false);
 
-        /* calculate md5 checksum of the decrypted plaintext */
-        char md5_out[16];
-        hash_state md;
-        md5_init(&md);
-        md5_process(&md, (const unsigned char *) plaintext, unpadded_length);
-        md5_done(&md, (unsigned char *) md5_out);
+    thrd_t *thread_ids;
+    thread_param *thread_params;
+    if ((thread_ids = malloc(sizeof(thrd_t) * threads)) == NULL) {
+        perror("malloc");
+        return 1;
+    }
+    if ((thread_params = malloc(sizeof(thread_param) * threads)) == NULL) {
+        perror("malloc");
+        return 1;
+    }
 
-        /* compare md5_out[0~3] with 8-byte ASCII hex string ctx.yield */
-        /* hex of first 4-byte of md5_out,
-         * 1 more byte to hold the '\0' terminator */
-        char md5_hex[8 + 1];
-        snprintf(md5_hex, 8 + 1, "%02X%02X%02X%02X",
-                 md5_out[0] & 0xFFu, md5_out[1] & 0xFFu,
-                 md5_out[2] & 0xFFu, md5_out[3] & 0xFFu);
-        if (!memcmp(md5_hex, (const char *) (&ctx.yield), 8)) {
-            printf("[+] FOUND KEY: %zu\n", ctx.yield);
+    /* assign search ranges to workers */
+    uint32_t range_size = 0xFFFFFFFFu / threads;
+    for (int i = 0; i < threads; ++i) {
+        thread_params[i].a = range_size * i;
+        thread_params[i].b = range_size * i + range_size;
+        thread_params[i].worker_id = i;
+    }
+    /* the last search range should warp */
+    thread_params[threads - 1].b = 0;
 
-            if (plaintext_save_path) {
-                FILE *fout = fopen(plaintext_save_path, "wb");
-                if (!fout) {
-                    perror("Cannot fopen for saving");
-                    return 1;
-                }
-                fwrite(plaintext, 1, unpadded_length, fout);
-                fclose(fout);
-                printf("Flash photo has been saved in: %s\n",
-                       plaintext_save_path);
-            }
+    /* start workers */
+    for (int i = 0; i < threads; ++i) {
+        if (thrd_create(
+                &thread_ids[i],
+                (thrd_start_t) thread_worker,
+                &thread_params[i]) != thrd_success) {
+            fprintf(stderr, "Cannot start thread %d.\n", i);
+            return 1;
+        }
+    }
 
-            return 0;
+    /* wait for all workers to terminate */
+    for (int i = 0; i < threads; ++i) {
+        int ret;
+        thrd_join(thread_ids[i], &ret);
+        if (ret) {
+            fprintf(stderr, "Worker terminated with error code %d.\n", ret);
         }
-        /* otherwise the key is incorrect, continue searching */
+    }
+
+    /* save decrypted data */
+    if (crack_result.plaintext != NULL && plaintext_save_path) {
+        FILE *fout = fopen(plaintext_save_path, "wb");
+        if (!fout) {
+            perror("Cannot open file for saving");
+            return 1;
+        }
+        fwrite(crack_result.plaintext, 1, crack_result.len, fout);
+        fclose(fout);
+        printf("Flash photo has been saved in: %s\n",
+               plaintext_save_path);
     }
 
     return 0;