module scrypt

import crypto.pbkdf2
import crypto.sha256

fn test_salsa20_8() {
    // The input_block and output_block values are taken from
    // [RFC7914](https://datatracker.ietf.org/doc/html/rfc7914#section-8)
    // section 8.

    // vfmt off
    mut input_block := [
        u8(0x7e), 0x87, 0x9a, 0x21,   0x4f, 0x3e, 0xc9, 0x86,
           0x7c,  0xa9, 0x40, 0xe6,   0x41, 0x71, 0x8f, 0x26,
           0xba,  0xee, 0x55, 0x5b,   0x8c, 0x61, 0xc1, 0xb5,
           0x0d,  0xf8, 0x46, 0x11,   0x6d, 0xcd, 0x3b, 0x1d,
           0xee,  0x24, 0xf3, 0x19,   0xdf, 0x9b, 0x3d, 0x85,
           0x14,  0x12, 0x1e, 0x4b,   0x5a, 0xc5, 0xaa, 0x32,
           0x76,  0x02, 0x1d, 0x29,   0x09, 0xc7, 0x48, 0x29,
           0xed,  0xeb, 0xc6, 0x8d,   0xb8, 0xb8, 0xc2, 0x5e]

    output_block := [
        u8(0xa4), 0x1f, 0x85, 0x9c,   0x66, 0x08, 0xcc, 0x99,
           0x3b,  0x81, 0xca, 0xcb,   0x02, 0x0c, 0xef, 0x05,
           0x04,  0x4b, 0x21, 0x81,   0xa2, 0xfd, 0x33, 0x7d,
           0xfd,  0x7b, 0x1c, 0x63,   0x96, 0x68, 0x2f, 0x29,
           0xb4,  0x39, 0x31, 0x68,   0xe3, 0xc9, 0xe6, 0xbc,
           0xfe,  0x6b, 0xc5, 0xb7,   0xa0, 0x6d, 0x96, 0xba,
           0xe4,  0x24, 0xcc, 0x10,   0x2c, 0x91, 0x74, 0x5c,
           0x24,  0xad, 0x67, 0x3d,   0xc7, 0x61, 0x8f, 0x81]
    // vfmt on

    salsa20_8(mut input_block)

    for i in 0 .. 64 {
        assert input_block[i] == output_block[i], 'assertion failed for i: ${i}'
    }
}

fn test_block_mix() {
    // The input_block and output_block values are taken from
    // [RFC7914](https://datatracker.ietf.org/doc/html/rfc7914#section-9)
    // section 9.

    // vfmt off
    mut input_block := [
        // B[0] - the first 64 bytes of input
        u8(0xf7), 0xce, 0x0b, 0x65,   0x3d, 0x2d, 0x72, 0xa4,
           0x10,  0x8c, 0xf5, 0xab,   0xe9, 0x12, 0xff, 0xdd,
           0x77,  0x76, 0x16, 0xdb,   0xbb, 0x27, 0xa7, 0x0e,
           0x82,  0x04, 0xf3, 0xae,   0x2d, 0x0f, 0x6f, 0xad,
           0x89,  0xf6, 0x8f, 0x48,   0x11, 0xd1, 0xe8, 0x7b,
           0xcc,  0x3b, 0xd7, 0x40,   0x0a, 0x9f, 0xfd, 0x29,
           0x09,  0x4f, 0x01, 0x84,   0x63, 0x95, 0x74, 0xf3,
           0x9a,  0xe5, 0xa1, 0x31,   0x52, 0x17, 0xbc, 0xd7,
        // B[1] - the second 64 bytes of input
           0x89,  0x49, 0x91, 0x44,   0x72, 0x13, 0xbb, 0x22,
           0x6c,  0x25, 0xb5, 0x4d,   0xa8, 0x63, 0x70, 0xfb,
           0xcd,  0x98, 0x43, 0x80,   0x37, 0x46, 0x66, 0xbb,
           0x8f,  0xfc, 0xb5, 0xbf,   0x40, 0xc2, 0x54, 0xb0,
           0x67,  0xd2, 0x7c, 0x51,   0xce, 0x4a, 0xd5, 0xfe,
           0xd8,  0x29, 0xc9, 0x0b,   0x50, 0x5a, 0x57, 0x1b,
           0x7f,  0x4d, 0x1c, 0xad,   0x6a, 0x52, 0x3c, 0xda,
           0x77,  0x0e, 0x67, 0xbc,   0xea, 0xaf, 0x7e, 0x89]

    output_block := [
        // B'[0] - the first 64 bytes of output
        u8(0xa4), 0x1f, 0x85, 0x9c,   0x66, 0x08, 0xcc, 0x99,
           0x3b,  0x81, 0xca, 0xcb,   0x02, 0x0c, 0xef, 0x05,
           0x04,  0x4b, 0x21, 0x81,   0xa2, 0xfd, 0x33, 0x7d,
           0xfd,  0x7b, 0x1c, 0x63,   0x96, 0x68, 0x2f, 0x29,
           0xb4,  0x39, 0x31, 0x68,   0xe3, 0xc9, 0xe6, 0xbc,
           0xfe,  0x6b, 0xc5, 0xb7,   0xa0, 0x6d, 0x96, 0xba,
           0xe4,  0x24, 0xcc, 0x10,   0x2c, 0x91, 0x74, 0x5c,
           0x24,  0xad, 0x67, 0x3d,   0xc7, 0x61, 0x8f, 0x81,
        // B'[1] - the second 64 bytes of output
           0x20,  0xed, 0xc9, 0x75,   0x32, 0x38, 0x81, 0xa8,
           0x05,  0x40, 0xf6, 0x4c,   0x16, 0x2d, 0xcd, 0x3c,
           0x21,  0x07, 0x7c, 0xfe,   0x5f, 0x8d, 0x5f, 0xe2,
           0xb1,  0xa4, 0x16, 0x8f,   0x95, 0x36, 0x78, 0xb7,
           0x7d,  0x3b, 0x3d, 0x80,   0x3b, 0x60, 0xe4, 0xab,
           0x92,  0x09, 0x96, 0xe5,   0x9b, 0x4d, 0x53, 0xb6,
           0x5d,  0x2a, 0x22, 0x58,   0x77, 0xd5, 0xed, 0xf5,
           0x84,  0x2c, 0xb9, 0xf1,   0x4e, 0xef, 0xe4, 0x25]
    // vfmt on

    // an array capable of holding r * 128 bytes used during
    // the block_mix operation.
    mut temp_block := []u8{len: 128, cap: 128}

    // for this test, r = 1
    block_mix(mut input_block, mut temp_block, 1)

    for i in 0 .. 128 {
        assert input_block[i] == output_block[i], 'assertion failed for i: ${i}'
    }
}

