const std = @import("std");

const builtin = std.builtin;
const mem = std.mem;

const argon2 = std.crypto.pwhash.argon2;
const ed25519 = std.crypto.sign.Ed25519;
const xx32 = std.hash.XxHash32;

pub const pawprint = struct {
    pub const Context = struct {
        alloc: mem.Allocator,
        io: std.Io,
    };

    pub const identity = struct {
        pub const UnlockError = error{
            ChecksumMismatch,
        };

        const salt_len = 16;

        pub fn passphrase_kdf(
            ctx: *const Context,
            salt: [salt_len]u8,
            passphrase: []const u8,
        ) ![36]u8 {
            var derived: [36]u8 = undefined;
            try argon2.kdf(
                ctx.alloc,
                &derived,
                passphrase,
                &salt,
                argon2.Params.sensitive_2id,
                argon2.Mode.argon2id,
                ctx.io,
            );

            return derived;
        }

        pub const Locked = struct {
            bytes: [36]u8,
            salt: [salt_len]u8,

            pub fn unlock(
                self: *const Locked,
                ctx: *const Context,
                passphrase: []const u8,
            ) !Unlocked {
                var seed = try passphrase_kdf(ctx, self.salt, passphrase);
                for (0..36) |i| {
                    seed[i] ^= self.bytes[i];
                }

                const csum = xx32.hash(0, seed[0..32]);

                var csum_bytes: [4]u8 = undefined;
                mem.writeInt(u32, &csum_bytes, csum, builtin.Endian.little);

                if (!mem.eql(u8, &csum_bytes, seed[32..36])) {
                    return UnlockError.ChecksumMismatch;
                }

                return Unlocked{
                    .keypair = try ed25519.KeyPair.generateDeterministic(
                        seed[0..32].*,
                    ),
                };
            }

            pub fn drop(self: *volatile Locked) void {
                self.* = std.mem.zeroes(Locked);
            }
        };

        pub const Unlocked = struct {
            keypair: ed25519.KeyPair,

            pub fn generate(io: std.Io) Unlocked {
                return Unlocked{
                    .keypair = ed25519.KeyPair.generate(io),
                };
            }

            pub fn lock(
                self: *const Unlocked,
                ctx: *const Context,
                passphrase: []const u8,
            ) !Locked {
                var salt: [salt_len]u8 = undefined;
                try ctx.io.randomSecure(&salt);

                const derived = try passphrase_kdf(ctx, salt, passphrase);
                var locked = Locked{
                    .bytes = derived,
                    .salt = salt,
                };

                const seed = self.keypair.secret_key.seed();
                for (0..32) |i| {
                    locked.bytes[i] ^= seed[i];
                }

                const csum = xx32.hash(0, seed);

                var csum_bytes: [4]u8 = undefined;
                mem.writeInt(u32, &csum_bytes, csum, builtin.Endian.little);

                for (0..4) |i| {
                    locked.bytes[i + 32] ^= csum_bytes[i];
                }

                return locked;
            }

            pub fn drop(self: *volatile Unlocked) void {
                self.* = std.mem.zeroes(Unlocked);
            }
        };
    };
};

test "lock and unlock identity" {
    const ctx = pawprint.Context{
        .alloc = std.testing.allocator,
        .io = std.testing.io,
    };

    const ident = pawprint.identity.Unlocked.generate(ctx.io);
    const locked = try ident.lock(&ctx, "foobar123");
    const unlocked = try locked.unlock(&ctx, "foobar123");

    try std.testing.expectEqualDeep(ident, unlocked);
}

test "unlock with wrong passphrase" {
    const ctx = pawprint.Context{
        .alloc = std.testing.allocator,
        .io = std.testing.io,
    };

    const ident = pawprint.identity.Unlocked.generate(ctx.io);
    const locked = try ident.lock(&ctx, "foobar123");
    const unlocked = locked.unlock(&ctx, "barbaz456");

    try std.testing.expectEqual(
        unlocked,
        pawprint.identity.UnlockError.ChecksumMismatch,
    );
}

test "drop unlocked identity" {
    const ctx = pawprint.Context{
        .alloc = std.testing.allocator,
        .io = std.testing.io,
    };

    var ident = pawprint.identity.Unlocked.generate(ctx.io);
    ident.drop();

    try std.testing.expectEqual(
        mem.zeroes(pawprint.identity.Unlocked),
        ident,
    );
}

test "drop locked identity" {
    const ctx = pawprint.Context{
        .alloc = std.testing.allocator,
        .io = std.testing.io,
    };

    const ident = pawprint.identity.Unlocked.generate(ctx.io);
    var locked = try ident.lock(&ctx, "foobar123");
    locked.drop();

    try std.testing.expectEqual(
        mem.zeroes(pawprint.identity.Locked),
        locked,
    );
}