1 files changed, 462 insertions, 0 deletions

diff --git a/src/znk.zig b/src/znk.zig
new file mode 100644
index 0000000..ab36c96
--- /dev/null
+++ b/src/znk.zig
@@ -0,0 +1,462 @@
+const std = @import("std");
+const Allocator = mem.Allocator;
+const ascii = std.ascii;
+const build_options = @import("build_options");
+const builtin = std.builtin;
+const fs = std.fs;
+const io = std.io;
+const mem = std.mem;
+const nkeys = @import("nkeys.zig");
+const process = std.process;
+const testing = std.testing;
+
+pub fn fatal(comptime format: []const u8, args: anytype) noreturn {
+    std.debug.print("error: " ++ format ++ "\n", args);
+    process.exit(1);
+}
+
+pub fn info(comptime format: []const u8, args: anytype) void {
+    std.debug.print(format ++ "\n", args);
+}
+
+const usage =
+    \\Usage: znk [command] [options]
+    \\
+    \\Commands:
+    \\
+    \\  gen            Generate a new key pair
+    \\  help           Print this help and exit
+    \\  sign           Sign a file
+    \\  verify         Verify a file with a signature
+    \\  version        Print version number and exit
+    \\
+    \\General Options:
+    \\
+    \\  -h, --help     Print this help and exit
+    \\
+;
+
+var general_purpose_allocator = std.heap.GeneralPurposeAllocator(.{}){};
+
+pub fn main() anyerror!void {
+    // Stolen from the Zig compiler
+    var gpa_need_deinit = false;
+    const gpa = gpa: {
+        if (!std.builtin.link_libc) {
+            gpa_need_deinit = true;
+            break :gpa &general_purpose_allocator.allocator;
+        }
+        // We would prefer to use raw libc allocator here, but cannot
+        // use it if it won't support the alignment we need.
+        if (@alignOf(std.c.max_align_t) < @alignOf(i128)) {
+            break :gpa std.heap.c_allocator;
+        }
+        break :gpa std.heap.raw_c_allocator;
+    };
+    defer if (gpa_need_deinit) {
+        std.debug.assert(!general_purpose_allocator.deinit());
+    };
+    var arena_instance = std.heap.ArenaAllocator.init(gpa);
+    defer arena_instance.deinit();
+    const arena = &arena_instance.allocator;
+
+    const args = try process.argsAlloc(arena);
+    return mainArgs(gpa, arena, args);
+}
+
+pub fn mainArgs(gpa: *Allocator, arena: *Allocator, args: []const []const u8) !void {
+    if (args.len <= 1) {
+        info("{s}", .{usage});
+        fatal("expected command argument", .{});
+    }
+
+    const cmd = args[1];
+    const cmd_args = args[2..];
+    if (mem.eql(u8, cmd, "gen")) {
+        return cmdGen(gpa, arena, cmd_args);
+    } else if (mem.eql(u8, cmd, "sign")) {
+        return cmdSign(gpa, arena, cmd_args);
+    } else if (mem.eql(u8, cmd, "verify")) {
+        return cmdVerify(gpa, arena, cmd_args);
+    } else if (mem.eql(u8, cmd, "version")) {
+        return io.getStdOut().writeAll(build_options.version ++ "\n");
+    } else if (mem.eql(u8, cmd, "help") or mem.eql(u8, cmd, "-h") or mem.eql(u8, cmd, "--help")) {
+        return io.getStdOut().writeAll(usage);
+    } else {
+        info("{s}", .{usage});
+        fatal("unknown command: {s}", .{cmd});
+    }
+}
+
+const usage_gen =
+    \\Usage: znk gen [options] <type>
+    \\
+    \\Supported Types:
+    \\
+    \\  account
+    \\  cluster
+    \\  operator
+    \\  server
+    \\  user
+    \\
+    \\General Options:
+    \\
+    \\  -h, --help     Print this help and exit
+    \\
+    \\Generate Options:
+    \\
+    \\  -o, --pub-out  Print the public key to stdout
+    \\  -p, --prefix   Vanity public key prefix, turns -o on
+    \\
+;
+
+pub fn cmdGen(gpa: *Allocator, arena: *Allocator, args: []const []const u8) !void {
+    const stdout = io.getStdOut();
+
+    var ty: ?nkeys.PublicPrefixByte = null;
+    var pub_out: bool = false;
+    var prefix: ?[]const u8 = null;
+
+    var i: usize = 0;
+    while (i < args.len) : (i += 1) {
+        const arg = args[i];
+        if (mem.startsWith(u8, arg, "-")) {
+            if (mem.eql(u8, arg, "-h") or mem.eql(u8, arg, "--help")) {
+                return stdout.writeAll(usage_gen);
+            } else if (mem.eql(u8, arg, "-o") or mem.eql(u8, arg, "--pub-out")) {
+                pub_out = true;
+            } else if (mem.eql(u8, arg, "-p") or mem.eql(u8, arg, "--prefix")) {
+                if (i + 1 >= args.len)
+                    fatal("expected argument after '{s}'", .{arg});
+                i += 1;
+                if (args[i].len > nkeys.text_public_len - 1)
+                    fatal("public key prefix '{s}' is too long", .{arg});
+                prefix = args[i];
+            } else {
+                fatal("unrecognized parameter: '{s}'", .{arg});
+            }
+        } else if (ty != null) {
+            fatal("more than one type to generate provided", .{});
+        } else if (mem.eql(u8, arg, "account")) {
+            ty = .account;
+        } else if (mem.eql(u8, arg, "cluster")) {
+            ty = .cluster;
+        } else if (mem.eql(u8, arg, "operator")) {
+            ty = .operator;
+        } else if (mem.eql(u8, arg, "server")) {
+            ty = .server;
+        } else if (mem.eql(u8, arg, "user")) {
+            ty = .user;
+        } else {
+            fatal("unrecognized extra parameter: '{s}'", .{arg});
+        }
+    }
+
+    if (ty == null) {
+        info("{s}", .{usage_gen});
+        fatal("no type to generate seed for provided", .{});
+    }
+
+    if (prefix != null) {
+        const capitalized_prefix = try toUpper(arena, prefix.?);
+
+        try PrefixKeyGenerator.init(arena, ty.?, capitalized_prefix).generate();
+    } else {
+        var kp = nkeys.SeedKeyPair.init(ty.?) catch |e| fatal("could not generate key pair: {e}", .{e});
+        defer kp.wipe();
+        try stdout.writeAll(&kp.seed);
+        try stdout.writeAll("\n");
+
+        var public_key = kp.publicKey() catch |e| fatal("could not generate public key: {e}", .{e});
+        if (pub_out) {
+            try stdout.writeAll(&public_key);
+            try stdout.writeAll("\n");
+        }
+    }
+}
+
+const usage_sign =
+    \\Usage: znk sign -k <file> [options] <file>
+    \\
+    \\General Options:
+    \\
+    \\  -h, --help     Print this help and exit
+    \\
+    \\Sign Options:
+    \\
+    \\  -k, --key      Path of private key/seed to sign with
+    \\
+;
+
+pub fn cmdSign(gpa: *Allocator, arena: *Allocator, args: []const []const u8) !void {
+    // TODO(rutgerbrf): support setting a custom entropy file?
+    const stdin = io.getStdIn();
+    const stdout = io.getStdOut();
+
+    var file_stdin = false;
+    var key_stdin = false;
+    var file: ?fs.File = null;
+    var key: ?fs.File = null;
+    defer if (!key_stdin) if (file) |f| f.close();
+    defer if (!file_stdin) if (key) |f| f.close();
+
+    var i: usize = 0;
+    while (i < args.len) : (i += 1) {
+        const arg = args[i];
+        if (mem.startsWith(u8, arg, "-") and arg.len > 1) {
+            if (mem.eql(u8, arg, "-h") or mem.eql(u8, arg, "--help")) {
+                return stdout.writeAll(usage_sign);
+            } else if (mem.eql(u8, arg, "-k") or mem.eql(u8, arg, "--key")) {
+                if (i + 1 >= args.len) fatal("expected argument after '{s}'", .{arg});
+                i += 1;
+                if (key != null) fatal("parameter '{s}' provided more than once", .{arg});
+                if (std.mem.eql(u8, args[i], "-")) {
+                    key = stdin;
+                    key_stdin = true;
+                } else {
+                    key = try fs.cwd().openFile(args[i], .{});
+                }
+            } else {
+                fatal("unrecognized parameter: '{s}'", .{arg});
+            }
+        } else if (file != null) {
+            fatal("more than one file to generate a signature for provided", .{});
+        } else if (mem.eql(u8, args[i], "-")) {
+            file = stdin;
+            file_stdin = true;
+        } else {
+            file = try fs.cwd().openFile(args[i], .{});
+        }
+    }
+
+    if (file == null) {
+        info("{s}", .{usage_sign});
+        fatal("no file to generate a signature for provided", .{});
+    }
+
+    if (key == null) {
+        info("{s}", .{usage_sign});
+        fatal("no key to sign with provided", .{});
+    }
+
+    if (file_stdin and key_stdin) {
+        fatal("can't use stdin for reading multiple files", .{});
+    }
+
+    const content = file.?.readToEndAlloc(arena, std.math.maxInt(usize)) catch {
+        fatal("could not read file to generate signature for", .{});
+    };
+    var kp = switch (readKeyFile(arena, key.?)) {
+        .seed_key_pair => |kp| kp,
+        else => |*k| {
+            k.wipe();
+            fatal("key was provided but is not a seed", .{});
+        },
+    };
+    defer kp.wipe();
+
+    const sig = kp.sign(content) catch fatal("could not generate signature", .{});
+    var encoded_sig = try arena.alloc(u8, std.base64.standard.Encoder.calcSize(sig.len));
+    _ = std.base64.standard.Encoder.encode(encoded_sig, &sig);
+    try stdout.writeAll(encoded_sig);
+    try stdout.writeAll("\n");
+}
+
+const usage_verify =
+    \\Usage: znk verify [options] <file>
+    \\
+    \\General Options:
+    \\
+    \\  -h, --help     Print this help and exit
+    \\
+    \\Verify Options:
+    \\
+    \\  -k, --key      Path of key to verify with
+    \\  -s, --sig      Path of signature to verify
+    \\
+;
+
+pub fn cmdVerify(gpa: *Allocator, arena: *Allocator, args: []const []const u8) !void {
+    const stdin = io.getStdIn();
+    const stdout = io.getStdOut();
+
+    var file_stdin = false;
+    var key_stdin = false;
+    var sig_stdin = false;
+    var key: ?fs.File = null;
+    var file: ?fs.File = null;
+    var sig: ?fs.File = null;
+    defer if (!file_stdin) if (file) |f| f.close();
+    defer if (!key_stdin) if (key) |f| f.close();
+    defer if (!sig_stdin) if (sig) |f| f.close();
+
+    var i: usize = 0;
+    while (i < args.len) : (i += 1) {
+        const arg = args[i];
+        if (mem.startsWith(u8, arg, "-") and arg.len > 1) {
+            if (mem.eql(u8, arg, "-h") or mem.eql(u8, arg, "--help")) {
+                return stdout.writeAll(usage_verify);
+            } else if (mem.eql(u8, arg, "-k") or mem.eql(u8, arg, "--key")) {
+                if (i + 1 >= args.len) fatal("expected argument after '{s}'", .{arg});
+                i += 1;
+                if (key != null) fatal("parameter '{s}' provided more than once", .{arg});
+                if (std.mem.eql(u8, args[i], "-")) {
+                    key = stdin;
+                    key_stdin = true;
+                } else {
+                    key = try fs.cwd().openFile(args[i], .{});
+                }
+            } else if (mem.eql(u8, arg, "-s") or mem.eql(u8, arg, "--sig")) {
+                if (i + 1 >= args.len) fatal("expected argument after '{s}'", .{arg});
+                i += 1;
+                if (sig != null) fatal("parameter '{s}' provided more than once", .{arg});
+                if (std.mem.eql(u8, args[i], "-")) {
+                    sig = stdin;
+                    sig_stdin = true;
+                } else {
+                    sig = try fs.cwd().openFile(args[i], .{});
+                }
+            } else {
+                fatal("unrecognized parameter: '{s}'", .{arg});
+            }
+        } else if (file != null) {
+            fatal("more than one file to verify signature of provided", .{});
+        } else if (mem.eql(u8, args[i], "-")) {
+            file = stdin;
+            file_stdin = true;
+        } else {
+            file = try fs.cwd().openFile(args[i], .{});
+        }
+    }
+
+    if (file == null) {
+        info("{s}", .{usage_verify});
+        fatal("no file to verify signature of provided", .{});
+    }
+
+    if (key == null) {
+        info("{s}", .{usage_verify});
+        fatal("no key to verify signature with provided", .{});
+    }
+
+    if (sig == null) {
+        info("{s}", .{usage_verify});
+        fatal("no file to generate a signature for provided", .{});
+    }
+
+    if (two(&.{ file_stdin, key_stdin, sig_stdin })) {
+        fatal("can't use stdin for reading multiple files", .{});
+    }
+
+    const content = file.?.readToEndAlloc(arena, std.math.maxInt(usize)) catch {
+        fatal("could not read file to generate signature for", .{});
+    };
+    const signature_b64 = sig.?.readToEndAlloc(arena, std.math.maxInt(usize)) catch {
+        fatal("could not read signature", .{});
+    };
+    var k = readKeyFile(arena, key.?);
+    defer k.wipe();
+
+    const trimmed_signature_b64 = mem.trim(u8, signature_b64, " \n\t\r");
+    const decoded_len = std.base64.standard.Decoder.calcSizeForSlice(trimmed_signature_b64) catch {
+        fatal("invalid signature encoding", .{});
+    };
+    if (decoded_len != std.crypto.sign.Ed25519.signature_length)
+        fatal("invalid signature length", .{});
+    const signature = try arena.alloc(u8, decoded_len);
+
+    _ = std.base64.standard.Decoder.decode(signature, trimmed_signature_b64) catch {
+        fatal("invalid signature encoding", .{});
+    };
+    k.verify(content, signature[0..std.crypto.sign.Ed25519.signature_length].*) catch {
+        fatal("bad signature", .{});
+    };
+
+    try stdout.writeAll("good signature\n");
+}
+
+const PrefixKeyGenerator = struct {
+    ty: nkeys.PublicPrefixByte,
+    prefix: []const u8,
+    allocator: *Allocator,
+    done: std.atomic.Bool,
+
+    const Self = @This();
+
+    pub fn init(allocator: *Allocator, ty: nkeys.PublicPrefixByte, prefix: []const u8) Self {
+        return .{
+            .ty = ty,
+            .prefix = prefix,
+            .allocator = allocator,
+            .done = std.atomic.Bool.init(false),
+        };
+    }
+
+    fn generatePrivate(self: *Self) void {
+        while (true) {
+            if (self.done.load(.SeqCst)) return;
+
+            var kp = nkeys.SeedKeyPair.init(self.ty) catch |e| fatal("could not generate key pair: {e}", .{e});
+            defer kp.wipe();
+            var public_key = kp.publicKey() catch |e| fatal("could not generate public key: {e}", .{e});
+            if (!mem.startsWith(u8, public_key[1..], self.prefix)) continue;
+
+            if (self.done.xchg(true, .SeqCst)) return; // another thread is already done
+
+            info("{s}", .{kp.seed});
+            info("{s}", .{public_key});
+
+            return;
+        }
+    }
+
+    pub usingnamespace if (builtin.single_threaded) struct {
+        pub fn generate(self: *Self) !void {
+            return self.generatePrivate();
+        }
+    } else struct {
+        pub fn generate(self: *Self) !void {
+            var cpu_count = try std.Thread.cpuCount();
+            var threads = try self.allocator.alloc(*std.Thread, cpu_count);
+            defer self.allocator.free(threads);
+            for (threads) |*thread| thread.* = try std.Thread.spawn(Self.generatePrivate, self);
+            for (threads) |thread| thread.wait();
+        }
+    };
+};
+
+pub fn readKeyFile(allocator: *Allocator, file: fs.File) nkeys.Key {
+    var bytes = file.readToEndAlloc(allocator, std.math.maxInt(usize)) catch fatal("could not read key file", .{});
+
+    var iterator = mem.split(bytes, "\n");
+    while (iterator.next()) |line| {
+        if (nkeys.isValidEncoding(line) and line.len == nkeys.text_seed_len) {
+            var k = nkeys.fromText(line) catch continue;
+            defer k.wipe();
+            allocator.free(bytes);
+            return k;
+        }
+    }
+
+    fatal("could not find a valid key", .{});
+}
+
+fn two(slice: []const bool) bool {
+    var one = false;
+    for (slice) |x| if (x and one) {
+        return true;
+    } else {
+        one = true;
+    };
+    return false;
+}
+
+fn toUpper(allocator: *Allocator, slice: []const u8) ![]u8 {
+    const result = try allocator.alloc(u8, slice.len);
+    for (slice) |c, i| result[i] = ascii.toUpper(c);
+    return result;
+}
+
+test {
+    testing.refAllDecls(@This());
+}