const std = @import("std");

pub const key_size = 256 / 8;

const test_key = [_]u8{
    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
    0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
};

/// Get key from password. The symmetric key will be derived using PBKDF2
/// Salt can be anything, including empty if a salt is not to be used.
/// Function execution is slow by design. Memory allocated for key will be
/// owned by the caller.
pub fn keyFromPassword(allocator: std.mem.Allocator, password: []const u8, salt: []const u8) !*[key_size]u8 {
    var key = try allocator.alloc(u8, key_size);
    const rounds = 4096;
    try std.crypto.pwhash.pbkdf2(key, password, salt, rounds, std.crypto.auth.hmac.HmacSha1);
    return key[0..key_size];
}

/// Symmetric encryption with a key. Key size is defined as a public constant.
/// Key can be provided directly or through the use of the keyFromPassword
/// function.
///
/// Encryption currently AES 256, though this should be considered an implementation
/// detail. Padding performed with PKCS#7.
///
/// Example:
///
/// const str = "0123456789ABCDEFG";
/// var buf: [str.len]u8 = undefined;
/// var data = std.mem.copy(u8, buf, str);
/// var key = try keyFromPassword(allocator, "foo", ""); // reuse key - this is slow
/// defer allocator.free(key);
/// const encrypted = try encryptWithKey(allocator, key.*, data);
///
pub fn encryptWithKey(allocator: std.mem.Allocator, key: [key_size]u8, data: []u8) ![]const u8 {
    const block_size = 16;
    var in: *[block_size]u8 = undefined;
    var out: [block_size]u8 = undefined;
    var ctx = std.crypto.core.aes.Aes256.initEnc(key);

    const total_blocks = data.len / block_size;
    var current_block: usize = 0;
    var encrypted = try std.ArrayList(u8).initCapacity(allocator, block_size * (total_blocks + 1));
    defer encrypted.deinit();
    while (current_block < total_blocks) {
        in = @ptrCast(*[block_size]u8, data[(current_block * block_size)..(((current_block + 1) * block_size) - 1)]);
        ctx.encrypt(out[0..], in.*[0..]);
        encrypted.appendSliceAssumeCapacity(out[0..]);
        current_block += 1;
    }

    // deal with final block, PKCS#7 padding
    {
        const padding: u8 = @intCast(u8, block_size - (data.len % block_size));
        var inx: u8 = 0;
        for (data[(total_blocks * block_size)..]) |b| {
            in[inx] = b;
            inx += 1;
        }
        while (inx < out.len) {
            in[inx] = padding;
            inx += 1;
        }
        ctx.encrypt(out[0..], in[0..]);
        encrypted.appendSliceAssumeCapacity(out[0..]);
    }
    return encrypted.toOwnedSlice();
}

/// Symmetric decryption with a key. Key size is defined as a public constant.
/// Key can be provided directly or through the use of the keyFromPassword
/// function.
///
/// Decryption currently AES 256, though this should be considered an implementation
/// detail. Padding assumed to be PKCS#7.
///
/// Example:
///
/// const str = "0123456789ABCDEFG";
/// var buf: [str.len]u8 = undefined;
/// var data = std.mem.copy(u8, buf, str);
/// var key = try keyFromPassword(allocator, "foo", ""); // reuse key - this is slow
/// defer allocator.free(key);
/// const encrypted = try encryptWithKey(allocator, key.*, data);
///
pub fn decryptWithKey(allocator: std.mem.Allocator, key: [key_size]u8, data: []const u8) ![]const u8 {
    const block_size = 16;
    var in: *const [block_size]u8 = undefined;
    var out: [block_size]u8 = undefined;
    var ctx = std.crypto.core.aes.Aes256.initDec(key);

    const total_blocks = data.len / block_size;
    var current_block: usize = 0;
    var decrypted = try std.ArrayList(u8).initCapacity(allocator, block_size * (total_blocks - 1));
    defer decrypted.deinit();
    while (current_block < total_blocks - 1) {
        in = @ptrCast(*const [block_size]u8, data[(current_block * block_size)..(((current_block + 1) * block_size) - 1)]);
        ctx.decrypt(out[0..], in.*[0..]);
        decrypted.appendSliceAssumeCapacity(out[0..]);
        current_block += 1;
    }

    // deal with final block, PKCS#7 padding
    in = @ptrCast(*const [block_size]u8, data[(current_block * block_size)..(((current_block + 1) * block_size) - 1)]);
    ctx.decrypt(out[0..], in.*[0..]);
    std.debug.assert(out[block_size - 1] <= block_size);
    try decrypted.appendSlice(out[0 .. block_size - out[block_size - 1]]);
    return decrypted.toOwnedSlice();
}

test "Appendix C.3" {
    var in = [_]u8{ 0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff };
    // const exp_out = [_]u8{ 0x8e, 0xa2, 0xb7, 0xca, 0x51, 0x67, 0x45, 0xbf, 0xea, 0xfc, 0x49, 0x90, 0x4b, 0x49, 0x60, 0x89 };
    var allocator = std.testing.allocator;
    const encrypted = try encryptWithKey(allocator, test_key, in[0..]);
    defer allocator.free(encrypted);
    const size: usize = 32;
    try std.testing.expectEqual(size, encrypted.len);
}

test "correct size, even block" {
    const str = "0123456789ABCDEF";
    var allocator = std.testing.allocator;
    var buf: [str.len]u8 = undefined;
    var data = try std.fmt.bufPrint(&buf, str, .{});
    const encrypted = try encryptWithKey(allocator, test_key, data);
    defer allocator.free(encrypted);
    const size: usize = 32;
    try std.testing.expectEqual(size, encrypted.len);
}

test "round trip, even block" {
    const str = "0123456789ABCDEF";
    var allocator = std.testing.allocator;
    var buf: [str.len]u8 = undefined;
    var data = try std.fmt.bufPrint(&buf, str, .{});
    const encrypted = try encryptWithKey(allocator, test_key, data);
    defer allocator.free(encrypted);
    const decrypted = try decryptWithKey(allocator, test_key, encrypted);
    defer allocator.free(decrypted);
    try std.testing.expectEqual(str.len, decrypted.len);
    try std.testing.expectEqualStrings(str, decrypted);
}

test "round trip, uneven block" {
    const str = "0123456789ABCDEFG";
    var allocator = std.testing.allocator;
    var buf: [str.len]u8 = undefined;
    var data = try std.fmt.bufPrint(&buf, str, .{});
    const encrypted = try encryptWithKey(allocator, test_key, data);
    defer allocator.free(encrypted);
    const decrypted = try decryptWithKey(allocator, test_key, encrypted);
    defer allocator.free(decrypted);
    try std.testing.expectEqualStrings(str, decrypted);
}

test "round trip, uneven block, using password" {
    const str = "0123456789ABCDEFG";
    var allocator = std.testing.allocator;
    var buf: [str.len]u8 = undefined;
    var data = try std.fmt.bufPrint(&buf, str, .{});
    var key = try keyFromPassword(allocator, "foo", "");
    defer allocator.free(key);
    const encrypted = try encryptWithKey(allocator, key.*, data);
    defer allocator.free(encrypted);
    const decrypted = try decryptWithKey(allocator, key.*, encrypted);
    defer allocator.free(decrypted);
    try std.testing.expectEqualStrings(str, decrypted);
}