/******************************************************************************
 * This program implements the ChaCha20 stream cipher.
 * 
 * Copyright © 2021 Richard Lesh.  All rights reserved.
 *****************************************************************************/

#undef NDEBUG

#include "Utils.hpp"
#include <clocale>
#include <codecvt>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>

std::locale utf8loc(std::locale(), new std::codecvt_utf8<wchar_t>);
using namespace Utils;
using namespace std;

void print_key_block(const vector<int> keyblock) noexcept {
    for (int i = 0; i < 16; ++i) {
        int x = keyblock[i];
        for (int j = 0; j < 4; ++j) {
            wstring const s = Utils::toHex<wchar_t>(x & 0xFF);
            int const POS = s.length() - 2;
            wcout << s.substr(POS, 2);
            x >>= 8;
        }
        wcout << L" ";
        if (i % 4 == 3) {
            wcout << endl;
        }
    }
    wcout << endl;
}

int char2_le(wstring s) noexcept {
    int x = 0;
    for (int i = 0; i < 4; ++i) {
        x <<= 8;
        x |= (int)(s[i]);
    }
    return x;
}

void keyblock2_bytes(const vector<int> keyblock, vector<unsigned char> keybytes) noexcept {
    for (int i = 0; i < 16; ++i) {
        int x = keyblock[i];
        for (int j = 0; j < 4; ++j) {
            keybytes[i * 4 + 3 - j] = (unsigned char)(x & 0xFF);
            x >>= 8;
        }
    }
}

int rotl(int a, int shift) noexcept {
    return ((a << shift) | ((unsigned int)(a) >> (32 - shift))) & 0xffffffff;
}

void quarter_round(vector<int> state, int ai, int bi, int ci, int di) noexcept {
    int a = state[ai];
    int b = state[bi];
    int c = state[ci];
    int d = state[di];

    a += b;
    d ^= a;
    d = rotl(d, 16);
    c += d;
    b ^= c;
    b = rotl(d, 12);
    a += b;
    d ^= a;
    d = rotl(d, 8);
    c += d;
    b ^= c;
    b = rotl(d, 7);

    state[ai] = a;
    state[bi] = b;
    state[ci] = c;
    state[di] = d;
}

void chacha20_block(const vector<int> in_block, vector<int> out) noexcept {
    for (int i = 0; i < 16; ++i) {
        out[i] = in_block[i];
    }
// 10 loops × 2 rounds/loop = 20 rounds
    for (int i = 0; i < 20; i += 2) {
// Odd round
        quarter_round(out, 0, 4, 8, 12);
        quarter_round(out, 1, 5, 9, 13);
        quarter_round(out, 2, 6, 10, 14);
        quarter_round(out, 3, 7, 11, 15);
// Even round
        quarter_round(out, 0, 5, 10, 15);
        quarter_round(out, 1, 6, 11, 12);
        quarter_round(out, 2, 7, 8, 13);
        quarter_round(out, 3, 4, 9, 14);
    }
    for (int i = 0; i < 16; ++i) {
        out[i] = out[i] + in_block[i];
    }
}

void init_block(vector<int> keyblock, wstring password, wstring nonce) noexcept {
    keyblock[0] = char2_le(L"expa");
    keyblock[1] = char2_le(L"nd 3");
    keyblock[2] = char2_le(L"2-by");
    keyblock[3] = char2_le(L"te k");
    wstring const KEY = Utils::repeat(password, 32 / password.length() + 1);
    for (int i = 0; i < 8; ++i) {
        int const POS = i * 4;
        keyblock[i + 4] = char2_le(KEY.substr(POS, 4));
    }
    keyblock[12] = 0;
    keyblock[13] = 0;
    keyblock[14] = char2_le(nonce.substr(0, 4));
    keyblock[15] = char2_le(nonce.substr(4, 4));
}

void process_file(wstring from_filespec, wstring to_filespec, wstring password) {
    ofstream ofh(filesystem::path(to_filespec).c_str(), ios_base::binary);
    if (!ofh.good()) {
        throw ios_base::failure(Utils::wchar_to_UTF8(L"Problem opening output file!"));
    }

    ifstream ifh(filesystem::path(from_filespec).c_str(), ios_base::binary);
    if (!ifh.good()) {
        throw ios_base::failure(Utils::wchar_to_UTF8(L"Problem opening input file!"));
    }

    vector<int> keyblock_initial(16);
    vector<int> keyblock(16);
    vector<unsigned char> keybytes(64);
    init_block(keyblock_initial, password, password + wstring(L"1234567"));
    chacha20_block(keyblock_initial, keyblock);
    keyblock2_bytes(keyblock, keybytes);
    int count = 0;
    unsigned char c;
    while (Utils::getbytes(ifh, c)) {
        ofh.put(c ^ keybytes[count % 64]);
        ++count;
        if (count % 64 == 0) {
// Increment the block count in 12 & 13
            keyblock_initial[12] = (keyblock_initial[12] + 1) & 0xffffffff;
            if (keyblock_initial[12] == 0) {
                keyblock_initial[13] = (keyblock_initial[13] + 1) & 0xffffffff;
            }
            chacha20_block(keyblock_initial, keyblock);
            keyblock2_bytes(keyblock, keybytes);
        }
    }
    ofh.close();
    ifh.close();
}

int main(int argc, char **argv) {
    setlocale(LC_ALL, "en_US.UTF-8");
    wcout.imbue(utf8loc);
    wcin.imbue(utf8loc);

    if (argc != 4) {
        wcout << L"Syntax: " << Utils::UTF8_to_wstring(argv[0]) << L" {fromFilespec} {toFilespec} {password}" << endl;
        exit(1);
    }
    wstring from_filespec = Utils::UTF8_to_wstring(argv[1]);
    wstring to_filespec = Utils::UTF8_to_wstring(argv[2]);
    wstring password = Utils::UTF8_to_wstring(argv[3]);

    try {
        process_file(from_filespec, to_filespec, password);
    } catch (ios_base::failure ex) {
        wcout << L"Error: " << Utils::exceptionMessage<wchar_t>(ex) << endl;
    }
    return 0;
}