// Compile with -O2 -std=c++17
#include <algorithm>
#include <array>
#include <cassert>
#include <cstdint>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <variant>
#include <string>
#include <tuple>
#include <vector>
#include <random>

#include <sys/mman.h>

using u8 = uint8_t;
using u16 = uint16_t;
using u32 = uint32_t;
using u64 = uint64_t;

using std::array;
using std::cin;
using std::cout;
using std::endl;
using std::get;
using std::holds_alternative;
using std::string;
using std::tuple;
using std::vector;

array<u8, 4> pack32(u32 val)
{
    return { u8(val), u8(val >> 8), u8(val >> 16), u8(val >> 24) };
}

array<u8, 2> pack16(u16 val)
{
    return { u8(val & 0xFF), u8(val >> 8) };
}

u16 unpack16(array<u8, 2> bytes)
{
    return bytes[0] + bytes[1] * 256;
}

array<u8, 8> pack_4x16(u16 a, u16 b, u16 c, u16 d)
{
    array<u8, 8> res;
    size_t pos = 0;
    for (auto& x : { pack16(a), pack16(b), pack16(c), pack16(d) })
    {
        std::copy(x.begin(), x.end(), res.begin() + pos);
        pos += 2;
    }
    return res;
}

std::tuple<u16, u16, u16, u16> unpack_4x16(const array<u8, 8>& in)
{
    return std::make_tuple(unpack16({ in[0], in[1] }),
                           unpack16({ in[2], in[3] }),
                           unpack16({ in[4], in[5] }),
                           unpack16({ in[6], in[7] }));
}

u8 rol8(u8 x, int shift)
{
    shift %= 8;
    return ((x << shift) & 0xFF) | (x >> (8 - shift));
}

u8 ror8(u8 x, int shift)
{
    return rol8(x, 8 - (shift % 8));
}

array<u8, 8> xor8(const array<u8, 8>& a, const array<u8, 8>& b)
{
    array<u8, 8> res;
    for (int i = 0; i < 8; i++)
        res[i] = a[i] ^ b[i];
    return res;
}

struct round_key
{
    u16 a, b;
    array<u8, 4> shuffle;
};

struct key
{
    array<u8, 256> xorbox;
    array<round_key, 16> round_keys;

    key(const array<u8, 256> k1, const array<u8, 128> k2)
    {
        xorbox = k1;
        for (int i = 0; i < 16; i++)
        {
            auto k2_it = &k2[i * sizeof(round_key)];
            round_keys[i].a = unpack16({ *(k2_it + 0), *(k2_it + 1) });
            round_keys[i].b = unpack16({ *(k2_it + 2), *(k2_it + 3) });
            memcpy(&round_keys[i].shuffle, k2_it + 4, 4);
        }
    }
};

std::tuple<u16, u16, u16, u16> tsb_cipher_round(u16 a, u16 b, u16 c, u16 d,
                                                const array<u8, 256> xorbox,
                                                const array<u8, 4> shuffle)
{
    array<u8, 8> block = pack_4x16(a, b, c, d);
    u8 x1 = rol8(xorbox[shuffle[1] ^ block[1] ^ block[3] ^ block[5] ^ block[7]], 3);
    u8 x2 = rol8(xorbox[shuffle[0] ^ x1 ^ block[0] ^ block[2]], 1);
    u8 x3 = ror8(xorbox[shuffle[2] ^ block[0] ^ block[2] ^ block[4] ^ block[6]], 3);
    u8 x4 = ror8(xorbox[shuffle[3] ^ x3 ^ block[5] ^ block[7]], 1);
    array<u8, 8> xors = { x1, x4, x1, x4, x2, x3, x2, x3 };
    for (int i = 0; i < 8; i++)
        block[i] ^= xors[i];
    return unpack_4x16(block);
}

array<u8, 8> tsb_encrypt_block(const array<u8, 8>& in, const key& k)
{
    u16 a, b, c, d;
    std::tie(a, b, c, d) = unpack_4x16(in);

    for (const round_key& round_key : k.round_keys)
    {
        std::tie(a, b, c, d) = std::make_tuple(d, a, b, c);
        b += round_key.a;
        d += round_key.b;
        std::tie(a, b, c, d) = tsb_cipher_round(a, b, c, d, k.xorbox, round_key.shuffle);
    }
    return pack_4x16(d, a, b, c);
}

