module chacha20;

private import std.string : format;
private import std.range : isInputRange, isForwardRange, ElementType;
private import std.array;

private import bitmanip;
private import endian;

public:

// TODO: Create unittests!!!!

auto chacha20Cipher(R)(R range, ubyte[32] key, ubyte[8] nonce)
    if(isInputRange!R && is(ElementType!R : ubyte))
{
    static struct rangeResult
    {
    private:
        ubyte[32] key;
        ubyte[8] nonce;
        ulong count;
        R range;
        ubyte[] salsaSection;
    public:
        this(R range, in ubyte[32] key, in ubyte[8] nonce)
        {
            this.range = range;
            this.key = key;
            this.nonce = nonce;
            this.count = 0;
            salsaSection = salsa20Exp(key, nonce ~ littleEndianInv(count));
        }

        bool empty() @property
        {
            return range.empty || (count == ulong.max && salsaSection.empty);
        }

        ubyte front() @property
        {
            assert(!empty);
            return range.front ^ salsaSection.front;
        }

        void popFront()
        {
            assert(!empty);
            salsaSection.popFront();
            if(salsaSection.empty)
            {
                ++count;
                salsaSection = salsa20Exp(key, nonce ~ littleEndianInv(count));
            }
            range.popFront();
        }
        static if(isForwardRange!R)
        {
            auto save() @property
            {
                rangeResult copy;
                copy.range = range.save;
                copy.key = key.dup;
                copy.nonce = nonce.dup;
                copy.count = count;
                copy.salsaSection = salsaSection.dup;
                return copy;
            }
        }
    }
    return rangeResult(range, key, nonce);
}

private:

enum string quarterRound(alias _x0, alias _x1, alias _x2, alias _x3) = q{
    %1$s += %2$s; %4$s ^= %1$s; %4$s = rotateLeft(%4$s, 16);
    %3$s += %4$s; %2$s ^= %3$s; %2$s = rotateLeft(%2$s, 12);
    %1$s += %2$s; %4$s ^= %1$s; %4$s = rotateLeft(%4$s, 8);
    %3$s += %4$s; %2$s ^= %3$s; %2$s = rotateLeft(%2$s, 7);
}.format(__traits(identifier, _x0), __traits(identifier, _x1),
         __traits(identifier, _x2), __traits(identifier, _x3));

/*unittest
{
    uint a1 = 0x00000000, a2 = 0x00000000, a3 = 0x00000000, a4 = 0x00000000,
         b1 = 0x00000001, b2 = 0x00000000, b3 = 0x00000000, b4 = 0x00000000,
         c1 = 0x00000000, c2 = 0x00000001, c3 = 0x00000000, c4 = 0x00000000,
         d1 = 0x00000000, d2 = 0x00000000, d3 = 0x00000001, d4 = 0x00000000,
         e1 = 0x00000000, e2 = 0x00000000, e3 = 0x00000000, e4 = 0x00000001,
         f1 = 0xe7e8c006, f2 = 0xc4f9417d, f3 = 0x6479b4b2, f4 = 0x68c67137,
         g1 = 0xd3917c5b, g2 = 0x55f1c407, g3 = 0x52a58a7a, g4 = 0x8f887a3b;

    mixin(quarterRound!(a1, a2, a3, a4));
    mixin(quarterRound!(b1, b2, b3, b4));
    mixin(quarterRound!(c1, c2, c3, c4));
    mixin(quarterRound!(d1, d2, d3, d4));
    mixin(quarterRound!(e1, e2, e3, e4));
    mixin(quarterRound!(f1, f2, f3, f4));
    mixin(quarterRound!(g1, g2, g3, g4));

    assert([a1, a2, a3, a4] == [0x00000000, 0x00000000, 0x00000000, 0x00000000]);
    assert([b1, b2, b3, b4] == [0x08008145, 0x00000080, 0x00010200, 0x20500000]);
    assert([c1, c2, c3, c4] == [0x88000100, 0x00000001, 0x00000200, 0x00402000]);
    assert([d1, d2, d3, d4] == [0x80040000, 0x00000000, 0x00000001, 0x00002000]);
    assert([e1, e2, e3, e4] == [0x00048044, 0x00000080, 0x00010000, 0x20100001]);
    assert([f1, f2, f3, f4] == [0xe876d72b, 0x9361dfd5, 0xf1460244, 0x948541a3]);
    assert([g1, g2, g3, g4] == [0x3e2f308c, 0xd90a8f36, 0x6ab2a923, 0x2883524c]);
}*/

