#!/usr/bin/env perl
use utf8;
###############################################################################
# This program implements the ChaCha20 stream cipher.
# 
# Copyright © 2021 Richard Lesh.  All rights reserved.
###############################################################################

use Error::Simple;
use Nice::Try;
use Utils;
use strict;
use warnings;

sub print_key_block {
    my ($keyblock) = @_;
    for (my $i = 0; $i < 16; ++$i) {
        my $x = $keyblock->[$i];
        for (my $j = 0; $j < 4; ++$j) {
            my $s = Utils::toHex(Utils::bitwiseAnd32($x, 0xff));
            my $POS = length($s) - 2;
            print substr($s, $POS, 2);
            $x = Utils::arshift32($x, 8);
        }
        print " ";
        if ($i % 4 == 3) {
            print "\n";
        }
    }
    print "\n";
}

sub char2_le {
    my ($s) = @_;
    my $x = 0;
    for (my $i = 0; $i < 4; ++$i) {
        $x = Utils::lshift32($x, 8);
        $x = Utils::bitwiseOr32($x, ord(substr($s, $i, 1)));
    }
    return $x;
}

sub keyblock2_bytes {
    my ($keyblock, $keybytes) = @_;
    for (my $i = 0; $i < 16; ++$i) {
        my $x = $keyblock->[$i];
        for (my $j = 0; $j < 4; ++$j) {
            $keybytes->[$i * 4 + 3 - $j] = int(Utils::bitwiseAnd32($x, 0xff));
            $x = Utils::arshift32($x, 8);
        }
    }
}

sub rotl {
    my ($a, $shift) = @_;
    return (Utils::bitwiseOr32((Utils::lshift32($a, $shift)), (Utils::rshift32($a, (32 - $shift))))) & 0xffffffff;
}

sub quarter_round {
    my ($state, $ai, $bi, $ci, $di) = @_;
    my $a = $state->[$ai];
    my $b = $state->[$bi];
    my $c = $state->[$ci];
    my $d = $state->[$di];

    $a += $b;
    $d = Utils::bitwiseXOr32($d, $a);
    $d = rotl($d, 16);
    $c += $d;
    $b = Utils::bitwiseXOr32($b, $c);
    $b = rotl($d, 12);
    $a += $b;
    $d = Utils::bitwiseXOr32($d, $a);
    $d = rotl($d, 8);
    $c += $d;
    $b = Utils::bitwiseXOr32($b, $c);
    $b = rotl($d, 7);

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

sub chacha20_block {
    my ($in_block, $out) = @_;
    for (my $i = 0; $i < 16; ++$i) {
        $out->[$i] = $in_block->[$i];
    }
# 10 loops × 2 rounds/loop = 20 rounds
    for (my $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 (my $i = 0; $i < 16; ++$i) {
        $out->[$i] = $out->[$i] + $in_block->[$i];
    }
}

sub init_block {
    my ($keyblock, $password, $nonce) = @_;
    $keyblock->[0] = char2_le("expa");
    $keyblock->[1] = char2_le("nd 3");
    $keyblock->[2] = char2_le("2-by");
    $keyblock->[3] = char2_le("te k");
    my $KEY = $password x (int(32 / length($password)) + 1);
    for (my $i = 0; $i < 8; ++$i) {
        my $POS = $i * 4;
        $keyblock->[$i + 4] = char2_le(substr($KEY, $POS, 4));
    }
    $keyblock->[12] = 0;
    $keyblock->[13] = 0;
    $keyblock->[14] = char2_le(substr($nonce, 0, 4));
    $keyblock->[15] = char2_le(substr($nonce, 4, 4));
}

sub process_file {
    my ($from_filespec, $to_filespec, $password) = @_;
    my $ofh;
    my $is_good = open($ofh, ">", $to_filespec);
    if (!$is_good) {
        die Error::IOException->new("Problem opening output file!");
    }
    binmode($ofh);

    my $ifh;
    $is_good = open($ifh, "<", $from_filespec);
    if (!$is_good) {
        die Error::IOException->new("Problem opening input file!");
    }
    binmode($ifh);

    my $keyblock_initial = [(undef) x 16];
    my $keyblock = [(undef) x 16];
    my $keybytes = [(undef) x 64];
    init_block($keyblock_initial, $password, $password . "1234567");
    chacha20_block($keyblock_initial, $keyblock);
    keyblock2_bytes($keyblock, $keybytes);
    my $count = 0;
    my $c;
    while (defined($c = Utils::getbyte(*$ifh))) {
        print $ofh chr(Utils::bitwiseXOr32($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);
        }
    }
    close($ofh);
    close($ifh);
}

MAIN:
{
    binmode(STDOUT, ":utf8");
    binmode(STDERR, ":utf8");
    binmode(STDIN, ":utf8");
    foreach my $arg (@ARGV) {
        utf8::decode($arg);
    }
    if (scalar(@ARGV) != 3) {
        print "Syntax: ", "ChaCha20Cipher", " \{fromFilespec\} \{toFilespec\} \{password\}\n";
        exit 1;
    }
    my $from_filespec = $ARGV[0];
    my $to_filespec = $ARGV[1];
    my $password = $ARGV[2];

    try {
        process_file($from_filespec, $to_filespec, $password);
    } catch (Error::IOException $ex) {
        print "Error: ", $ex, "\n";
    }
}