array<u8, 8> tsb_decrypt_block(const array<u8, 8>& in, const key& k)
{
    u16 a, b, c, d;
    std::tie(a, b, c, d) = unpack_4x16(in);

    for (int i = k.round_keys.size() - 1; i >= 0; i--)
    {
        const round_key& round_key = k.round_keys[i];
        std::tie(a, b, c, d) = std::make_tuple(b, c, d, a);
        std::tie(a, b, c, d) = tsb_cipher_round(a, b, c, d, k.xorbox, round_key.shuffle);
        b -= round_key.a;
        d -= round_key.b;
    }
    return pack_4x16(b, c, d, a);
}

array<u8, 8> cbc_mac(vector<u8> data, const key& key)
{
    // Add padding
    u8 pad = 8 - data.size() % 8;
    for (int i = 0; i < pad; i++)
        data.push_back(pad);
    // Calculate tag
    array<u8, 8> res = {};
    for (size_t i = 0; i < data.size(); i += 8)
    {
        array<u8, 8> block;
        std::copy_n(data.begin() + i, 8, block.begin());
        res = xor8(res, block);
        res = tsb_encrypt_block(res, key);
    }
    return res;
}

int nibble_to_int(char nibble)
{
    if (isdigit(nibble))
        return nibble - '0';
    if (isupper(nibble))
        return nibble - 'A' + 10;
    if (islower(nibble))
        return nibble - 'a' + 10;
    return -1;
}

vector<u8> read_from_hex()
{
    vector<u8> res;
    string hex;
    std::getline(cin, hex);
    for (size_t i = 0; i < hex.size(); i += 2)
    {
        int a = nibble_to_int(hex[i]);
        int b = nibble_to_int(hex[i + 1]);
        if (a == -1 || b == -1)
            break;
        res.push_back(a * 16 + b);
    }
    return res;
}

void compile(const key& key)
{
    // Only basic operations for now, we'll add more in the next version.
    // Example: 100+5^123
    cout << "Expression: " << std::flush;
    string expr;
    bool ok = false;
    while (!ok)
    {
        std::getline(cin, expr);
        if (std::any_of(expr.begin(), expr.end(), [](char x) {return !isspace(x); }))
            ok = true;
    }

    std::vector<std::variant<char, u32>> tokens;
    for (char ch : expr)
    {
        if (isspace(ch))
            continue;
        if (isdigit(ch))
        {
            if (!tokens.size() || holds_alternative<char>(tokens.back()))
                tokens.push_back(u32(0));
            auto& val = get<u32>(tokens.back());
            val *= 10;
            val += ch - '0';
        }
        else if (tokens.size() && holds_alternative<u32>(tokens.back()) && ch == '+')
        {
            tokens.push_back('+');
        }
        else if (tokens.size() && holds_alternative<u32>(tokens.back()) && ch == '^')
        {
            tokens.push_back('^');
        }
        else
        {
            cout << "Invalid expression!" << endl;
            return;
        }
    }
    if (!tokens.size() || holds_alternative<char>(tokens.back()))
    {
        cout << "Invalid expression!" << endl;
        return;
    }

    vector<u8> code;

    // mov eax, imm32
    code.push_back(0xB8);
    for (u8 x : pack32(get<u32>(tokens[0])))
        code.push_back(x);

    for (size_t i = 1; i < tokens.size(); i += 2)
    {
        auto op = get<char>(tokens[i]);
        auto arg = get<u32>(tokens[i + 1]);
        if (op == '+')
        {
            // add eax, imm32
            code.push_back(0x48);
            code.push_back(0x05);
        }
        else
        {
            // xor eax, imm32
            assert(op == '^');
            code.push_back(0x35);
        }
        for (u8 x : pack32(arg))
            code.push_back(x);
    }
    // ret
    code.push_back(0xC3);

    auto tag = cbc_mac(code, key);

    cout << "Code: ";
    for (u8 byte : code)
        cout << std::hex << std::setw(2) << std::setfill('0') << (u32)byte;
    cout << endl;

    cout << "Tag: ";
    for (u8 byte : tag)
        cout << std::hex << std::setw(2) << std::setfill('0') << (u32)byte;
    cout << std::dec;
    cout << endl;
}

