From 6c6d5a51a89cf8eedbbe8f042ce34c77107b515a Mon Sep 17 00:00:00 2001 From: Andrey Golovizin Date: Sat, 15 Apr 2023 13:34:25 +0200 Subject: [PATCH] Save tokens for each machine --- src/dirs.rs | 43 ++++++++++++++++++++++++++++++++++++++++++- src/keys.rs | 27 +++++++-------------------- src/main.rs | 40 ++++++++++++++++++++++++++++++++++------ 3 files changed, 83 insertions(+), 27 deletions(-) diff --git a/src/dirs.rs b/src/dirs.rs index 03a38b1..4d80e00 100644 --- a/src/dirs.rs +++ b/src/dirs.rs @@ -1,7 +1,8 @@ use directories::ProjectDirs; +use gethostname::gethostname; use once_cell::sync::OnceCell; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; const QUALIFIER: &str = "com"; const ORGANIZATION: &str = "Ero-sennin"; @@ -16,3 +17,43 @@ pub(crate) fn get_data_dir() -> PathBuf { }); project_dirs.data_dir().to_path_buf() } + +pub(crate) struct MachineConfig { + path: PathBuf, +} + +impl MachineConfig { + pub(crate) fn new(machine_name: Option<&PathBuf>) -> anyhow::Result { + let hostname: PathBuf; + let machine: &PathBuf = if let Some(machine_name) = machine_name { + machine_name + } else { + hostname = PathBuf::from(gethostname()); + &hostname + }; + let path = get_data_dir().join("machines").join(machine); + std::fs::create_dir_all(&path)?; + Ok(Self { path }) + } + + pub(crate) fn destroy(self) -> anyhow::Result<()> { + std::fs::remove_dir_all(self.path)?; + Ok(()) + } + + pub(crate) fn path(&self) -> &Path { + &self.path + } + + pub(crate) fn username_path(&self) -> PathBuf { + self.path().join("username") + } + + pub(crate) fn token_path(&self) -> PathBuf { + self.path().join("token") + } + + pub(crate) fn key_path(&self) -> PathBuf { + self.path().join("key") + } +} diff --git a/src/keys.rs b/src/keys.rs index c894dd5..227495c 100644 --- a/src/keys.rs +++ b/src/keys.rs @@ -1,14 +1,10 @@ use base64::prelude::{Engine as _, BASE64_STANDARD}; -use gethostname::gethostname; use log::debug; use x25519_dalek::{PublicKey, StaticSecret}; -use std::{ - io::Write, - path::{Path, PathBuf}, -}; +use std::{io::Write, path::Path}; -use crate::dirs::get_data_dir; +use crate::dirs::MachineConfig; const KEY_SIZE: usize = 32; @@ -46,24 +42,15 @@ pub(crate) struct WireguardKeyPair { pub private_key: Key, } -pub(crate) fn get_keys(machine: Option<&PathBuf>) -> Result { - let hostname: PathBuf; - let machine_subdir: &PathBuf = if let Some(machine) = machine { - machine - } else { - hostname = PathBuf::from(gethostname()); - &hostname - }; - let key_path = get_data_dir().join("keys"); +pub(crate) fn get_keys(machine_config: &MachineConfig) -> Result { + let key_path = machine_config.key_path(); debug!("key path = {:?}", &key_path); - std::fs::create_dir_all(&key_path)?; - let private_key_path = key_path.join(machine_subdir); - let private_key = if private_key_path.is_file() { - Key::load(&private_key_path)? + let private_key = if key_path.is_file() { + Key::load(&key_path)? } else { let key = generate_private_key(); - key.save(&private_key_path)?; + key.save(&key_path)?; key }; diff --git a/src/main.rs b/src/main.rs index faf9a11..5d7ed12 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,15 +1,15 @@ mod api; -use std::net::IpAddr; mod dirs; mod keys; use std::io::Write; +use std::net::IpAddr; use std::path::PathBuf; use clap::{Parser, Subcommand}; use log::debug; -use crate::keys::{get_keys, WireguardKeyPair}; +use crate::{keys::{get_keys, WireguardKeyPair}, dirs::MachineConfig}; /// AzireVPN client #[derive(Parser, Debug)] @@ -27,10 +27,14 @@ struct Opts { } #[derive(Parser, Debug)] -struct ConfigOpts { - location: String, +struct LoginOpts { username: String, token: String, +} + +#[derive(Parser, Debug)] +struct ConfigOpts { + location: String, #[arg(short, long)] no_dns: bool, @@ -41,6 +45,12 @@ struct ConfigOpts { #[derive(Subcommand, Debug)] enum Command { + /// Add machine token + Login(LoginOpts), + + /// Logout + Logout, + /// Checks connection status Check, @@ -55,6 +65,8 @@ fn main() -> Result<(), anyhow::Error> { env_logger::init(); let opts = Opts::parse(); match &opts.command { + Command::Login(login_opts) => login(&opts, login_opts)?, + Command::Logout => logout(&opts)?, Command::Check => check(&opts)?, Command::Locations => list_locations(&opts)?, Command::Config(get_config_opts) => get_config(&opts, get_config_opts)?, @@ -62,6 +74,19 @@ fn main() -> Result<(), anyhow::Error> { Ok(()) } +fn login(opts: &Opts, login_opts: &LoginOpts) -> Result<(), anyhow::Error> { + let machine_config = dirs::MachineConfig::new(opts.machine.as_ref())?; + std::fs::write(machine_config.username_path(), &login_opts.username)?; + std::fs::write(machine_config.token_path(), &login_opts.token)?; + Ok(()) +} + +fn logout(opts: &Opts) -> Result<(), anyhow::Error> { + let machine_config = dirs::MachineConfig::new(opts.machine.as_ref())?; + machine_config.destroy()?; + Ok(()) +} + fn check(opts: &Opts) -> Result<(), anyhow::Error> { let check_result = api::check()?; if opts.json { @@ -97,6 +122,7 @@ fn list_locations(opts: &Opts) -> Result<(), anyhow::Error> { } fn get_config(opts: &Opts, config_opts: &ConfigOpts) -> Result<(), anyhow::Error> { + let machine_config = MachineConfig::new(opts.machine.as_ref())?; let locations = api::get_locations()?; let location = locations .locations @@ -104,9 +130,11 @@ fn get_config(opts: &Opts, config_opts: &ConfigOpts) -> Result<(), anyhow::Error .find(|location| location.name == config_opts.location) .ok_or_else(|| anyhow::anyhow!("no such location: {}", config_opts.location))?; debug!("location = {:?}", &location); - let keys = get_keys(opts.machine.as_ref())?; + let keys = get_keys(&machine_config)?; + let username = std::fs::read_to_string(machine_config.username_path())?; + let token = std::fs::read_to_string(machine_config.token_path())?; debug!("keys = {:?}", &keys); - let addresses = api::add_ip(&config_opts.username, &config_opts.token, &keys.public_key)?; + let addresses = api::add_ip(&username, &token, &keys.public_key)?; write_config( &mut std::io::stdout().lock(),