Skip to content

Commit

Permalink
Yes
Browse files Browse the repository at this point in the history
  • Loading branch information
shavit committed Jun 29, 2023
0 parents commit 0afb90f
Show file tree
Hide file tree
Showing 7 changed files with 404 additions and 0 deletions.
7 changes: 7 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
zig-cache/
zig-out/
/release/
/debug/
/build/
/build-*/
/docgen_tmp/
32 changes: 32 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
# TOTP

> Generate TOTP codes RFC-6238
# CLI Quick Start

Add a provider and generate code:
```
$ zotp a <server> <token>
$ zotp g <server>
123456
```

List all providers:
```
$ zotp l
$ zotp list
```

Delete a provider:
```
$ zotp delete <server>
$ zotp d <server>
```

Uninstall:
```
$ zotp uninstall
```
This will delete `.config/zotp`

49 changes: 49 additions & 0 deletions build.zig
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
const std = @import("std");

pub fn build(b: *std.Build) void {
const target = b.standardTargetOptions(.{});
const optimize = b.standardOptimizeOption(.{});

const pkg_base32 = b.addModule("base32", .{
.source_file = .{ .path = "../zig-base32/src/base32.zig" },
});

const lib = b.addStaticLibrary(.{
.name = "totp",
.root_source_file = .{ .path = "src/totp.zig" },
.target = target,
.optimize = optimize,
});
lib.addModule("base32", pkg_base32);
b.installArtifact(lib);

const unit_tests = b.addTest(.{
.root_source_file = .{ .path = "src/tests.zig" },
.target = target,
.optimize = optimize,
});

const run_unit_tests = b.addRunArtifact(unit_tests);
const test_step = b.step("test", "Run library unit tests");
test_step.dependOn(&run_unit_tests.step);

// Add executable
const exe = b.addExecutable(.{
.name = "zotp",
.root_source_file = .{ .path = "src/cli.zig" },
.target = target,
.optimize = optimize,
});
exe.addModule("base32", pkg_base32);

b.installArtifact(exe);

const run_cmd = b.addRunArtifact(exe);
run_cmd.step.dependOn(b.getInstallStep());
if (b.args) |args| {
run_cmd.addArgs(args);
}

const run_step = b.step("run", "Run the command line app");
run_step.dependOn(&run_cmd.step);
}
139 changes: 139 additions & 0 deletions src/cli.zig
Original file line number Diff line number Diff line change
@@ -0,0 +1,139 @@
const std = @import("std");
const testing = std.testing;
const mem = std.mem;
const process = std.process;
const path = std.fs.path;
const Allocator = std.mem.Allocator;
const Dir = std.fs.Dir;
const totp = @import("totp.zig");
const storage_ = @import("storage.zig");
const Storage = storage_.Storage;
const Provider = storage_.Provider;

const CmdType = enum {
add,
delete,
list,
generate,
uninstall,
invalid,
};

fn readCmdType(a: []const u8) ?CmdType {
const table = std.ComptimeStringMap(CmdType, .{
.{ "add", .add },
.{ "a", .add },
.{ "delete", .delete },
.{ "d", .delete },
.{ "list", .list },
.{ "l", .list },
.{ "generate", .generate },
.{ "g", .generate },
.{ "uninstall", .uninstall },
});

return table.get(a);
}

fn cmdAdd(storage: *Storage, args: *process.ArgIterator) !void {
if (args.inner.count < 4) return error.NotEnoughArguments;
const name = args.next() orelse unreachable;
const token = args.next() orelse unreachable;

var p = Provider{ .name = name, .token = token };
try storage.put(p);
try storage.commit();
}

fn cmdDelete(storage: *Storage, args: *process.ArgIterator) !void {
if (args.inner.count < 3) return error.NotEnoughArguments;
const name = args.next() orelse unreachable;

var provider = storage.get(name);
if (provider == null) return error.ProviderNotFound;

try storage.delete(&provider.?);
try storage.commit();
}

fn cmdList(storage: *Storage, args: *process.ArgIterator) !void {
if (args.inner.count > 2) return error.InvalidArgument;
storage.list_print();
}

fn cmdGenerate(storage: *Storage, args: *process.ArgIterator) !void {
const arg_name = args.next() orelse unreachable;

if (storage.get(arg_name)) |p| {
const code: u32 = try totp.generate(p.token);
var buf: [6]u8 = undefined;
const output = try std.fmt.bufPrint(&buf, "{d}", .{code});
try write_stdout(output);
} else {
return error.ProviderNotFound;
}
}
fn cmdUninstall(storage: *Storage, args: *process.ArgIterator) !void {
if (args.inner.count > 2) return error.InvalidArgument;

try std.fs.cwd().deleteFile(storage.config_path);
if (std.fs.path.dirname(storage.config_path)) |x| {
try std.fs.deleteDirAbsolute(x);
}
}

fn start_cli(args: *process.ArgIterator, a: []const u8) !void {
if (mem.eql(u8, a, "-h")) {
println(help);
std.process.exit(0);
}

var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
defer arena.deinit();
const allocator = arena.allocator();

const config_path = try Storage.find_config(allocator);
var storage: Storage = try Storage.init(allocator, config_path);

switch (readCmdType(a) orelse .invalid) {
.list => try cmdList(&storage, args),
.generate => try cmdGenerate(&storage, args),
.add => try cmdAdd(&storage, args),
.delete => try cmdDelete(&storage, args),
.uninstall => try cmdUninstall(&storage, args),
else => println(help),
}
}

