use crate::error::{FluxError, Result};
use crate::keys::{KeyPair, PrivateKey, PublicKey};
use std::fs::{File, OpenOptions};
use std::io::{Read, Write};
use std::path::Path;
#[cfg(unix)]
use std::os::unix::fs::PermissionsExt;
#[derive(Debug)]
pub struct KeyStorage {
encrypt_private_keys: bool,
}
#[derive(Debug, Clone)]
pub struct StorageOptions {
pub file_mode: Option<u32>,
pub overwrite: bool,
pub password: Option<String>,
}
impl Default for StorageOptions {
fn default() -> Self {
Self {
file_mode: Some(0o600), overwrite: false,
password: None,
}
}
}
impl KeyStorage {
pub fn new() -> Self {
Self {
encrypt_private_keys: false,
}
}
pub fn with_encryption() -> Self {
Self {
encrypt_private_keys: true,
}
}
pub fn save_keypair(
&self,
keypair: &KeyPair,
public_key_path: &Path,
private_key_path: &Path,
options: &StorageOptions,
) -> Result<()> {
self.save_public_key(keypair.public_key(), public_key_path, options)?;
self.save_private_key(keypair.private_key(), private_key_path, options)?;
Ok(())
}
pub fn save_public_key(
&self,
public_key: &PublicKey,
path: &Path,
options: &StorageOptions,
) -> Result<()> {
if path.exists() && !options.overwrite {
return Err(FluxError::invalid_input(format!(
"File already exists: {}",
path.display()
)));
}
let pem_data = public_key.to_pem()?;
let mut file = self.create_file_with_permissions(path, options.file_mode)?;
file.write_all(pem_data.as_bytes())?;
log::info!("Public key saved to: {}", path.display());
Ok(())
}
pub fn save_private_key(
&self,
private_key: &PrivateKey,
path: &Path,
options: &StorageOptions,
) -> Result<()> {
if path.exists() && !options.overwrite {
return Err(FluxError::invalid_input(format!(
"File already exists: {}",
path.display()
)));
}
let pem_data = if self.encrypt_private_keys && options.password.is_some() {
return Err(FluxError::config(
"Encrypted private key storage not yet implemented",
));
} else {
private_key.to_pem()?
};
let mut file = self.create_file_with_permissions(path, Some(0o600))?;
file.write_all(pem_data.as_bytes())?;
log::info!("Private key saved to: {}", path.display());
Ok(())
}
pub fn load_public_key(&self, path: &Path) -> Result<PublicKey> {
let mut file = File::open(path).map_err(|e| {
FluxError::invalid_input(format!(
"Cannot open public key file {}: {}",
path.display(),
e
))
})?;
let mut contents = String::new();
file.read_to_string(&mut contents)?;
crate::keys::parsing::parse_public_key_from_str(&contents)
}
pub fn load_private_key(&self, path: &Path, password: Option<&str>) -> Result<PrivateKey> {
let mut file = File::open(path).map_err(|e| {
FluxError::invalid_input(format!(
"Cannot open private key file {}: {}",
path.display(),
e
))
})?;
let mut contents = String::new();
file.read_to_string(&mut contents)?;
if contents.contains("ENCRYPTED") && password.is_none() {
return Err(FluxError::invalid_input(
"Private key is encrypted but no password provided",
));
}
if password.is_some() {
return Err(FluxError::config(
"Encrypted private key loading not yet implemented",
));
}
crate::keys::parsing::parse_private_key_from_str(&contents)
}
pub fn load_keypair(
&self,
public_key_path: &Path,
private_key_path: &Path,
password: Option<&str>,
) -> Result<KeyPair> {
let public_key = self.load_public_key(public_key_path)?;
let private_key = self.load_private_key(private_key_path, password)?;
KeyPair::from_keys(public_key, private_key)
}
fn create_file_with_permissions(&self, path: &Path, mode: Option<u32>) -> Result<File> {
let file = OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(path)?;
#[cfg(unix)]
if let Some(mode) = mode {
let metadata = file.metadata()?;
let mut permissions = metadata.permissions();
permissions.set_mode(mode);
std::fs::set_permissions(path, permissions)?;
}
Ok(file)
}
}
impl Default for KeyStorage {
fn default() -> Self {
Self::new()
}
}
pub fn save_keypair_default(keypair: &KeyPair, base_name: &str) -> Result<()> {
let storage = KeyStorage::new();
let options = StorageOptions::default();
let public_name = format!("{}.pub", base_name);
let private_name = format!("{}.pem", base_name);
let public_path = Path::new(&public_name);
let private_path = Path::new(&private_name);
storage.save_keypair(keypair, public_path, private_path, &options)
}
pub fn load_keypair_default(base_name: &str, password: Option<&str>) -> Result<KeyPair> {
let storage = KeyStorage::new();
let public_name = format!("{}.pub", base_name);
let private_name = format!("{}.pem", base_name);
let public_path = Path::new(&public_name);
let private_path = Path::new(&private_name);
storage.load_keypair(public_path, private_path, password)
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
#[test]
fn test_storage_options_default() {
let options = StorageOptions::default();
assert_eq!(options.file_mode, Some(0o600));
assert!(!options.overwrite);
assert!(options.password.is_none());
}
#[test]
fn test_key_storage_creation() {
let storage = KeyStorage::new();
assert!(!storage.encrypt_private_keys);
let storage = KeyStorage::with_encryption();
assert!(storage.encrypt_private_keys);
}
#[test]
fn test_save_public_key() {
use crate::keys::KeyPair;
let keypair = KeyPair::generate(2048).unwrap();
let storage = KeyStorage::new();
let options = StorageOptions::default();
let temp_dir = tempdir().unwrap();
let key_path = temp_dir.path().join("test_key.pub");
let result = storage.save_public_key(keypair.public_key(), &key_path, &options);
assert!(result.is_ok());
assert!(key_path.exists());
}
#[test]
fn test_save_and_load_public_key() {
use crate::keys::KeyPair;
let keypair = KeyPair::generate(2048).unwrap();
let storage = KeyStorage::new();
let options = StorageOptions::default();
let temp_dir = tempdir().unwrap();
let key_path = temp_dir.path().join("test_key.pub");
storage
.save_public_key(keypair.public_key(), &key_path, &options)
.unwrap();
let loaded_key = storage.load_public_key(&key_path).unwrap();
assert_eq!(
loaded_key.key_size_bits(),
keypair.public_key().key_size_bits()
);
}
}