1 files changed, 257 insertions, 87 deletions

diff --git a/utils/nwztools/upgtools/keysig_search.c b/utils/nwztools/upgtools/keysig_search.c
index 7652efa233..2740dc1461 100644
--- a/utils/nwztools/upgtools/keysig_search.c
+++ b/utils/nwztools/upgtools/keysig_search.c
@@ -23,6 +23,118 @@
 #include "mg.h"
 #include <string.h>
 #include <stdio.h>
+#include <stdlib.h>
+#include <pthread.h>
+
+/** Generic search code */
+
+/* The generator sends chunks to the workers. The exact type of chunks depends
+ * on the method used. */
+static struct
+{
+    pthread_mutex_t mutex; /* mutex for the whole structure */
+    pthread_cond_t avail_cond; /* condition to signal available or stop */
+    pthread_cond_t req_cond; /* condition to signal request or stop */
+    bool stop; /* if true, stop searcg */
+    void *chunk; /* pointer to chunk (NULL if not available) */
+    size_t chunk_sz; /* chunk size */
+}g_producer;
+
+/* init producer */
+static void producer_init(void)
+{
+    pthread_cond_init(&g_producer.avail_cond, NULL);
+    pthread_cond_init(&g_producer.req_cond, NULL);
+    pthread_mutex_init(&g_producer.mutex, NULL);
+    g_producer.stop = false;
+    g_producer.chunk = NULL;
+    g_producer.chunk_sz = 0;
+}
+
+/* consumer get: called by worker to get a new chunk, return NULL to stop */
+static void *consumer_get(size_t *sz)
+{
+    pthread_mutex_lock(&g_producer.mutex);
+    /* loop until stop or new chunk */
+    while(true)
+    {
+        /* stop if requested */
+        if(g_producer.stop)
+        {
+            pthread_mutex_unlock(&g_producer.mutex);
+            return NULL;
+        }
+        if(g_producer.chunk)
+            break;
+        /* request a new chunk */
+        pthread_cond_signal(&g_producer.req_cond);
+        /* wait for availability */
+        pthread_cond_wait(&g_producer.avail_cond, &g_producer.mutex);
+    }
+    void *c = g_producer.chunk;
+    if(sz)
+        *sz = g_producer.chunk_sz;
+    g_producer.chunk = NULL;
+    pthread_mutex_unlock(&g_producer.mutex);
+    /* request a new chunk, so that if other consumers are waiting, the producer
+     * will also wake them up */
+    pthread_cond_signal(&g_producer.req_cond);
+    return c;
+}
+
+/* stop: called by worker to stop the search */
+static void consumer_stop(void)
+{
+    pthread_mutex_lock(&g_producer.mutex);
+    /* set stop */
+    g_producer.stop = true;
+    /* wake up everyone */
+    pthread_cond_broadcast(&g_producer.req_cond);
+    pthread_cond_broadcast(&g_producer.avail_cond);
+    pthread_mutex_unlock(&g_producer.mutex);
+}
+
+/* producer yield: called by generator to give a new chunk, return true to stop */
+static bool producer_yield(void *chunk, size_t sz)
+{
+    pthread_mutex_lock(&g_producer.mutex);
+    /* wait until stop or request */
+    while(true)
+    {
+        /* stop if requested */
+        if(g_producer.stop)
+        {
+            pthread_mutex_unlock(&g_producer.mutex);
+            return true;
+        }
+        /* if the chunk is empty, fill it */
+        if(g_producer.chunk == NULL)
+            break;
+        /* otherwise wait for request */
+        pthread_cond_wait(&g_producer.req_cond, &g_producer.mutex);
+    }
+    g_producer.chunk = malloc(sz);
+    memcpy(g_producer.chunk, chunk, sz);
+    g_producer.chunk_sz = sz;
+    /* signal availability */
+    pthread_cond_signal(&g_producer.avail_cond);
+    pthread_mutex_unlock(&g_producer.mutex);
+    return false;
+}
+
+static void producer_stop(void)
+{
+    pthread_mutex_lock(&g_producer.mutex);
+    /* if we are not already stopping and there is a chunk still waiting to
+     * be consumed, wait until next request */
+    if(!g_producer.stop && g_producer.chunk)
+        pthread_cond_wait(&g_producer.req_cond, &g_producer.mutex);
+    /* set stop */
+    g_producer.stop = true;
+    /* wake up everyone */
+    pthread_cond_broadcast(&g_producer.avail_cond);
+    pthread_mutex_unlock(&g_producer.mutex);
+}
 
 /* Key search methods
  *
@@ -41,16 +153,18 @@
  * towards very unbalanced strings (only digits or only letters).
  */
 
