#include "aes128.h"

#include <cstring>

namespace roro::meshcore {

namespace {
// The S-box and its inverse, worked out once rather than typed in.
struct Tables {
    uint8_t s[256], inv[256];
    Tables() {
        auto rotl = [](uint8_t x, int n) { return static_cast<uint8_t>(x << n | x >> (8 - n)); };
        uint8_t p = 1, q = 1;
        do {
            p = static_cast<uint8_t>(p ^ (p << 1) ^ (p & 0x80 ? 0x1B : 0));  // p times 3
            q ^= static_cast<uint8_t>(q << 1);                              // q divided by 3
            q ^= static_cast<uint8_t>(q << 2);
            q ^= static_cast<uint8_t>(q << 4);
            if (q & 0x80) q ^= 0x09;
            uint8_t x = static_cast<uint8_t>(q ^ rotl(q, 1) ^ rotl(q, 2) ^ rotl(q, 3) ^ rotl(q, 4) ^ 0x63);
            s[p] = x;
            inv[x] = p;
        } while (p != 1);
        s[0] = 0x63;
        inv[0x63] = 0;
    }
};
const Tables& tables() {
    static const Tables t;
    return t;
}

uint8_t mul(uint8_t a, uint8_t b) {  // in GF(2^8)
    uint8_t r = 0;
    for (; b; b >>= 1) {
        if (b & 1) r ^= a;
        a = static_cast<uint8_t>(a << 1 ^ (a & 0x80 ? 0x1B : 0));
    }
    return r;
}
}  // namespace

Aes128::Aes128(const uint8_t key[16]) {
    const Tables& t = tables();
    std::memcpy(roundKeys_, key, 16);
    uint8_t rcon = 1;
    for (int i = 16; i < 176; i += 4) {
        uint8_t w[4];
        std::memcpy(w, roundKeys_ + i - 4, 4);
        if (i % 16 == 0) {
            uint8_t first = w[0];
            w[0] = static_cast<uint8_t>(t.s[w[1]] ^ rcon);
            w[1] = t.s[w[2]];
            w[2] = t.s[w[3]];
            w[3] = t.s[first];
            rcon = static_cast<uint8_t>(rcon << 1 ^ (rcon & 0x80 ? 0x1B : 0));
        }
        for (int j = 0; j < 4; ++j) roundKeys_[i + j] = roundKeys_[i - 16 + j] ^ w[j];
    }
}

void Aes128::decryptBlock(const uint8_t in[16], uint8_t out[16]) const {
    const Tables& t = tables();
    uint8_t s[16];
    for (int i = 0; i < 16; ++i) s[i] = in[i] ^ roundKeys_[160 + i];
    for (int round = 9; round >= 0; --round) {
        // Rows shifted back and bytes substituted back, in one go: column c, row r came from column c - r.
        uint8_t u[16];
        for (int c = 0; c < 4; ++c)
            for (int r = 0; r < 4; ++r) u[4 * c + r] = t.inv[s[4 * ((c - r + 4) % 4) + r]];
        for (int i = 0; i < 16; ++i) u[i] ^= roundKeys_[16 * round + i];
        if (round == 0) {
            std::memcpy(s, u, 16);
            break;
        }
        for (int c = 0; c < 4; ++c) {
            const uint8_t* a = u + 4 * c;
            s[4 * c + 0] = mul(a[0], 14) ^ mul(a[1], 11) ^ mul(a[2], 13) ^ mul(a[3], 9);
            s[4 * c + 1] = mul(a[0], 9) ^ mul(a[1], 14) ^ mul(a[2], 11) ^ mul(a[3], 13);
            s[4 * c + 2] = mul(a[0], 13) ^ mul(a[1], 9) ^ mul(a[2], 14) ^ mul(a[3], 11);
            s[4 * c + 3] = mul(a[0], 11) ^ mul(a[1], 13) ^ mul(a[2], 9) ^ mul(a[3], 14);
        }
    }
    std::memcpy(out, s, 16);
}

}  // namespace roro::meshcore