u32 call_shellcode(vector<u8> shellcode)
{
    u8* payload_addr = (u8*)mmap(nullptr, shellcode.size(), PROT_READ|PROT_WRITE|PROT_EXEC, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
    if (payload_addr == (u8*)-1)
    {
        cout << "mmap() failed, please contact admin" << endl;
        return -1;
    }
    std::copy(shellcode.begin(), shellcode.end(), payload_addr);
    u32 (*payload)();
    payload = (decltype(payload))payload_addr;
    return payload();
}

void test(const key& key)
{
    cout << "Code: " << std::flush;
    vector<u8> code = read_from_hex();
    cout << "Tag: " << std::flush;
    vector<u8> tag = read_from_hex();
    auto correct_tag = cbc_mac(code, key);
    if (tag.size() != 8 || !std::equal(tag.begin(), tag.end(), correct_tag.begin()))
    {
        cout << "Incorrect tag!" << std::endl;
        return;
    }
    // Execute the payload
    u32 ret = call_shellcode(code);
    cout << "Result: " << ret << endl;
}

int main()
{
    std::ios_base::sync_with_stdio(false);

    // Self-test
    array<u8, 256> test_key_a = { 172, 195, 40, 161, 104, 150, 205, 114, 247, 221, 247, 142, 209, 54, 148, 145, 246, 94, 218, 142, 122, 169, 22, 200, 145, 246, 196, 251, 52, 248, 56, 255, 49, 228, 71, 217, 51, 26, 164, 194, 166, 143, 113, 179, 15, 136, 235, 40, 231, 123, 231, 141, 176, 40, 79, 12, 3, 204, 182, 111, 90, 74, 195, 81, 83, 35, 101, 30, 58, 36, 28, 165, 193, 140, 99, 13, 64, 129, 128, 84, 133, 139, 127, 169, 126, 77, 132, 140, 69, 209, 222, 243, 216, 47, 0, 145, 1, 235, 240, 206, 160, 230, 224, 32, 134, 13, 126, 177, 171, 43, 172, 173, 202, 169, 198, 199, 185, 93, 176, 187, 175, 242, 200, 202, 178, 27, 152, 114, 230, 248, 125, 58, 206, 142, 175, 16, 1, 230, 4, 210, 23, 197, 250, 13, 62, 33, 68, 211, 89, 108, 138, 208, 236, 23, 77, 71, 130, 248, 129, 197, 86, 230, 99, 58, 148, 30, 178, 19, 244, 79, 73, 87, 91, 153, 86, 40, 29, 125, 66, 222, 193, 66, 223, 61, 84, 225, 35, 67, 251, 72, 79, 224, 248, 0, 59, 216, 254, 147, 170, 243, 212, 216, 6, 23, 207, 51, 15, 57, 203, 95, 61, 240, 166, 220, 38, 156, 158, 244, 208, 87, 1, 204, 155, 66, 32, 207, 49, 115, 202, 162, 244, 227, 66, 169, 171, 216, 221, 130, 167, 174, 51, 63, 157, 225, 150, 179, 113, 27, 45, 74, 196, 215, 83, 20, 150, 180 };
    array<u8, 128> test_key_b = { 105, 244, 227, 54, 185, 151, 105, 32, 128, 75, 254, 45, 128, 138, 96, 62, 8, 58, 178, 48, 205, 245, 77, 242, 113, 177, 220, 243, 55, 196, 28, 121, 162, 29, 26, 150, 185, 13, 31, 220, 52, 6, 37, 5, 126, 184, 81, 180, 166, 106, 86, 111, 191, 207, 107, 240, 29, 188, 134, 89, 208, 242, 142, 227, 232, 226, 204, 124, 189, 154, 146, 187, 4, 122, 16, 151, 120, 9, 105, 150, 77, 56, 255, 75, 53, 219, 163, 36, 95, 111, 106, 161, 108, 191, 80, 35, 164, 196, 224, 243, 110, 61, 111, 229, 247, 248, 64, 14, 250, 106, 67, 254, 65, 39, 212, 190, 200, 136, 176, 77, 182, 45, 17, 173, 74, 201, 122, 113 };
    array<u8, 8> test_ct = { 48, 49, 50, 51, 52, 53, 54, 55 };
    array<u8, 8> test_pt = { 0x08, 0x11, 0x84, 0xe3, 0xd8, 0x06, 0xc1, 0x56 };
    array<u8, 8> test_out = tsb_decrypt_block(test_ct, key{ test_key_a, test_key_b });
    assert(test_out == test_pt);
    test_out = tsb_encrypt_block(test_pt, key{ test_key_a, test_key_b });
    assert(test_out == test_ct);

    // Initialization
    std::random_device rd;
    array<u8, 256> main_key_a;
    array<u8, 128> main_key_b;
    for (auto& byte : main_key_a)
        byte = rd();
    for (auto& byte : main_key_b)
        byte = rd();
    key main_key { main_key_a, main_key_b };

    cout << "Welcome to our expression compiler!" << endl;

    // Menu
    while (1)
    {
        cout << "1. Compile an expression" << endl;
        cout << "2. Test the result" << endl;
        cout << "3. Exit" << endl;
        string opt;
        std::getline(cin, opt);
        if (opt == "1")
            compile(main_key);
        else if (opt == "2")
            test(main_key);
        else if (opt == "3")
            break;
        else
            cout << "Invalid option!" << endl;
    }
}