fn test_smix() {
    // The input_block and output_block values are taken from
    // [RFC7914](https://datatracker.ietf.org/doc/html/rfc7914#section-10)
    // section 10.

    // vfmt off
    mut input_block := [
        u8(0xf7), 0xce, 0x0b, 0x65,   0x3d, 0x2d, 0x72, 0xa4,
           0x10,  0x8c, 0xf5, 0xab,   0xe9, 0x12, 0xff, 0xdd,
           0x77,  0x76, 0x16, 0xdb,   0xbb, 0x27, 0xa7, 0x0e,
           0x82,  0x04, 0xf3, 0xae,   0x2d, 0x0f, 0x6f, 0xad,
           0x89,  0xf6, 0x8f, 0x48,   0x11, 0xd1, 0xe8, 0x7b,
           0xcc,  0x3b, 0xd7, 0x40,   0x0a, 0x9f, 0xfd, 0x29,
           0x09,  0x4f, 0x01, 0x84,   0x63, 0x95, 0x74, 0xf3,
           0x9a,  0xe5, 0xa1, 0x31,   0x52, 0x17, 0xbc, 0xd7,
           0x89,  0x49, 0x91, 0x44,   0x72, 0x13, 0xbb, 0x22,
           0x6c,  0x25, 0xb5, 0x4d,   0xa8, 0x63, 0x70, 0xfb,
           0xcd,  0x98, 0x43, 0x80,   0x37, 0x46, 0x66, 0xbb,
           0x8f,  0xfc, 0xb5, 0xbf,   0x40, 0xc2, 0x54, 0xb0,
           0x67,  0xd2, 0x7c, 0x51,   0xce, 0x4a, 0xd5, 0xfe,
           0xd8,  0x29, 0xc9, 0x0b,   0x50, 0x5a, 0x57, 0x1b,
           0x7f,  0x4d, 0x1c, 0xad,   0x6a, 0x52, 0x3c, 0xda,
           0x77,  0x0e, 0x67, 0xbc,   0xea, 0xaf, 0x7e, 0x89]

    output_block := [
        u8(0x79), 0xcc, 0xc1, 0x93,   0x62, 0x9d, 0xeb, 0xca,
           0x04,  0x7f, 0x0b, 0x70,   0x60, 0x4b, 0xf6, 0xb6,
           0x2c,  0xe3, 0xdd, 0x4a,   0x96, 0x26, 0xe3, 0x55,
           0xfa,  0xfc, 0x61, 0x98,   0xe6, 0xea, 0x2b, 0x46,
           0xd5,  0x84, 0x13, 0x67,   0x3b, 0x99, 0xb0, 0x29,
           0xd6,  0x65, 0xc3, 0x57,   0x60, 0x1f, 0xb4, 0x26,
           0xa0,  0xb2, 0xf4, 0xbb,   0xa2, 0x00, 0xee, 0x9f,
           0x0a,  0x43, 0xd1, 0x9b,   0x57, 0x1a, 0x9c, 0x71,
           0xef,  0x11, 0x42, 0xe6,   0x5d, 0x5a, 0x26, 0x6f,
           0xdd,  0xca, 0x83, 0x2c,   0xe5, 0x9f, 0xaa, 0x7c,
           0xac,  0x0b, 0x9c, 0xf1,   0xbe, 0x2b, 0xff, 0xca,
           0x30,  0x0d, 0x01, 0xee,   0x38, 0x76, 0x19, 0xc4,
           0xae,  0x12, 0xfd, 0x44,   0x38, 0xf2, 0x03, 0xa0,
           0xe4,  0xe1, 0xc4, 0x7e,   0xc3, 0x14, 0x86, 0x1f,
           0x4e,  0x90, 0x87, 0xcb,   0x33, 0x39, 0x6a, 0x68,
           0x73,  0xe8, 0xf9, 0xd2,   0x53, 0x9a, 0x4b, 0x8e]
    // vfmt on

    r := u32(1)
    n := u64(16)

    // len and cap are 128 * r * n = 2048
    mut v_block := []u8{len: 2048, cap: 2048}

    // len and cap are 256 * r = 246
    mut temp_block := []u8{len: 256, cap: 256}

    smix(mut input_block, r, n, mut v_block, mut temp_block)

    for i in 0 .. 128 {
        assert input_block[i] == output_block[i], 'assertion failed for i: ${i}'
    }
}