pub fn main() !void {
var args = process.args();
if (args.inner.count <= 1) {
println(help);
goodbye("Missing arguments", .{});
}
_ = args.next(); // program name

const a: []const u8 = args.next() orelse unreachable;
try start_cli(&args, a);
}

const help =
"Usage: zotp [command] [options]\n" ++ "\t-h Prints this message\n" ++ "\nOptions: list generate add delete uninstall\n";

fn println(text: []const u8) void {
std.debug.print("{s}\n", .{text});
}

fn write_stdout(msg: []const u8) !void {
const stdout_file = std.io.getStdOut().writer();
var bw = std.io.bufferedWriter(stdout_file);
const writer = bw.writer();

try writer.print("{s}\n", .{msg});
try bw.flush();
}

fn goodbye(comptime format: []const u8, args: anytype) noreturn {
std.log.err(format, args);
std.process.exit(1);
}
120 changes: 120 additions & 0 deletions src/storage.zig
Original file line number Diff line number Diff line change
@@ -0,0 +1,120 @@
const std = @import("std");
const path = std.fs.path;
const mem = std.mem;
const io = std.io;
const fs = std.fs;
const StringHashMap = std.StringHashMap;

pub const Provider = struct {
name: []const u8,
token: []const u8,
};

pub const Storage = struct {
const Self = @This();

allocator: mem.Allocator,
providers: StringHashMap(Provider),
config_path: []const u8,

pub fn init(allocator: mem.Allocator, config_path: []const u8) !Self {
var storage = Self{
.allocator = allocator,
.providers = StringHashMap(Provider).init(allocator),
.config_path = config_path,
};
errdefer storage.providers.deinit();
try storage.load();

return storage;
}

fn load(self: *Self) !void {
var f = try fs.cwd().openFile(self.config_path, .{ .mode = .read_only });
defer f.close();

var reader = io.bufferedReader(f.reader());
var stream = reader.reader();
var buf: [1024]u8 = undefined;
while (try stream.readUntilDelimiterOrEof(&buf, '\n')) |l| {
var line = try self.allocator.alloc(u8, l.len);
@memcpy(line, l);

const p = try parse_fs_line(line);
try self.providers.put(p.name, p);
}
}

pub fn find_config(allocator: mem.Allocator) ![]const u8 {
const HOME = std.os.getenv("HOME") orelse "~";
const config_dir = try std.fmt.allocPrint(allocator, "{s}{u}{s}{u}{s}", .{ HOME, path.sep, ".config", path.sep, "zotp" });
defer allocator.free(config_dir);

const config_path = try std.fmt.allocPrint(allocator, "{s}{u}{s}", .{ config_dir, path.sep, "zotp.conf" });
_ = fs.cwd().statFile(config_path) catch blk: {
_ = fs.makeDirAbsolute(config_dir) catch {};
_ = fs.cwd().createFile(config_path, .{ .truncate = true }) catch {};
break :blk try fs.cwd().statFile(config_path);
};

return config_path;
}

fn parse_fs_line(line: []const u8) !Provider {
var it = mem.split(u8, line, ":");
const provider = it.next() orelse unreachable;
if (line.len <= provider.len + 1) return error.CorruptedFile;

return Provider{
.name = provider,
.token = line[(provider.len + 1)..(line.len)],
};
}

pub fn get(self: *Self, name: []const u8) ?Provider {
return self.providers.get(name);
}

pub fn put(self: *Self, p: Provider) !void {
try self.providers.put(p.name, p);
}

pub fn delete(self: *Self, p: *Provider) !void {
_ = self.providers.remove(p.name);
}

pub fn list_print(self: *Self) void {
var it = self.providers.iterator();
var i: u8 = 1;
while (it.next()) |x| {
std.debug.print("{d}/{d}: {s}\n", .{ i, self.providers.count(), x.key_ptr.* });
i += 1;
}
it.index = 0;
}

pub fn commit(self: *Self) !void {
if (self.config_path.len == 0) {
std.debug.print("no config path is set\n", .{});
}

const mem_req = self.providers.count() * 128;
var buf: []u8 = try self.allocator.alloc(u8, mem_req);
var it = self.providers.iterator();
var i: usize = 0;
while (it.next()) |p| {
const pv = p.value_ptr.*;
@memcpy(buf[i .. i + pv.name.len], pv.name);
i += pv.name.len;
@memcpy(buf[i .. i + 1], ":");
i += 1;
@memcpy(buf[i .. i + pv.token.len], pv.token);
i += pv.token.len;

@memcpy(buf[i .. i + 1], "\n");
i += 1;
}

try fs.cwd().writeFile(self.config_path, buf[0..i]);
}
};
5 changes: 5 additions & 0 deletions src/tests.zig
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
comptime {
_ = @import("totp.zig");
_ = @import("cli.zig");
_ = @import("storage.zig");
}
Loading

0 comments on commit 0afb90f

Please sign in to comment.