// Copyright (C) 2019 Robin Krahl <robin.krahl@ireas.org>
// SPDX-License-Identifier: MIT

#![warn(missing_docs, rust_2018_compatibility, rust_2018_idioms, unused)]

//! Reads OTP configuration from a QR code and writes it to an OTP slot on a Nitrokey device.

mod error;

use std::borrow;
use std::fs;
use std::path;
use std::process;
use std::str;

use dialog::DialogBox;

use crate::error::Error;

#[derive(Debug, PartialEq)]
struct UrlData {
    otp_type: String,
    label: String,
    secret: String,
    issuer: Option<String>,
    digits: Option<String>,
    counter: Option<String>,
    period: Option<String>,
}

#[derive(Debug)]
struct Options {
    slot: u8,
    file: Option<String>,
    name: Option<String>,
}

fn parse_options() -> Result<Options, i32> {
    let mut options = Options {
        slot: 0,
        file: None,
        name: None,
    };
    let mut parser = argparse::ArgumentParser::new();
    parser.set_description(
        "Reads OTP configuration from a QR code and writes it to an OTP slot on a Nitrokey device.",
    );
    parser.refer(&mut options.slot).required().add_argument(
        "slot",
        argparse::Store,
        "The slot to write the OTP data to",
    );
    parser.refer(&mut options.file).add_argument(
        "file",
        argparse::StoreOption,
        "The file to read the QR code from",
    );
    parser.refer(&mut options.name).add_option(
        &["-n", "--name"],
        argparse::StoreOption,
        "The name to store in the OTP slot",
    );
    parser.parse_args()?;
    drop(parser);
    Ok(options)
}

fn import_qr_code() -> Result<path::PathBuf, Error> {
    let mut temp = mktemp::Temp::new_file()?;
    let path = temp.to_path_buf();

    let status = process::Command::new("import").arg(&path).status()?;

    if status.success() {
        temp.release();
        Ok(path)
    } else {
        Err(Error::from(("import", status)))
    }
}

fn decode_qr_code(path: &path::Path) -> Result<String, Error> {
    let output = process::Command::new("zbarimg")
        .arg("--quiet")
        .arg("--raw")
        .arg(path)
        .output()?;
    if output.status.success() {
        let output = String::from_utf8(output.stdout).map_err(|err| err.utf8_error())?;
        let urls = output
            .split("\n")
            .filter(|url| url.starts_with("otpauth://"))
            .collect::<Vec<&str>>();
        if urls.is_empty() {
            Err(Error::from("Could not find an otpauth QR code"))
        } else {
            if urls.len() > 1 {
                println!("Found more than otpauth QR code, using the first one.");
            }
            Ok(urls[0].to_string())
        }
    } else {
        Err(Error::from(("zbarimg", output.status)))
    }
}

fn strip_issuer_prefix(label: String, issuer: &str) -> String {
    let prefix = format!("{}:", issuer);
    if label.starts_with(&prefix) {
        let label = label.trim_left_matches(&prefix);
        let label = label.trim_left_matches(" ");
        label.to_string()
    } else {
        label
    }
}

fn parse_url(url: &str) -> Result<UrlData, Error> {
    let url = url::Url::parse(url)?;
    let scheme = url.scheme();
    if scheme != "otpauth" {
        return Err(Error::Error(format!("Unexpected URL scheme: {}", scheme)));
    }
    let otp_type = match url.host_str() {
        Some(host) => host.to_string(),
        None => return Err(Error::from("otpauth URL does not contain type")),
    };
    let label = url.path();
    let label = percent_encoding::percent_decode(label.as_bytes()).decode_utf8()?;
    let label = label.trim_start_matches("/").to_string();
    let mut secret: Option<String> = None;
    let mut issuer: Option<String> = None;
    let mut digits: Option<String> = None;
    let mut counter: Option<String> = None;
    let mut period: Option<String> = None;
    for (key, value) in url.query_pairs() {
        let field = match key {
            borrow::Cow::Borrowed("secret") => Some(&mut secret),
            borrow::Cow::Borrowed("issuer") => Some(&mut issuer),
            borrow::Cow::Borrowed("digits") => Some(&mut digits),
            borrow::Cow::Borrowed("counter") => Some(&mut counter),
            borrow::Cow::Borrowed("period") => Some(&mut period),
            _ => None,
        };
        if let Some(field) = field {
            *field = Some(value.into_owned());
        }
    }
    let label = match issuer {
        Some(ref issuer) => strip_issuer_prefix(label, issuer),
        None => label,
    };
    match secret {
        Some(secret) => Ok(UrlData {
            otp_type,
            label,
            secret,
            issuer,
            digits,
            counter,
            period,
        }),
        None => Err(Error::from("otpauth URL did not contain a secret")),
    }
}