fn test_pbkdf2_hmac_sha256() {
    // The input_block and output_block values are taken from
    // [RFC7914](https://datatracker.ietf.org/doc/html/rfc7914#section-11)
    // section 11.

    // vfmt off
    output_block := [
        [u8(0x55), 0xac, 0x04, 0x6e,   0x56, 0xe3, 0x08, 0x9f,
            0xec,  0x16, 0x91, 0xc2,   0x25, 0x44, 0xb6, 0x05,
            0xf9,  0x41, 0x85, 0x21,   0x6d, 0xde, 0x04, 0x65,
            0xe6,  0x8b, 0x9d, 0x57,   0xc2, 0x0d, 0xac, 0xbc,
            0x49,  0xca, 0x9c, 0xcc,   0xf1, 0x79, 0xb6, 0x45,
            0x99,  0x16, 0x64, 0xb3,   0x9d, 0x77, 0xef, 0x31,
            0x7c,  0x71, 0xb8, 0x45,   0xb1, 0xe3, 0x0b, 0xd5,
            0x09,  0x11, 0x20, 0x41,   0xd3, 0xa1, 0x97, 0x83
        ],
        [u8(0x4d), 0xdc, 0xd8, 0xf6,   0x0b, 0x98, 0xbe, 0x21,
            0x83,  0x0c, 0xee, 0x5e,   0xf2, 0x27, 0x01, 0xf9,
            0x64,  0x1a, 0x44, 0x18,   0xd0, 0x4c, 0x04, 0x14,
            0xae,  0xff, 0x08, 0x87,   0x6b, 0x34, 0xab, 0x56,
            0xa1,  0xd4, 0x25, 0xa1,   0x22, 0x58, 0x33, 0x54,
            0x9a,  0xdb, 0x84, 0x1b,   0x51, 0xc9, 0xb3, 0x17,
            0x6a,  0x27, 0x2b, 0xde,   0xbb, 0xa1, 0xd0, 0x78,
            0x47,  0x8f, 0x62, 0xb3,   0x97, 0xf3, 0x3c, 0x8d
        ]
    ]
    // vfmt on

    d0 := pbkdf2.key('passwd'.bytes(), 'salt'.bytes(), 1, 64, sha256.new())!

    assert d0 == output_block[0]

    d1 := pbkdf2.key('Password'.bytes(), 'NaCl'.bytes(), 80000, 64, sha256.new())!

    assert d1 == output_block[1]
}

struct ScryptTestData {
    name            string
    password        []u8
    salt            []u8
    n               u64
    r               u32
    p               u32
    dk_len          u64
    expected_result []u8
}

// The scrypt test vectors are taken from
// [RFC7914](https://datatracker.ietf.org/doc/html/rfc7914#section-12)
// section 12.
const scrypt_test_cases = [
    ScryptTestData{
        name:            'test case 1'
        password:        ''.bytes()
        salt:            ''.bytes()
        n:               16
        r:               1
        p:               1
        dk_len:          64
        expected_result: [u8(0x77), 0xd6, 0x57, 0x62, 0x38, 0x65, 0x7b, 0x20, 0x3b, 0x19, 0xca,
            0x42, 0xc1, 0x8a, 0x04, 0x97, 0xf1, 0x6b, 0x48, 0x44, 0xe3, 0x07, 0x4a, 0xe8, 0xdf,
            0xdf, 0xfa, 0x3f, 0xed, 0xe2, 0x14, 0x42, 0xfc, 0xd0, 0x06, 0x9d, 0xed, 0x09, 0x48,
            0xf8, 0x32, 0x6a, 0x75, 0x3a, 0x0f, 0xc8, 0x1f, 0x17, 0xe8, 0xd3, 0xe0, 0xfb, 0x2e,
            0x0d, 0x36, 0x28, 0xcf, 0x35, 0xe2, 0x0c, 0x38, 0xd1, 0x89, 0x06]
    },
    // test cases 2, 3, and 4 are moved to the slow test repo.
]

fn test_scrypt() {
    for c in scrypt_test_cases {
        results := scrypt(c.password, c.salt, c.n, c.r, c.p, c.dk_len)!
        assert results == c.expected_result, '${c.name} failed'
    }
}