-static uint8_t g_cipher[8];
-static keysig_notify_fn_t g_notify;
-static uint8_t g_key[8];
-static void *g_user;
+static struct
+{
+    pthread_mutex_t mutex;
+    uint8_t *enc_buf;
+    size_t enc_buf_sz;
+    bool found_keysig;
+    uint8_t key[NWZ_KEY_SIZE]; /* result */
+    uint8_t sig[NWZ_SIG_SIZE]; /* result */
+}g_keysig_search;
+
 static bool is_hex[256];
 static bool is_init = false;
-static uint64_t g_tot_count;
-static uint64_t g_cur_count;
-static int g_last_percent;
-static int g_last_subpercent;
 
 static void keysig_search_init()
 {
@@ -71,105 +185,169 @@ static inline bool is_full_ascii(uint8_t *arr)
     return true;
 }
 
-static inline bool check_stupid()
+struct upg_header_t
 {
-    uint8_t res[8];
-    // display progress
-    g_cur_count++;
-    int percent = (g_cur_count * 100ULL) / g_tot_count;
-    int subpercent = ((g_cur_count * 1000ULL) / g_tot_count) % 10;
-    if(percent != g_last_percent)
-    {
-        cprintf(RED, "%d%%", percent);
-        fflush(stdout);
-        g_last_subpercent = 0;
-    }
-    if(subpercent != g_last_subpercent)
+    uint8_t sig[NWZ_SIG_SIZE];
+    uint32_t nr_files;
+    uint32_t pad; // make sure structure size is a multiple of 8
+} __attribute__((packed));
+
+static bool check_key(uint8_t key[NWZ_KEY_SIZE])
+{
+    struct upg_header_t hdr;
+    mg_decrypt_fw(g_keysig_search.enc_buf, sizeof(hdr.sig), (void *)&hdr, key);
+    if(is_full_ascii(hdr.sig))
     {
-        cprintf(WHITE, ".");
-        fflush(stdout);
+        /* the signature looks correct, so decrypt the header futher to be sure */
+        mg_decrypt_fw(g_keysig_search.enc_buf, sizeof(hdr), (void *)&hdr, key);
+        /* we expect the number of files to be small and the padding to be 0 */
+        if(hdr.nr_files == 0 || hdr.nr_files > 10 || hdr.pad != 0)
+            return false;
+        cprintf(RED, "    Found key: %.8s (sig=%.8s, nr_files=%u)\n", key, hdr.sig, (unsigned)hdr.nr_files);
+        pthread_mutex_lock(&g_keysig_search.mutex);
+        g_keysig_search.found_keysig = true;
+        memcpy(g_keysig_search.key, key, NWZ_KEY_SIZE);
+        memcpy(g_keysig_search.sig, hdr.sig, NWZ_SIG_SIZE);
+        pthread_mutex_unlock(&g_keysig_search.mutex);
+        consumer_stop();
+        return true;
     }
-    g_last_percent = percent;
-    g_last_subpercent = subpercent;
-
-    mg_decrypt_fw(g_cipher, 8, res, g_key);
-    if(is_full_ascii(res))
-        return g_notify(g_user, g_key, res);
     return false;
 }
 
