extern crate base32; extern crate clap; extern crate oath; extern crate rand; extern crate serde_json; #[macro_use] extern crate serde_derive; use clap::{App, Arg, SubCommand}; use std::collections::HashMap; use std::fs::{create_dir_all, File, OpenOptions}; use std::io::ErrorKind; use std::path::{Path, PathBuf}; use std::thread; use std::time::Duration; const CONFIG_PATH: &str = ".rustotpony/db.json"; const DATABASE_VERSION: u8 = 1; struct RusTOTPony { database: DB, applications: HashMap, } impl RusTOTPony { fn new(db: DB) -> RusTOTPony { RusTOTPony { applications: db.get_applications(), database: db, } } fn add_application(&mut self, name: &str, secret: &str) -> Result<(), String> { if let Some(secret_bytes) = GenApp::base32_to_bytes(secret) { let new_app = GenApp::new(name, secret, secret_bytes); if self.applications.contains_key(name) { Err(format!("Application with name '{}' already exists!", name)) } else { &self.applications.insert(String::from(name), new_app); Ok(()) } } else { return Err(String::from("Couldn't decode secret key")) } } fn delete_application(&mut self, name: &str) -> Result<(), String> { if let Some(_) = self.applications.remove(name) { Ok(()) } else { Err(format!( "Application with the name '{}' doesn't exist", name )) } } fn rename_application(&mut self, name: &str, newname: &str) -> Result<(), String> { if let Some(app) = self.applications.get_mut(name) { app.name = String::from(newname); Ok(()) } else { Err(format!("Application '{}' wasn't found", name)) } } fn get_applications(&self) -> Result<&HashMap, String> { if self.applications.len() == 0 { Err(String::from("There are no applications")) } else { Ok(&self.applications) } } fn get_application(&self, name:&str) -> Result<&GenApp, String> { if let Some(app) = self.applications.get(name) { Ok(app) } else { Err(format!("Application '{}' wasn't found", name)) } } fn delete_all_applications(&mut self) { self.applications = HashMap::new(); } fn flush(&self) { &self.database.save_applications(&self.applications); } } trait Database { fn get_applications(&self) -> HashMap; fn save_applications(&self, applications: &HashMap); } impl Database for JsonDatabase { fn get_applications(&self) -> HashMap { let db_content = self.read_database_file(); db_content.content.applications } fn save_applications(&self, applications: &HashMap) { let mut db_content = Self::get_empty_schema(); db_content.content.applications = applications.clone(); self.save_database_file(db_content); } } #[derive(Serialize, Deserialize)] struct JsonDatabaseSchema { version: u8, content: DatabaseContentSchema, } #[derive(Serialize, Deserialize)] struct DatabaseContentSchema { applications: HashMap, } struct JsonDatabase { file_path: PathBuf, } impl JsonDatabase { fn new(path: PathBuf) -> JsonDatabase { JsonDatabase { file_path: path } } fn read_database_file(&self) -> JsonDatabaseSchema { let file = match File::open(&self.file_path) { Ok(f) => f, Err(ref err) if err.kind() == ErrorKind::NotFound => return Self::get_empty_schema(), Err(err) => panic!("There was a problem opening file: {:?}", err), }; serde_json::from_reader(file).expect("Couldn't parse JSON from database file") } fn save_database_file(&self, content: JsonDatabaseSchema) { let file = match self.open_database_file_for_write() { Ok(f) => f, Err(ref err) if err.kind() == ErrorKind::NotFound => self.create_database_file() .expect("Couldn't create database file"), Err(err) => panic!("Couldn't open database file: {:?}", err), }; serde_json::to_writer(file, &content).expect("Couldn't write JSON data to database file"); } fn create_database_file(&self) -> Result { let dir = std::env::home_dir().unwrap_or(PathBuf::from(".")); if let Some(parent_dir) = Path::new(&self.file_path).parent() { let dir = dir.join(parent_dir); create_dir_all(dir)?; } self.open_database_file_for_write() } fn open_database_file_for_write(&self) -> Result { OpenOptions::new() .write(true) .truncate(true) .create(true) .open(&self.file_path) } fn get_empty_schema() -> JsonDatabaseSchema { JsonDatabaseSchema { version: DATABASE_VERSION, content: DatabaseContentSchema { applications: HashMap::new(), }, } } } #[derive(Serialize, Deserialize, Debug, Clone)] struct GenApp { name: String, secret: String, username: String, secret_bytes: Vec, } impl GenApp { fn new(name: &str, secret: &str, secret_bytes: Vec) -> Self { GenApp { name: String::from(name), secret: String::from(secret), username: String::from(""), secret_bytes: secret_bytes, } } fn get_code(&self) -> u64 { Self::totp(&self.secret_bytes) } fn base32_to_bytes(secret:&str) -> Option> { base32::decode(base32::Alphabet::RFC4648 {padding: false}, secret) } fn totp(secret_bytes: &[u8]) -> u64 { oath::totp_raw_now(&secret_bytes, 6, 0, 30, &oath::HashType::SHA1) } } struct Cli { app: RusTOTPony, } impl Cli { fn new() -> Self { let db = JsonDatabase::new(Self::get_database_path()); Self { app: RusTOTPony::new(db), } } fn run(&mut self) { match self.get_cli_api_matches().subcommand() { ("dash", Some(_)) => { self.show_dashboard(); } ("list", Some(_)) => { self.show_applications_list(false); } ("show-all", Some(_)) => { self.show_applications_list(true); } ("show", Some(sub_app)) => { let app_name: &str = sub_app .value_of("APPNAME") .expect("Couldn't read APPNAME for 'show' command"); self.show_application(app_name); } ("add", Some(sub_app)) => { let app_name: &str = sub_app .value_of("APPNAME") .expect("Couldn't read APPNAME for 'add' command"); let key: &str = sub_app .value_of("KEY") .expect("Couldn't read KEY for 'add' command"); self.add_application(app_name, key); } ("delete", Some(sub_app)) => { let app_name: &str = sub_app .value_of("APPNAME") .expect("Couldn't read APPNAME for 'delete' command"); self.delete_application(app_name); } ("rename", Some(sub_app)) => { let app_name: &str = sub_app .value_of("APPNAME") .expect("Couldn't read APPNAME for 'rename' command"); let new_name: &str = sub_app .value_of("NEWNAME") .expect("Couldn't read NEWNAME for 'rename' command"); self.rename_application(app_name, new_name); } ("eradicate", Some(_)) => { self.eradicate_database(); } _ => { self.show_applications_list(true); } } } fn get_cli_api_matches(&self) -> clap::ArgMatches<'static> { App::new("Main") .version("0.1.0") .author("German Lashevich ") .about("TOTP code generator written with Rust") .subcommand( SubCommand::with_name("dash").about("Shows realtime dashboard with all generators"), ) .subcommand(SubCommand::with_name("list").about("Lists all generators")) .subcommand( SubCommand::with_name("show-all") .about("Shows all generators with their's current values"), ) .subcommand( SubCommand::with_name("show") .about("Shows generator with it's current value") .arg(Arg::with_name("APPNAME").required(true)), ) .subcommand( SubCommand::with_name("add") .about("Adds new generator") .arg(Arg::with_name("APPNAME").required(true)) .arg(Arg::with_name("KEY").required(true)), ) .subcommand( SubCommand::with_name("delete") .about("Deletes generator") .arg(Arg::with_name("APPNAME").required(true)), ) .subcommand( SubCommand::with_name("rename") .about("Renames generator") .arg(Arg::with_name("APPNAME").required(true)) .arg(Arg::with_name("NEWNAME").required(true)), ) .subcommand(SubCommand::with_name("eradicate").about("Deletes all generators")) .get_matches() } fn get_database_path() -> PathBuf { let home = std::env::home_dir().unwrap_or(PathBuf::from(".")); home.join(Path::new(CONFIG_PATH)) } fn show_dashboard(&self) { match self.app.get_applications(){ Ok(apps) => { let mut is_first_iteration = true; let lines_count = apps.len(); loop { if is_first_iteration { is_first_iteration = false; } else { print!("\x1B[{}A", lines_count); } for (_, app) in apps { println!{"{} {}", app.get_code(), app.name}; } thread::sleep(Duration::from_millis(1000)); } }, Err(err) => println!("{}", err), } } fn show_applications_list(&self, _: bool) { // TODO Create Table structure with HashMap as follows and metadata about columns - width, titles, names let mut output_table: HashMap<&str, Vec> = HashMap::new(); let mut applications_count = 0; let apps = match self.app.get_applications() { Ok(v) => v, Err(e) => { println!("{}", e); return; }, }; for (_, application) in apps { applications_count += 1; output_table .entry("name") .or_insert(Vec::new()) .push(application.name.clone()); output_table .entry("key") .or_insert(Vec::new()) .push(application.secret.clone()); output_table .entry("username") .or_insert(Vec::new()) .push(application.username.clone()); } let name_max_length = output_table["name"] .iter() .fold("name".len(), |max, val| std::cmp::max(max, val.len())); let key_max_length = output_table["key"] .iter() .fold("key".len(), |max, val| std::cmp::max(max, val.len())); let username_max_length = output_table["username"] .iter() .fold("username".len(), |max, val| std::cmp::max(max, val.len())); let header_row_delimiter = format!( "+-{}-+-{}-+-{}-+", "-".repeat(name_max_length), "-".repeat(key_max_length), "-".repeat(username_max_length) ); println!("{}", header_row_delimiter); println!( "| {name: { self.app.flush(); println!("New application created: {}", name) }, Err(err) => println!("{} Aborting…", err), } } fn delete_application(&mut self, name: &str) { match self.app.delete_application(name) { Ok(_) => { self.app.flush(); println!("Application '{}' successfully deleted", name) }, Err(err) => println!("Couldn't delete application '{}': {}", name, err), }; } fn rename_application(&mut self, name: &str, newname: &str) { match self.app.rename_application(name, newname) { Ok(_) => { self.app.flush(); println!("Application '{}' successfully renamed to '{}'", name, newname) }, Err(err) => println!("Couldn't rename application '{}': {}", name, err), }; } fn eradicate_database(&mut self) { self.app.delete_all_applications(); self.app.flush(); println!("Done."); } } fn main() { Cli::new().run(); }