enum string rowRound(alias _x00, alias _x01, alias _x02, alias _x03,
                     alias _x04, alias _x05, alias _x06, alias _x07,
                     alias _x08, alias _x09, alias _x10, alias _x11,
                     alias _x12, alias _x13, alias _x14, alias _x15,) = q{
    mixin(quarterRound!(%1$s, %6$s, %11$s, %16$s));
    mixin(quarterRound!(%2$s, %7$s, %12$s, %13$s));
    mixin(quarterRound!(%3$s, %8$s, %9$s, %14$s));
    mixin(quarterRound!(%4$s, %5$s, %10$s, %15$s));
}.format(__traits(identifier, _x00), __traits(identifier, _x01),
         __traits(identifier, _x02), __traits(identifier, _x03),
         __traits(identifier, _x04), __traits(identifier, _x05),
         __traits(identifier, _x06), __traits(identifier, _x07),
         __traits(identifier, _x08), __traits(identifier, _x09),
         __traits(identifier, _x10), __traits(identifier, _x11),
         __traits(identifier, _x12), __traits(identifier, _x13),
         __traits(identifier, _x14), __traits(identifier, _x15));

/*unittest
{
    uint y00 = 0x00000001, y01 = 0x00000000, y02 = 0x00000000, y03 = 0x00000000,
         y04 = 0x00000001, y05 = 0x00000000, y06 = 0x00000000, y07 = 0x00000000,
         y08 = 0x00000001, y09 = 0x00000000, y10 = 0x00000000, y11 = 0x00000000,
         y12 = 0x00000001, y13 = 0x00000000, y14 = 0x00000000, y15 = 0x00000000;

    uint x00 = 0x08521bd6, x01 = 0x1fe88837, x02 = 0xbb2aa576, x03 = 0x3aa26365,
         x04 = 0xc54c6a5b, x05 = 0x2fc74c2f, x06 = 0x6dd39cc3, x07 = 0xda0a64f6,
         x08 = 0x90a2f23d, x09 = 0x067f95a6, x10 = 0x06b35f61, x11 = 0x41e4732e,
         x12 = 0xe859c100, x13 = 0xea4d84b7, x14 = 0x0f619bff, x15 = 0xbc6e965a;

    mixin(rowRound!(y00, y01, y02, y03, y04, y05, y06, y07,
                     y08, y09, y10, y11, y12, y13, y14, y15));

    mixin(rowRound!(x00, x01, x02, x03, x04, x05, x06, x07,
                     x08, x09, x10, x11, x12, x13, x14, x15));

    uint[] test0 = [y00, y01, y02, y03, y04, y05, y06, y07,
                    y08, y09, y10, y11, y12, y13, y14, y15];

    uint[] test1 = [x00, x01, x02, x03, x04, x05, x06, x07,
                    x08, x09, x10, x11, x12, x13, x14, x15];

    assert(test0 == [0x08008145, 0x00000080, 0x00010200, 0x20500000,
                     0x20100001, 0x00048044, 0x00000080, 0x00010000,
                     0x00000001, 0x00002000, 0x80040000, 0x00000000,
                     0x00000001, 0x00000200, 0x00402000, 0x88000100]);

    assert(test1 == [0xa890d39d, 0x65d71596, 0xe9487daa, 0xc8ca6a86,
                     0x949d2192, 0x764b7754, 0xe408d9b9, 0x7a41b4d1,
                     0x3402e183, 0x3c3af432, 0x50669f96, 0xd89ef0a8,
                     0x0040ede5, 0xb545fbce, 0xd257ed4f, 0x1818882d]);   
}*/

enum string colRound(alias _x00, alias _x01, alias _x02, alias _x03,
                     alias _x04, alias _x05, alias _x06, alias _x07,
                     alias _x08, alias _x09, alias _x10, alias _x11,
                     alias _x12, alias _x13, alias _x14, alias _x15,) = q{
    mixin(quarterRound!(%1$s, %5$s, %9$s, %13$s));
    mixin(quarterRound!(%2$s, %6$s, %10$s, %14$s));
    mixin(quarterRound!(%3$s, %7$s, %11$s, %15$s));
    mixin(quarterRound!(%4$s, %8$s, %12$s, %16$s));
}.format(__traits(identifier, _x00), __traits(identifier, _x01),
         __traits(identifier, _x02), __traits(identifier, _x03),
         __traits(identifier, _x04), __traits(identifier, _x05),
         __traits(identifier, _x06), __traits(identifier, _x07),
         __traits(identifier, _x08), __traits(identifier, _x09),
         __traits(identifier, _x10), __traits(identifier, _x11),
         __traits(identifier, _x12), __traits(identifier, _x13),
         __traits(identifier, _x14), __traits(identifier, _x15));

