#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>

void * safe_realloc(void * p, size_t size) {
  p = realloc(p, size);
  if (p == NULL) exit(1);
  return p;
}

/* This has been pre bit reversed */
static const sbox[64] = {0x3, 0xb, 0xd, 0x5, 0x8, 0xd, 0xc, 0xd, 0xe, 0x6, 0x5, 0x0, 0xf, 0xb, 0x4, 0x9, 0xc, 0x9, 0xb, 0x2, 0x3, 0x9, 0xa, 0x0, 0x4, 0x8, 0x1, 0xf, 0x6, 0xa, 0xf, 0x5, 0xe, 0x8, 0x4, 0x7, 0x1, 0x6, 0x3, 0x7, 0x0, 0x4, 0x8, 0x7, 0xc, 0x3, 0xa, 0xe, 0xc, 0x2, 0xa, 0x2, 0x5, 0x6, 0x1, 0xf, 0x0, 0x9, 0x2, 0x7, 0xe, 0xb, 0x1, 0xd };

static uint8_t rf(uint8_t in, uint16_t rk) {
    uint16_t ex = ((uint16_t)(in & 0x07) | ((uint16_t)(in & 0x1c) << 1) | ((uint16_t)(in & 0x70) << 2) | ((uint16_t)(in & 0xc0) << 3) | ((uint16_t)(in & 0x01) << 11)) ^ rk;
    return (sbox[ex & 0x3f] << 4) | sbox[ex >> 6];
}

static uint16_t fre(uint16_t in, uint16_t rk) {
   return (in << 8) | ((in >> 8) ^ ((uint16_t) rf(in, rk)));
}

static uint16_t frd(uint16_t in, uint16_t rk) {
   return ((in ^ ((uint16_t) rf(in >> 8, rk))) << 8) | (in >> 8);
}

static uint16_t enc(uint16_t key, uint16_t data) {
//     fprintf(stderr, "Key: %04hx Data: %04hx\n", key, data);
    key = ((key << 3) & 0xf8f8) | ((key >> 5) & 0x0707);
    data = fre(data, ((key >> 4) & 0x0fc0) | ((key >> 2) & 0x003f));
    key = ((key << 3) & 0xf8f8) | ((key >> 5) & 0x0707);
    data = fre(data, ((key >> 4) & 0x0fc0) | ((key >> 2) & 0x003f));
    key = ((key << 3) & 0xf8f8) | ((key >> 5) & 0x0707);
    data = fre(data, ((key >> 4) & 0x0fc0) | ((key >> 2) & 0x003f));
    key = ((key << 3) & 0xf8f8) | ((key >> 5) & 0x0707);
    data = fre(data, ((key >> 4) & 0x0fc0) | ((key >> 2) & 0x003f));
//     fprintf(stderr, "enc: %04hx\n", data);
    return data;
}

static uint16_t dec(uint16_t dkey, uint16_t ddata) {
//     fprintf(stderr, "Key: %04hx Data: %04hx\n", dkey, ddata);
    dkey = ((dkey << 4) & 0xf0f0) | ((dkey >> 4) & 0x0f0f);
    ddata = frd(ddata, ((dkey >> 4) & 0x0fc0) | ((dkey >> 2) & 0x003f));
    dkey = ((dkey << 5) & 0xe0e0) | ((dkey >> 3) & 0x1f1f);
    ddata = frd(ddata, ((dkey >> 4) & 0x0fc0) | ((dkey >> 2) & 0x003f));
    dkey = ((dkey << 5) & 0xe0e0) | ((dkey >> 3) & 0x1f1f);
    ddata = frd(ddata, ((dkey >> 4) & 0x0fc0) | ((dkey >> 2) & 0x003f));
    dkey = ((dkey << 5) & 0xe0e0) | ((dkey >> 3) & 0x1f1f);
    ddata = frd(ddata, ((dkey >> 4) & 0x0fc0) | ((dkey >> 2) & 0x003f));
//     fprintf(stderr, "dec: %04hx\n", ddata);
}

struct middle {
  size_t count;
  uint16_t * keys;
};

int main (int argc, char ** argv) {
    uint16_t m1, c1, m2, c2;
    uint16_t k1, k2, i;
    uint16_t r;
    int encs = 0, decs = 0, j;
    struct middle a[1<<16];
    if (argc < 4) {
        fprintf(stderr, "Usage: %s m1 c1 m2 c2\n", argv[0]);
        exit(1);
    }
    if (sscanf(argv[1], "%04hx", &m1) != 1) {
        fprintf(stderr, "Couldn't parse m1\n");
        exit(1);
    }
    if (sscanf(argv[2], "%04hx", &c1) != 1) {
        fprintf(stderr, "Couldn't parse c1\n");
        exit(1);
    }
    if (sscanf(argv[3], "%04hx", &m2) != 1) {
        fprintf(stderr, "Couldn't parse m2\n");
        exit(1);
    }
    if (sscanf(argv[4], "%04hx", &c2) != 1) {
        fprintf(stderr, "Couldn't parse c2\n");
        exit(1);
    }

    i = 0;
    do {
      a[i].count = 0;
      a[i].keys = NULL;
      i++;
    } while (i != 0);

    k1 = 0;
    do {
      r = enc(k1, m1);
      encs++;
      a[r].count++;
      a[r].keys = safe_realloc(a[r].keys, a[r].count * sizeof(uint16_t));
      a[r].keys[a[r].count-1] = k1;
      k1++;
    } while (k1 != 0);

    k2 = 0;
    do {
      r = dec(k2, c1);
      decs++;
      for (j = 0; j < a[r].count; j++) {
        k1 = a[r].keys[j];
//      printf("Possible MITM: k1=%04hx k2=%04hx\n", k1, k2);
        if (enc(k2, enc(k1, m2)) == c2) {
          printf("MITM: k1=%04hx k2=%04hx\n", k1, k2);
        }
        encs += 2;
      }
      k2++;
    } while (k2 != 0);
    printf("Enc was run %d times and Dec was run %d times.\n", encs, decs);
    return 0;
}