fn query_name(label: &str, issuer: Option<&str>) -> Result<String, Error> {
    let title = "Enter OTP slot name";
    let mut text = "Please enter a name for the OTP secret:".to_string();
    text.push_str(&format!("\n\tlabel:\t{}", label));
    if let Some(issuer) = issuer {
        text.push_str(&format!("\n\tissuer:\t{}", issuer));
    };
    let text = text;
    let default = issuer.unwrap_or(label);
    let name = dialog::Input::new(text)
        .title(title)
        .default(default)
        .show()?;
    if let Some(name) = name {
        if name.is_empty() {
            Err(Error::from("The OTP name may not be empty"))
        } else {
            Ok(name.trim_end_matches(&"\n").to_string())
        }
    } else {
        Err(Error::from("You canceled the name input dialog"))
    }
}

fn store_data(slot: u8, name: &str, data: &UrlData) -> Result<(), Error> {
    let slot = slot.to_string();
    let mut args = vec!["otp", "set", slot.as_ref(), name, data.secret.as_ref()];
    args.append(&mut vec!["--algorithm", data.otp_type.as_ref()]);
    args.append(&mut vec!["--format", "base32"]);
    // TODO: parse digits
    if let Some(ref counter) = data.counter {
        args.append(&mut vec!["--counter", counter]);
    }
    if let Some(ref period) = data.period {
        args.append(&mut vec!["--time-window", period]);
    }

    let status = process::Command::new("nitrocli").args(args).status()?;
    if status.success() {
        Ok(())
    } else {
        Err(Error::from("Could not set OTP data using nitrocli"))
    }
}

fn run(options: Options) -> Result<(), Error> {
    let path = match options.file {
        Some(ref file) => path::PathBuf::from(file),
        None => import_qr_code()?,
    };
    let url = decode_qr_code(&path)?;
    let url_data = parse_url(&url)?;
    let name = match options.name {
        Some(name) => name,
        None => query_name(&url_data.label, url_data.issuer.as_ref().map(|x| &**x))?,
    };
    store_data(options.slot, &name, &url_data)?;
    if options.file.is_none() {
        fs::remove_file(&path)?;
    }
    Ok(())
}

fn main() {
    let status = match parse_options() {
        Ok(options) => match run(options) {
            Ok(()) => 0,
            Err(err) => {
                println!("{}", err);
                1
            }
        },
        Err(err) => err,
    };
    process::exit(status);
}

#[cfg(test)]
mod tests {
    #[test]
    fn parse_url() -> Result<(), super::Error> {
        let result = super::parse_url(
            "otpauth://totp/Example:alice@google.com?secret=JBSWY3DPEHPK3PXP&issuer=Example",
        )?;
        let expected = super::UrlData {
            otp_type: "totp".to_string(),
            label: "alice@google.com".to_string(),
            secret: "JBSWY3DPEHPK3PXP".to_string(),
            issuer: Some("Example".to_string()),
            digits: None,
            counter: None,
            period: None,
        };
        assert_eq!(result, expected);

        let result = super::parse_url("otpauth://totp/ACME%20Co:john.doe@email.com?secret=HXDMVJECJJWSRB3HWIZR4IFUGFTMXBOZ&issuer=ACME%20Co&algorithm=SHA1&digits=6&period=30")?;
        let expected = super::UrlData {
            otp_type: "totp".to_string(),
            label: "john.doe@email.com".to_string(),
            secret: "HXDMVJECJJWSRB3HWIZR4IFUGFTMXBOZ".to_string(),
            issuer: Some("ACME Co".to_string()),
            digits: Some("6".to_string()),
            counter: None,
            period: Some("30".to_string()),
        };
        assert_eq!(result, expected);

        assert!(super::parse_url("otpauth://totp/test?secret=blubb").is_ok());
        assert!(super::parse_url("otauth://totp/test?secret=blubb").is_err());
        assert!(super::parse_url("otpauth://totp/test").is_err());

        Ok(())
    }
}