Reuse private keys

src/main.rs

@@ -1,14 +1,23 @@
 use std::io::Write;
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
+use std::path::PathBuf;
 use std::process;
 
-use log::debug;
-
 use clap::{Parser, Subcommand};
+use directories::ProjectDirs;
+use gethostname::gethostname;
+use log::debug;
+use once_cell::sync::OnceCell;
 use serde::{Deserialize, Serialize};
 
 const BASE_URL: &str = "https://api.azirevpn.com/v2";
 
+const QUALIFIER: &str = "com";
+const ORGANIZATION: &str = "Ero-sennin";
+const APPLICATION: &str = "AzireVPN";
+
+static PROJECT_DIRS: OnceCell<ProjectDirs> = OnceCell::new();
+
 /// AzireVPN client
 #[derive(Parser, Debug)]
 #[command(version)]
@@ -32,6 +41,9 @@ struct ConfigOpts {
 
     #[arg(short = '4', long)]
     no_ipv6: bool,
+
+    #[arg(short, long)]
+    machine: Option<PathBuf>,
 }
 
 #[derive(Subcommand, Debug)]
@@ -94,6 +106,14 @@ fn main() -> Result<(), anyhow::Error> {
     Ok(())
 }
 
+fn get_data_dir() -> PathBuf {
+    let project_dirs = PROJECT_DIRS.get_or_init(|| {
+        ProjectDirs::from(QUALIFIER, ORGANIZATION, APPLICATION)
+            .expect("cannot get project data dir")
+    });
+    project_dirs.data_dir().to_path_buf()
+}
+
 fn list_locations(opts: &Opts) -> Result<(), anyhow::Error> {
     let locations: Locations = get_locations()?;
     if opts.json {
@@ -122,7 +142,7 @@ fn get_config(_opts: &Opts, config_opts: &ConfigOpts) -> Result<(), anyhow::Erro
         .find(|location| location.name == config_opts.location)
         .ok_or_else(|| anyhow::anyhow!("no such location: {}", config_opts.location))?;
     debug!("location = {:?}", &location);
-    let keys = generage_keys()?;
+    let keys = get_keys(config_opts.machine.as_ref())?;
     debug!("keys = {:?}", &keys);
     let url = format!("{}/ip/add", BASE_URL);
     let config: WireguardConfig = ureq::post(&url)
@@ -208,8 +228,45 @@ fn get_locations() -> Result<Locations, anyhow::Error> {
     Ok(locations)
 }
 
-fn generage_keys() -> Result<WireguardKeyPair, anyhow::Error> {
+fn get_keys(machine: Option<&PathBuf>) -> Result<WireguardKeyPair, anyhow::Error> {
+    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").join(machine_subdir);
+    debug!("key path = {:?}", &key_path);
+    std::fs::create_dir_all(&key_path)?;
+    let private_key_path = key_path.join("key");
+    let private_key = if private_key_path.is_file() {
+        std::fs::read_to_string(private_key_path)?
+    } else {
+        let key = generate_private_key()?;
+        std::fs::write(private_key_path, key.as_bytes())?;
+        key
+    };
+    let public_key_path = key_path.join("pubkey");
+    let public_key = if public_key_path.is_file() {
+        std::fs::read_to_string(public_key_path)?
+    } else {
+        let key = generate_public_key(&private_key)?;
+        std::fs::write(public_key_path, key.as_bytes())?;
+        key
+    };
+    Ok(WireguardKeyPair {
+        private_key,
+        public_key,
+    })
+}
+
+fn generate_private_key() -> anyhow::Result<String> {
     let privkey = process::Command::new("wg").arg("genkey").output()?.stdout;
+    Ok(String::from_utf8(privkey)?.trim_end().to_string())
+}
+
+fn generate_public_key(private_key: &str) -> anyhow::Result<String> {
     let mut pubkey_cmd = process::Command::new("wg")
         .arg("pubkey")
         .stdin(process::Stdio::piped())
@@ -219,10 +276,7 @@ fn generage_keys() -> Result<WireguardKeyPair, anyhow::Error> {
         .stdin
         .as_mut()
         .expect("no stdin")
-        .write_all(&privkey)?;
+        .write_all(private_key.as_bytes())?;
     let pubkey = pubkey_cmd.wait_with_output()?.stdout;
-    Ok(WireguardKeyPair {
-        private_key: String::from_utf8(privkey)?.trim_end().to_string(),
-        public_key: String::from_utf8(pubkey)?.trim_end().to_string(),
-    })
+    Ok(String::from_utf8(pubkey)?.trim_end().to_string())
 }