/*unittest{
    uint y00 = 0x00000001, y01 = 0x00000000, y02 = 0x00000000, y03 = 0x00000000,
         y04 = 0x00000001, y05 = 0x00000000, y06 = 0x00000000, y07 = 0x00000000,
         y08 = 0x00000001, y09 = 0x00000000, y10 = 0x00000000, y11 = 0x00000000,
         y12 = 0x00000001, y13 = 0x00000000, y14 = 0x00000000, y15 = 0x00000000;

    uint x00 = 0x08521bd6, x01 = 0x1fe88837, x02 = 0xbb2aa576, x03 = 0x3aa26365,
         x04 = 0xc54c6a5b, x05 = 0x2fc74c2f, x06 = 0x6dd39cc3, x07 = 0xda0a64f6,
         x08 = 0x90a2f23d, x09 = 0x067f95a6, x10 = 0x06b35f61, x11 = 0x41e4732e,
         x12 = 0xe859c100, x13 = 0xea4d84b7, x14 = 0x0f619bff, x15 = 0xbc6e965a;


    mixin(colRound!(y00, y01, y02, y03, y04, y05, y06, y07,
                    y08, y09, y10, y11, y12, y13, y14, y15));

    mixin(colRound!(x00, x01, x02, x03, x04, x05, x06, x07,
                    x08, x09, x10, x11, x12, x13, x14, x15));

    uint[] test0 = [y00, y01, y02, y03, y04, y05, y06, y07,
                    y08, y09, y10, y11, y12, y13, y14, y15];

    uint[] test1 = [x00, x01, x02, x03, x04, x05, x06, x07,
                    x08, x09, x10, x11, x12, x13, x14, x15];

    assert(test0 == [0x10090288, 0x00000000, 0x00000000, 0x00000000,
                     0x00000101, 0x00000000, 0x00000000, 0x00000000,
                     0x00020401, 0x00000000, 0x00000000, 0x00000000,
                     0x40a04001, 0x00000000, 0x00000000, 0x00000000]);

    assert(test1 == [0x8c9d190a, 0xce8e4c90, 0x1ef8e9d3, 0x1326a71a,
                     0x90a20123, 0xead3c4f3, 0x63a091a0, 0xf0708d69,
                     0x789b010c, 0xd195a681, 0xeb7d5504, 0xa774135c,
                     0x481c2027, 0x53a8e4b5, 0x4c1f89c5, 0x3f78c9c8]);
}*/

ubyte[] chacha20(in ubyte[] input) @safe nothrow pure
in
{
    assert(input.length == 64);
}
out(result)
{
    assert(result.length == 64);
}
body
{
    auto x00 = littleEndian(input[0..4]),   x01 = littleEndian(input[4..8]), 
         x02 = littleEndian(input[8..12]),  x03 = littleEndian(input[12..16]), 
         x04 = littleEndian(input[16..20]), x05 = littleEndian(input[20..24]), 
         x06 = littleEndian(input[24..28]), x07 = littleEndian(input[28..32]), 
         x08 = littleEndian(input[32..36]), x09 = littleEndian(input[36..40]), 
         x10 = littleEndian(input[40..44]), x11 = littleEndian(input[44..48]), 
         x12 = littleEndian(input[48..52]), x13 = littleEndian(input[52..56]), 
         x14 = littleEndian(input[56..60]), x15 = littleEndian(input[60..64]);

    auto y00 = x00, y01 = x01, y02 = x02, y03 = x03,
         y04 = x04, y05 = x05, y06 = x06, y07 = x07,
         y08 = x08, y09 = x09, y10 = x10, y11 = x11,
         y12 = x12, y13 = x13, y14 = x14, y15 = x15;

    foreach(i; 0..10)
    {
        mixin(colRound!(x00, x01, x02, x03, x04, x05, x06, x07,
                        x08, x09, x10, x11, x12, x13, x14, x15));

        mixin(rowRound!(x00, x01, x02, x03, x04, x05, x06, x07,
                        x08, x09, x10, x11, x12, x13, x14, x15));
    }

    return littleEndianInv(x00 + y00) ~ littleEndianInv(x01 + y01) ~
           littleEndianInv(x02 + y02) ~ littleEndianInv(x03 + y03) ~ 
           littleEndianInv(x04 + y04) ~ littleEndianInv(x05 + y05) ~ 
           littleEndianInv(x06 + y06) ~ littleEndianInv(x07 + y07) ~ 
           littleEndianInv(x08 + y08) ~ littleEndianInv(x09 + y09) ~ 
           littleEndianInv(x10 + y10) ~ littleEndianInv(x11 + y11) ~ 
           littleEndianInv(x12 + y12) ~ littleEndianInv(x13 + y13) ~ 
           littleEndianInv(x14 + y14) ~ littleEndianInv(x15 + y15);
}

