Google CTF: solving food the hard way

food.c

/* Made for the Google 2017 CTF
 * Author: Morgan Jones <me at numin dot it>
 *
 * Compile: clang -std=gnu99 -fopenmp -O3 -funroll-loops -fomit-frame-pointer -ofood food.c
 * Run: ./food <start percentage> <end percentage> [num threads=autodetect]
 */

#include <unistd.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <inttypes.h>
#include <stdbool.h>
#include <stdlib.h>
#include <signal.h>

#include <omp.h>

typedef uint64_t u64;
typedef int64_t s64;
typedef uint8_t u8;

const u64 MAX_SEED = (1ULL << 40ULL) - 1;
const u64 CHECKIN_INTERVAL = 10000000;

#define SWAP(_a, _x, _y) \
    do { \
        u8 _v = _a[_x]; \
        _a[_x] = _a[_y]; \
        _a[_y] = _v; \
    } while (0)

static void de(const u8 *ct, size_t ctl, const u8 *key, size_t kl, u8 *pt) {
    u8 magic[256];
    u8 expanded_key[256];

    #pragma omp simd
    for (size_t i = 0; i < 256; i++) {
        magic[i] = (u8)i;
        expanded_key[i] = key[i % kl];
    }

    for (size_t i = 0, j = 0; i < 256; i++) {
        j = ((j + magic[i]) + expanded_key[i]) & 0xff;
        SWAP(magic, i, j);
    }

    for (size_t i = 0, j = 0, k = 0; i < ctl; i++) {
        j = (j + 1) & 0xff;
        k = (k + magic[j]) & 0xff;
        SWAP(magic, j, k);
        pt[i] = (u8)(ct[i] ^ magic[(magic[j] + magic[k]) & 0xff]);
    }
}

static void s2k(u64 seed, u8 *key) {
    // Expand each 5 bits of the seed into the key
    #pragma omp simd
    for (size_t i = 0; i < 8; i++) {
        key[i] = seed & 0x1f;
        seed >>= 5;
    }
}

static void go(u64 start, u64 end) {
    static const u8 enc[] = {
        0xed, 0x74, 0x3a, 0x6c, 0xff, 0x21, 0x09, 0x3d,
        0xc3, 0xdb, 0x6c, 0x85, 0x03, 0x23, 0x61, 0xf6,
        0xf1, 0x0f, 0xab, 0xbe, 0xe1, 0xbf, 0x11, 0x4f,
        0x1f, 0x19, 0xd9, 0x5f, 0x5d, 0x01, 0x92, 0x99,
        0x8a, 0xda, 0xc7, 0xc6, 0xcd, 0xb1
    };
    u8 key[8];
    u8 flag[sizeof(enc)];
    u64 thread_seed = start;
    u64 iterations = 0, max_iterations = 0;

    #pragma omp parallel firstprivate(key, flag, thread_seed, iterations, max_iterations)
    {
        bool first = omp_get_thread_num() == 0;
        bool last = omp_get_thread_num() == omp_get_num_threads() - 1;

        // Compute the max iteration count
        max_iterations = (end - start) / omp_get_num_threads();

        // Start them out with different seeds
        thread_seed = start + max_iterations * omp_get_thread_num();

        if (last) {
            // Make sure to get the leftovers
            max_iterations += (end - start) % omp_get_num_threads();

            if (end == MAX_SEED) {
                // Don't forget the end
                max_iterations++;
            }
        }

        #pragma omp critical(printing)
        {
            printf("[%d:%02d] starting with %d threads; %" PRIu64 " iterations per thread; handling 0x%016" PRIx64 "..0x%016" PRIx64 " (%.2f..%.2f%%)%s\n", getpid(), omp_get_thread_num(), omp_get_num_threads(), max_iterations, thread_seed, thread_seed + max_iterations - 1, (thread_seed / (double)MAX_SEED) * 100, ((thread_seed + max_iterations - 1) / (double)MAX_SEED) * 100, first ? " <== lower bound" : (last ? " <== upper bound" : ""));
            fflush(stdout);
        }

        #pragma omp barrier
        while (iterations < max_iterations) {
            s2k(thread_seed, key);
            de(enc, sizeof(enc), key, sizeof(key), flag);
#ifdef TESTING
            #pragma omp critical(printing)
            {
                printf("Key for seed %" PRIu64 ": ", thread_seed);
                for (size_t i = 0; i < sizeof(key); i++) {
                    printf("%02x ", key[i]);
                }
                printf("\n");
                printf("Flag for seed %" PRIu64 ": ", thread_seed);
                for (size_t i = 0; i < sizeof(flag); i++) {
                    printf("%02x ", flag[i]);
                }
                printf("\n");
                fflush(stdout);
            }
            break;
#endif
            if (flag[0] == 'C' && flag[1] == 'T' && flag[2] == 'F' && flag[3] == '{') {
                #pragma omp critical(printing)
                {
                    printf("[%d:%02d] POSSIBLE FLAG AT SEED 0x%016" PRIx64 ": ", getpid(), omp_get_thread_num(), thread_seed);
                    for (size_t i = 0; i < sizeof(flag); i++) {
                        u8 c = flag[i];
                        printf(c >= 32 && c <= 127 ? "%c" : "[%02x]", c);
                    }
                    printf("\n");
                    fflush(stdout);
                }
            } else if (iterations % CHECKIN_INTERVAL == 0) {
                #pragma omp critical(printing)
                {
                    fprintf(stderr, "[%d:%02d] %.2f%% complete; current seed: 0x%016" PRIx64 "\n", getpid(), omp_get_thread_num(), (iterations / (double)max_iterations) * 100, thread_seed);
                    fflush(stderr);
                }
            }
            thread_seed++;
            iterations++;
        }
    }
}

static void sigint_handler(int sig) {
    if (sig == SIGINT) {
        printf("Interrupted.\n");
        exit(0);
    }
}

int main(int argc, char *const argv[]) {
    if (argc < 3) {
        fprintf(stderr, "Usage: %s <starting keyspace percentage> <ending keyspace percentage> <num threads=autodetect>\n", argv[0]);
        return 1;
    }

    s64 start = (strtod(argv[1], NULL) / 100.0) * MAX_SEED;
    s64 end = (strtod(argv[2], NULL) / 100.0) * MAX_SEED;

    if (start < 0 || end < 0 || start > MAX_SEED || end > MAX_SEED) {
        fprintf(stderr, "start or end are out of range\n");
        return 1;
    } else if (end < start) {
        fprintf(stderr, "start must >= end\n");
        return 1;
    }

    int threads = argc > 3 ? atoi(argv[3]) : 0;
    if (threads > 0) {
        omp_set_num_threads(threads);
    }

    signal(SIGINT, sigint_handler);
    go((u64)start, (u64)end);
    return 0;
}