const std = @import("std");
const builtin = @import("builtin");
const Io = std.Io;

const totp = @import("totp");
const hidapi = @import("hidapi");

const CLIError = error {
    MissingParameterDeviceId,
    MissingParameterVendorId,
    MissingKeyParameter,
    InvalidBase32,
    PasswordTooLong,
    UnknowDevice
};

/// Receives the base32 encoded code decodes it and stores the output in output. Returns the size of output.
pub fn base32decode(code: []const u8, output: []u8) !usize {
    var position:usize = 0;
    var value:u64 = 0;
    var bits:usize = 0;
    if (std.mem.findScalar(u8, code, ' ')) |_| {
        std.log.warn("Key contains whitespace which is not valid base32, ignoring",.{});
    }
    for(code,0..) |b,pos| {
        if (b==' ') {
            continue;
        }
        else if (b>'Z' or b < '2' or (b>'9' and b < 'A')) {
            std.log.err("The character at position {}: {c} is invalid base32",.{pos,b});
            return CLIError.InvalidBase32;
        }
        const val = if (b >= 'A') b-'A' else (b-'2')+26;

        value = (value<<5) | val;
        bits += 5;
        if (bits >= 8) {
            bits -= 8;
            output[position] = @as(u8,@intCast((value>>@intCast(bits))&0xff));
            position+=1;
        }
    }
    return position;
}

/// Prints the given prompt before asking for user input. reader is assumed to be a reader for stdin.
/// if with_delim is set to true returns the user input in addition to the final carriage return
pub fn read_password(prompt: []const u8, reader: *std.Io.Reader, with_delim: bool) ![]u8 {
    std.debug.print("{s}",.{prompt});
    const stdin = std.Io.File.stdin();
    const original_termios = try std.posix.tcgetattr(stdin.handle);
    var new_termios = original_termios;
    new_termios.lflag.ECHO = false;
    try std.posix.tcsetattr(stdin.handle, .FLUSH, new_termios);
    defer std.posix.tcsetattr(stdin.handle, .FLUSH, original_termios) catch {};

    const password = if (with_delim) try reader.takeDelimiterInclusive('\n') else (try reader.takeDelimiter('\n')).?;
    std.debug.print("\n",.{});
    return password;
}

