1 use clap::{App, Arg, ArgMatches, SubCommand};
2 use data_encoding::BASE32_NOPAD;
3 use fs;
4 
5 // Create arguments for `add` subcommand
subcommand<'a, 'b>() -> App<'a, 'b>6 pub fn subcommand<'a, 'b>() -> App<'a, 'b> {
7     SubCommand::with_name("add")
8         .about("Add a new account")
9         .arg(
10             Arg::with_name("account")
11                 .required(true)
12                 .help("Name of the account"),
13         )
14         .arg(
15             Arg::with_name("key")
16                 .required(true)
17                 .help("Secret key of the OTP")
18                 .validator(is_base32_key),
19         )
20         .arg(
21             Arg::with_name("totp")
22                 .long("totp")
23                 .conflicts_with("hotp")
24                 .help("Time based account (default)"),
25         )
26         .arg(
27             Arg::with_name("hotp")
28                 .long("hotp")
29                 .help("Counter based account"),
30         )
31         .arg(
32             Arg::with_name("algorithm")
33                 .short("a")
34                 .long("algorithm")
35                 .takes_value(true)
36                 .possible_values(&["SHA1", "SHA256", "SHA384", "SHA512", "SHA512_256"])
37                 .value_name("ALGORITHM")
38                 .help("Algorithm to use to generate the OTP code"),
39         )
40 }
41 
42 // Validate key provided in arguments is a valid base32 encoding
is_base32_key(value: String) -> Result<(), String>43 fn is_base32_key(value: String) -> Result<(), String> {
44     let value = value.to_uppercase();
45     match BASE32_NOPAD.decode(value.as_bytes()) {
46         Ok(_) => Ok(()),
47         Err(_) => Err(String::from("the key is not a valid base32 encoding")),
48     }
49 }
50 
51 // Implementation for the `add` subcommand
run(args: &ArgMatches)52 pub fn run(args: &ArgMatches) {
53     let totp = !args.is_present("hotp");
54     let hash_function = match args.value_of("algorithm") {
55         Some(algorithm) => algorithm,
56         None => "SHA1",
57     };
58     let account_name = args.value_of("account").unwrap();
59     let key = args.value_of("key").unwrap().to_uppercase();
60     match fs::read() {
61         Ok(mut accounts) => {
62             let counter = if !totp { Some(0) } else { None };
63             let account = fs::Account {
64                 key,
65                 totp,
66                 hash_function: hash_function.to_string(),
67                 counter,
68             };
69 
70             if accounts.get(account_name).is_some() {
71                 println!("Account already exists");
72             } else {
73                 accounts.insert(account_name.to_string(), account);
74                 match fs::write(&accounts) {
75                     Ok(_) => println!("Account successfully created"),
76                     Err(err) => eprintln!("{}", err),
77                 };
78             }
79         }
80         Err(err) => eprintln!("{}", err),
81     };
82 }
83 
84 #[cfg(test)]
85 mod tests {
86     #[test]
test_is_base32_key()87     fn test_is_base32_key() {
88         let result = super::is_base32_key(String::from("12123EQ"));
89         assert!(result.is_err());
90         assert_eq!(
91             result.err(),
92             Some(String::from("the key is not a valid base32 encoding"))
93         );
94 
95         let result = super::is_base32_key(String::from("4AZJFQFIGYM2KMTOO72I6FAOZ6ZFWJR6"));
96         assert!(result.is_ok());
97         assert_eq!(result.ok(), Some(()));
98     }
99 }
100