-static bool search_stupid_rec(int rem_digit, int rem_letter, int pos)
+/** Hex search */
+
+struct hex_chunk_t
 {
-    if(pos == 8)
-        return check_stupid();
-    if(rem_digit > 0)
+    uint8_t key[NWZ_KEY_SIZE]; /* partially pre-filled key */
+    int pos;
+    int rem_letters;
+    int rem_digits;
+};
+
+static bool hex_rec(bool producer, struct hex_chunk_t *ch)
+{
+    /* we list the first 4 pos in generator, and remaining 4 in workers */
+    if(producer && ch->pos == 4)
+    {
+        //printf("yield(%.8s,%d,%d,%d)\n", ch->key, ch->pos, ch->rem_digits, ch->rem_letters);
+        return producer_yield(ch, sizeof(struct hex_chunk_t));
+    }
+    /* filled the key ? */
+    if(!producer && ch->pos == NWZ_KEY_SIZE)
+        return check_key(ch->key);
+    /* list next possibilities
+     *
+     * NOTE (42) Since the cipher is DES, the key is actually 56-bit: the least
+     * significant bit of each byte is an (unused) parity bit. We thus only
+     * generate keys where the least significant bit is 0. */
+    int p = ch->pos++;
+    if(ch->rem_digits > 0)
     {
-        for(int i = '0'; i <= '9'; i++)
+        ch->rem_digits--;
+        /* NOTE (42) */
+        for(int i = '0'; i <= '9'; i += 2)
         {
-            g_key[pos] = i;
-            if(search_stupid_rec(rem_digit - 1, rem_letter, pos + 1))
+            ch->key[p] = i;
+            if(hex_rec(producer, ch))
                 return true;
         }
+        ch->rem_digits++;
     }
-    if(rem_letter > 0)
+    if(ch->rem_letters > 0)
     {
-        for(int i = 'a' - 1; i <= 'f'; i++)
+        ch->rem_letters--;
+        /* NOTE (42) */
+        for(int i = 'a'; i <= 'f'; i += 2)
         {
-            g_key[pos] = i;
-            if(search_stupid_rec(rem_digit, rem_letter - 1, pos + 1))
+            ch->key[p] = i;
+            if(hex_rec(producer, ch))
                 return true;
         }
+        ch->rem_letters++;
     }
+    ch->pos--;
     return false;
 }
 
-static bool search_stupid(int rem_digit, int rem_letter)
+static void *hex_worker(void *arg)
 {
-    cprintf(WHITE, "\n    Looking for keys with ");
-    cprintf(RED, "%d", rem_digit);
-    cprintf(WHITE, " digits and ");
-    cprintf(RED, "%d", rem_letter);
-    cprintf(WHITE, " letters...");
-    fflush(stdout);
-    return search_stupid_rec(rem_digit, rem_letter, 0);
+    (void) arg;
+    while(true)
+    {
+        struct hex_chunk_t *ch = consumer_get(NULL);
+        if(ch == NULL)
+            break;
+        hex_rec(false, ch);
+    }
+    return NULL;
 }
 
-bool keysig_search_ascii_stupid(uint8_t *cipher, keysig_notify_fn_t notify, void *user)
+static bool hex_producer_list(int nr_digits, int nr_letters)
 {
-    keysig_search_init();
-    memcpy(g_cipher, cipher, 8);
-    g_notify = notify;
-    g_user = user;
-    // compute number of possibilities
-    g_cur_count = 0;
-    g_tot_count = 1;
-    g_last_percent = -1;
-    for(int i = 0; i < 8; i++)
-        g_tot_count *= 16ULL;
-    cprintf(WHITE, "    Search space:");
-    cprintf(RED, " %llu", (unsigned long long)g_tot_count);
+    struct hex_chunk_t ch;
+    cprintf(BLUE, "    Listing keys with %d letters and %d digits\n", nr_letters,
+        nr_digits);
+    memset(ch.key, ' ', 8);
+    ch.pos = 0;
+    ch.rem_letters = nr_letters;
+    ch.rem_digits = nr_digits;
+    return hex_rec(true, &ch);
+}
+
+void *hex_producer(void *arg)
+{
+    (void) arg;
     // sorted by probability:
-    bool ret = search_stupid(5, 3) // 5 digits, 3 letters: 0.281632
-        || search_stupid(6, 2) // 6 digits, 2 letters: 0.234693
-        || search_stupid(4, 4) // 4 digits, 4 letters: 0.211224
-        || search_stupid(7, 1) // 7 digits, 1 letters: 0.111759
-        || search_stupid(3, 5) // 3 digits, 5 letters: 0.101388
-        || search_stupid(2, 6) // 2 digits, 6 letters: 0.030416
-        || search_stupid(8, 0) // 8 digits, 0 letters: 0.023283
-        || search_stupid(1, 7) // 1 digits, 7 letters: 0.005214
-        || search_stupid(0, 8);// 0 digits, 8 letters: 0.000391
-    cprintf(OFF, "\n");
-    return ret;
+    bool stop = hex_producer_list(5, 3) // 5 digits, 3 letters: 0.281632
+        || hex_producer_list(6, 2) // 6 digits, 2 letters: 0.234693
+        || hex_producer_list(4, 4) // 4 digits, 4 letters: 0.211224
+        || hex_producer_list(7, 1) // 7 digits, 1 letters: 0.111759
+        || hex_producer_list(3, 5) // 3 digits, 5 letters: 0.101388
+        || hex_producer_list(2, 6) // 2 digits, 6 letters: 0.030416
+        || hex_producer_list(8, 0) // 8 digits, 0 letters: 0.023283
+        || hex_producer_list(1, 7) // 1 digits, 7 letters: 0.005214
+        || hex_producer_list(0, 8);// 0 digits, 8 letters: 0.000391
+    if(!stop)
+        producer_stop();
+    return NULL;
 }
 