pub fn main(init: std.process.Init) !void {
    var iter = init.minimal.args.iterate();
    _ = iter.next();

    var buffer: [256]u8 = undefined;
    if(iter.next()) |arg| {
        const stdin = std.Io.File.stdin();
        var reader = stdin.reader(init.io, &buffer);
        if (std.mem.eql(u8,arg, "add")) {
            var output_buffer: [128]u8 = undefined;
            const key_enc = if (iter.next()) |key| key else return CLIError.MissingKeyParameter;
            const key_len = try base32decode(key_enc, &output_buffer);
            const key = try init.gpa.dupe(u8,output_buffer[0..key_len]);
            defer init.gpa.free(key);
            const password = try read_password("Password: ", &reader.interface, false);

            var output: [32]u8 = undefined;
            std.crypto.hash.sha2.Sha256.hash(password, &output, .{});
            std.crypto.aead.chacha_poly.XChaCha20Poly1305.encrypt(output_buffer[0..key.len], output_buffer[key.len..][0..16], key, "", totp.nonce, output);
            const slice = std.base64.standard_no_pad.Encoder.encode(&buffer, output_buffer[0..key.len+16]);
            std.debug.print("Add this struct entry to the configuration {{KEY,6,{s}}}\n",.{slice});
        }
        else if (std.mem.eql(u8,arg, "decode")) {
            const key_enc = if (iter.next()) |key| key else return CLIError.MissingKeyParameter;
            const password = try read_password("Password: ", &reader.interface, true);
            password[password.len-1] = 0;
            var output: [32]u8 = undefined;
            _ = totp.decode_secret_and_generate_2fa(password.ptr, key_enc,@as(u64,@intCast(std.Io.Timestamp.now(init.io,.real).toSeconds())), 6, output[0..]);
            std.debug.print("2FA-Key: {s}\n", .{output[0..6]});
        }
        else if (std.mem.eql(u8,arg, "reencode")) {
            var output_buffer: [128]u8 = undefined;
            const key_enc = if (iter.next()) |key| key else return CLIError.MissingKeyParameter;
            const password = try read_password("Old Password: ", &reader.interface, true);
            password[password.len-1] = 0;
            totp.sha256(password.ptr, output_buffer[0..]);
            var secret_2fa: [128]u8 = undefined;
            const result = totp.decode_2fa_key(key_enc, output_buffer[0..], secret_2fa[0..]);
            if (result < 0) {
                std.debug.print("Error decoding\n",.{});
                return;
            }
            const length:usize = @intCast(result);
            const newpassword = try read_password("New Password: ", &reader.interface, false);
            var output: [32]u8 = undefined;
            std.crypto.hash.sha2.Sha256.hash(newpassword, &output, .{});
            std.crypto.aead.chacha_poly.XChaCha20Poly1305.encrypt(output_buffer[0..length], output_buffer[length..][0..16], secret_2fa[0..length], "", totp.nonce, output);
            const slice = std.base64.standard_no_pad.Encoder.encode(&buffer, output_buffer[0..length+16]);
            std.debug.print("Add this struct entry to the configuration {{KEY,6,{s}}}\n",.{slice});
        }
        else if(std.mem.eql(u8,arg,"nonce")) {
            var op = try std.Io.Dir.openFileAbsolute(init.io,"/dev/urandom",.{});
            var reader_nonce = op.reader(init.io, &buffer);
            var bytes:[24]u8 = undefined;
            _ = try reader_nonce.interface.readSliceShort(&bytes);
            std.debug.print("Nonce {any}\n",.{bytes});
        }
        else if(std.mem.eql(u8,arg,"hid")) {
            const vendorid = if (iter.next()) |vendorid| try std.fmt.parseInt(u16, vendorid, 16) else return CLIError.MissingParameterVendorId;
            const deviceid = if (iter.next()) |devid| try std.fmt.parseInt(u16, devid, 16) else return CLIError.MissingParameterDeviceId;
            _ = hidapi.hid_init();
            defer _ = hidapi.hid_exit();
            var enumerate = hidapi.hid_enumerate(vendorid, deviceid);
            while(enumerate) |interface| {
                if (interface.*.usage == 0x61 and interface.*.usage_page == 0xFF60) {
                    break;
                }
                enumerate = enumerate.*.next;
            }
            if(enumerate == null) {
                std.debug.print("Could not find device.\n", .{});
                return CLIError.UnknowDevice;
            }
            const dev: ?*hidapi.hid_device = hidapi.hid_open_path(enumerate.*.path);
            std.debug.print("Found device at {s}\n", .{enumerate.*.path});
            if (dev) |device_handle| {
                defer hidapi.hid_close(device_handle);
                var write_buffer: [33]u8 = undefined;
                // Report ID
                write_buffer[0] = 0;

                const now:u64 = @intCast(std.Io.Timestamp.now(init.io, .real).toSeconds());

                // Payload
                inline for(0..8) |x| {
                    write_buffer[1+x] = @intCast((now>>(x*8))&0xFF);
                }
                const password = try read_password("Password: ", &reader.interface, false);
                // Password length is bounded since HIDAPI report length is bounded
                if (password.len > 23) {
                    std.debug.print("Only passwords lengths of at most 23 bytes supported.\n", .{});
                    return CLIError.PasswordTooLong;
                }

                @memcpy(write_buffer[9..9+password.len], password);
                write_buffer[9+password.len] = 0;
                const res2 = hidapi.hid_write(device_handle, &write_buffer, 33);
                if(res2 == 33) {
                        std.debug.print("Set up keyboard.\n", .{});
                }
            }
            else {
                std.debug.print("Could not open device with vendor id: {x}, device id: {x}\n", .{vendorid, deviceid});
            }
        }
    }
    else {
        std.debug.print("Usage: totp [command]\nCommunicates and sets up time base one time passwords for programmable keyboards.\nCommands:\n\tadd key\n\t\tAdds a new base32 encoded key.\n\n\tdecode key\n\t\tDecodes the key and generates one time password based on it.\n\n\tnonce\n\t\tGenerates a new nonce for the symmetric cipher.\n\n\thid vendorid deviceid\n\t\tSets up HID device. IDS are interpreted as base 16 numbers.\n",.{});
    }
}


test "base32 decode" {
    var output: [20]u8 = undefined;
    const encoded = "AABBDDEEFFHH";
    const size = try base32decode(encoded, &output);

    const encoded_err = "AABBDDEEFFHH--";
    try std.testing.expectError(CLIError.InvalidBase32,base32decode(encoded_err,&output));

    var output_2: [20]u8 = undefined;
    const encoded_spaces = "AABBD DEEFF HH";
    const size2 = try base32decode(encoded_spaces, &output_2);
    try std.testing.expectEqual(size, size2);
    try std.testing.expect(std.mem.eql(u8, output[0..size], output_2[0..size2]));
}