/*unittest
{
    ubyte[] test0 = [  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
                       0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
                       0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
                       0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0];
    test0 = salsa20(test0);

    ubyte[] test1 = [211,159, 13,115, 76, 55, 82,183,  3,117,222, 37,191,187,234,136,
                      49,237,179, 48,  1,106,178,219,175,199,166, 48, 86, 16,179,207,
                      31,240, 32, 63, 15, 83, 93,161,116,147, 48,113,238, 55,204, 36,
                      79,201,235, 79,  3, 81,156, 47,203, 26,244,243, 88,118,104, 54];
    test1 = salsa20(test1);

    ubyte[] test2 = [ 88,118,104, 54, 79,201,235, 79,  3, 81,156, 47,203, 26,244,243,
                     191,187,234,136,211,159, 13,115, 76, 55, 82,183,  3,117,222, 37,
                      86, 16,179,207, 49,237,179, 48,  1,106,178,219,175,199,166, 48,
                     238, 55,204, 36, 31,240, 32, 63, 15, 83, 93,161,116,147, 48,113];
    test2 = salsa20(test2);

    ubyte[] test3 = [  6,124, 83,146, 38,191,  9, 50,  4,161, 47,222,122,182,223,185,
                      75, 27,  0,216, 16,122,  7, 89,162,104,101,147,213, 21, 54, 95,
                     225,253,139,176,105,132, 23,116, 76, 41,176,207,221, 34,157,108,
                      94, 94, 99, 52, 90,117, 91,220,146,190,239,143,196,176,130,186];
    foreach(i; 0..1000000)
        test3 = salsa20(test3);

    assert(test0 == [  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
                       0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
                       0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
                       0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0]);

    assert(test1 == [109, 42,178,168,156,240,248,238,168,196,190,203, 26,110,170,154,
                      29, 29,150, 26,150, 30,235,249,190,163,251, 48, 69,144, 51, 57,
                     118, 40,152,157,180, 57, 27, 94,107, 42,236, 35, 27,111,114,114,
                     219,236,232,135,111,155,110, 18, 24,232, 95,158,179, 19, 48,202]);
    
    assert(test2 == [179, 19, 48,202,219,236,232,135,111,155,110, 18, 24,232, 95,158,
                      26,110,170,154,109, 42,178,168,156,240,248,238,168,196,190,203,
                      69,144, 51, 57, 29, 29,150, 26,150, 30,235,249,190,163,251, 48,
                      27,111,114,114,118, 40,152,157,180, 57, 27, 94,107, 42,236, 35]);

    assert(test3 == [  8, 18, 38,199,119, 76,215, 67,173,127,144,162,103,212,176,217,
                     192, 19,233, 33,159,197,154,160,128,243,219, 65,171,136,135,225,
                     123, 11, 68, 86,237, 82, 20,155,133,189,  9, 83,167,116,194, 78,
                     122,127,195,185,185,204,188, 90,245,  9,183,248,226, 85,245,104]);
}*/

enum ubyte[4] σ0 = [101, 120, 112,  97];
enum ubyte[4] σ1 = [110, 100,  32,  51];
enum ubyte[4] σ2 = [ 50,  45,  98, 121];
enum ubyte[4] σ3 = [116, 101,  32, 107];

ubyte[] chacha20Exp(in ubyte[] key, in ubyte[] n) @safe nothrow pure
in
{
    assert(key.length == 32);
    assert(n.length == 16);
}
out(result)
{
    assert(result.length == 64);
}
body
{
    return chacha20(σ0 ~ key[0..16] ~ σ1 ~ n ~ σ2 ~ key[16..$] ~ σ3);
}

/*unittest
{
    ubyte[] key;
    ubyte[] n;
    key.length = 32;
    n.length = 16;
    foreach(i; 0..16)
        key[i] = cast(ubyte)(i + 1);
    foreach(i; 16..32)
        key[i] = cast(ubyte)(i + 200 - 15);

    foreach(i; 0..16)
        n[i] = cast(ubyte)(i + 1+ 100);

    assert(salsa20Exp(key, n) == [ 69, 37, 68, 39, 41, 15,107,193,255,139,122,  6,170,233,217, 98,
                                   89,144,182,106, 21, 51,200, 65,239, 49,222, 34,215,114, 40,126,
                                  104,197,  7,225,197,153, 31,  2,102, 78, 76,176, 84,245,246,184,
                                  177,160,133,130,  6, 72,149,119,192,195,132,236,234,103,246, 74]);
}*/