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

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import org.pureprogrammer.Utils;

public class ChaCha20Cipher {

    static void printKeyBlock(final List<Integer> keyblock) {
        for (int i = 0; i < 16; ++i) {
            int x = keyblock.get(i);
            for (int j = 0; j < 4; ++j) {
                final String s = Utils.toHex(x & 0xFF);
                final int POS = Utils.cpLength(s) - 2;
                System.out.print(Utils.cpSubstring(s, POS, POS + 2));
                x >>= 8;
            }
            System.out.print(" ");
            if (i % 4 == 3) {
                System.out.println();
            }
        }
        System.out.println();
    }

    static int char2Le(String s) {
        int x = 0;
        for (int i = 0; i < 4; ++i) {
            x <<= 8;
            x |= Utils.cpAt(s, i);
        }
        return x;
    }

    static void keyblock2Bytes(final List<Integer> keyblock, List<Byte> keybytes) {
        for (int i = 0; i < 16; ++i) {
            int x = keyblock.get(i);
            for (int j = 0; j < 4; ++j) {
                keybytes.set(i * 4 + 3 - j, (byte)(x & 0xFF));
                x >>= 8;
            }
        }
    }

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

    static void quarterRound(List<Integer> state, int ai, int bi, int ci, int di) {
        int a = state.get(ai);
        int b = state.get(bi);
        int c = state.get(ci);
        int d = state.get(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.set(ai, a);
        state.set(bi, b);
        state.set(ci, c);
        state.set(di, d);
    }

    static void chacha20Block(final List<Integer> inBlock, List<Integer> out) {
        for (int i = 0; i < 16; ++i) {
            out.set(i, inBlock.get(i));
        }
// 10 loops × 2 rounds/loop = 20 rounds
        for (int i = 0; i < 20; i += 2) {
// Odd round
            quarterRound(out, 0, 4, 8, 12);
            quarterRound(out, 1, 5, 9, 13);
            quarterRound(out, 2, 6, 10, 14);
            quarterRound(out, 3, 7, 11, 15);
// Even round
            quarterRound(out, 0, 5, 10, 15);
            quarterRound(out, 1, 6, 11, 12);
            quarterRound(out, 2, 7, 8, 13);
            quarterRound(out, 3, 4, 9, 14);
        }
        for (int i = 0; i < 16; ++i) {
            out.set(i, out.get(i) + inBlock.get(i));
        }
    }

    static void initBlock(List<Integer> keyblock, String password, String nonce) {
        keyblock.set(0, char2Le("expa"));
        keyblock.set(1, char2Le("nd 3"));
        keyblock.set(2, char2Le("2-by"));
        keyblock.set(3, char2Le("te k"));
        final String KEY = password.repeat(32 / Utils.cpLength(password) + 1);
        for (int i = 0; i < 8; ++i) {
            final int POS = i * 4;
            keyblock.set(i + 4, char2Le(Utils.cpSubstring(KEY, POS, POS + 4)));
        }
        keyblock.set(12, 0);
        keyblock.set(13, 0);
        keyblock.set(14, char2Le(Utils.cpSubstring(nonce, 0, 4)));
        keyblock.set(15, char2Le(Utils.cpSubstring(nonce, 4, 8)));
    }

    static void processFile(String fromFilespec, String toFilespec, String password) throws IOException {
        BufferedOutputStream ofh = new BufferedOutputStream(new FileOutputStream(toFilespec));

        BufferedInputStream ifh = new BufferedInputStream(new FileInputStream(fromFilespec));

        List<Integer> keyblockInitial = new ArrayList<>(Collections.nCopies(16, null));
        List<Integer> keyblock = new ArrayList<>(Collections.nCopies(16, null));
        List<Byte> keybytes = new ArrayList<>(Collections.nCopies(64, null));
        initBlock(keyblockInitial, password, Utils.join("", password, "1234567"));
        chacha20Block(keyblockInitial, keyblock);
        keyblock2Bytes(keyblock, keybytes);
        int count = 0;
        int c;
        while ((c = Utils.getbyte(ifh)) != -1) {
            ofh.write(c ^ keybytes.get(count % 64));
            ++count;
            if (count % 64 == 0) {
// Increment the block count in 12 & 13
                keyblockInitial.set(12, (keyblockInitial.get(12) + 1) & 0xffffffff);
                if (keyblockInitial.get(12) == 0) {
                    keyblockInitial.set(13, (keyblockInitial.get(13) + 1) & 0xffffffff);
                }
                chacha20Block(keyblockInitial, keyblock);
                keyblock2Bytes(keyblock, keybytes);
            }
        }
        try {ofh.close();} catch (IOException ex) {};
        try {ifh.close();} catch (IOException ex) {};
    }

    public static void main(String[] args) {
        if (args.length != 3) {
            System.out.println(Utils.join("", "Syntax: ", "ChaCha20Cipher", " {fromFilespec} {toFilespec} {password}"));
            System.exit(1);
        }
        String fromFilespec = args[0];
        String toFilespec = args[1];
        String password = args[2];

        try {
            processFile(fromFilespec, toFilespec, password);
        } catch (IOException ex) {
            System.out.println(Utils.join("", "Error: ", Utils.exceptionMessage(ex)));
        }
    }
}