-bool keysig_search_ascii_brute(uint8_t *cipher, keysig_notify_fn_t notify, void *user)
+typedef void *(*routine_t)(void *);
+
+bool keysig_search(int method, uint8_t *enc_buf, size_t buf_sz,
+    keysig_notify_fn_t notify, void *user, int nr_threads)
 {
-    (void) cipher;
-    (void) notify;
-    (void) user;
+    /* init producer */
+    producer_init();
+    /* init search */
     keysig_search_init();
-    cprintf(RED, "Unimplemented\n");
-    return false;
+    pthread_mutex_init(&g_keysig_search.mutex, NULL);
+    g_keysig_search.enc_buf = enc_buf;
+    g_keysig_search.enc_buf_sz = buf_sz;
+    g_keysig_search.found_keysig = false;
+    /* get methods */
+    routine_t worker_fn = NULL;
+    routine_t producer_fn = NULL;
+    if(method == KEYSIG_SEARCH_ASCII_HEX)
+    {
+        worker_fn = hex_worker;
+        producer_fn = hex_producer;
+    }
+    /* create workers */
+    pthread_t *worker = malloc(sizeof(pthread_t) * nr_threads);
+    pthread_t producer;
+    for(int i = 0; i < nr_threads; i++)
+        pthread_create(&worker[i], NULL, worker_fn, NULL);
+    pthread_create(&producer, NULL, producer_fn, NULL);
+    /* wait for all threads */
+    pthread_join(producer, NULL);
+    for(int i = 0; i < nr_threads; i++)
+        pthread_join(worker[i], NULL);
+    free(worker);
+    if(g_keysig_search.found_keysig)
+        notify(user, g_keysig_search.key, g_keysig_search.sig);
+    return g_keysig_search.found_keysig;
 }
 
 struct keysig_search_desc_t keysig_search_desc[KEYSIG_SEARCH_LAST] =
@@ -177,19 +355,11 @@ struct keysig_search_desc_t keysig_search_desc[KEYSIG_SEARCH_LAST] =
     [KEYSIG_SEARCH_NONE] =
     {
         .name = "none",
-        .fn = NULL,
         .comment = "don't use",
     },
-    [KEYSIG_SEARCH_ASCII_STUPID] =
+    [KEYSIG_SEARCH_ASCII_HEX] =
     {
         .name = "ascii-hex",
-        .fn = keysig_search_ascii_stupid,
         .comment = "Try to find an hexadecimal ascii string keysig"
     },
-    [KEYSIG_SEARCH_ASCII_BRUTE] =
-    {
-        .name = "ascii-brute",
-        .fn = keysig_search_ascii_brute,
-        .comment = "Brute force all ASCII keys"
-    },
 };