ssh_utils_lib/config/
app_vault.rs

1use anyhow::Result;
2use anyhow::Context;
3use serde::{Deserialize, Serialize};
4use sha2::Digest;
5use std::fs;
6use std::fs::Permissions;
7use std::os::unix::fs::PermissionsExt;
8use std::path::PathBuf;
9use hmac::{Hmac, Mac};
10use sha2::Sha256;
11use crate::config::crypto::*;
12use crate::helper::get_file_path;
13use crate::helper::ENCRYPTED_FILE;
14
15type HmacSha256 = Hmac<Sha256>;
16pub type EncryptionKey = Vec<u8>;
17
18#[derive(Serialize, Deserialize, Debug, PartialEq, Eq)]
19pub struct Server {
20    pub id: String,
21    pub password: String,
22}
23
24impl Server {
25    pub fn new(id: String, password: String) -> Self {
26        Self { id, password }
27    }
28}
29
30#[derive(Serialize, Deserialize, Debug, Default, PartialEq, Eq)]
31pub struct Vault {
32    pub servers: Vec<Server>,
33}
34
35impl Vault {
36    pub fn save(&self, encryption_key: &[u8; 32]) -> Result<()> {
37        let encrypt_data = encrypt_vault(self, encryption_key)?;
38        let file_path = get_file_path(ENCRYPTED_FILE)?;
39        
40        // Ensure the directory exists
41        let path = PathBuf::from(&file_path);
42        if let Some(parent) = path.parent() {
43            if !parent.exists() {
44                fs::create_dir_all(parent).context(format!(
45                    "Failed to create config directory at {:?}",
46                    parent
47                ))?;
48            }
49        }
50
51        // Write the encrypted data to the file
52        fs::write(&file_path, encrypt_data)
53            .context(format!("Failed to write encrypted data to file at {:?}", file_path))?;
54        
55        // Set file permissions to 0600
56        // Only owner have permissions to read and write 
57        let permissions = Permissions::from_mode(0o600);
58        fs::set_permissions(&file_path, permissions)
59            .context(format!("Failed to set permissions for file at {:?}", file_path))?;
60
61        Ok(())
62    }
63
64    pub fn modify_server(&mut self, id: &str, new_server: Server, encryption_key: &[u8; 32]) -> Result<()> {
65        if let Some(server) = self.servers.iter_mut().find(|server| server.id == id) {
66            server.password = new_server.password.clone();
67            self.save(encryption_key)?;
68        } else {
69            return Err(anyhow::anyhow!("Server with id {} not found", id));
70        }
71        Ok(())
72    }
73
74    pub fn add_server(&mut self, new_server: Server, encryption_key: &[u8; 32]) -> Result<()> {
75        self.servers.push(new_server);
76        self.save(encryption_key)?;
77        Ok(())
78    }
79
80    pub fn delete_server(&mut self, id: &str, encryption_key: &[u8; 32]) -> Result<()> {
81        if let Some(pos) = self.servers.iter().position(|server| server.id == id) {
82            self.servers.remove(pos);
83            self.save(encryption_key)?;
84        } else {
85            return Err(anyhow::anyhow!("Server with id {} not found", id));
86        }
87        Ok(())
88    }
89}
90
91/**
92    check if config file and its directory exists
93*/
94pub fn check_if_vault_bin_exists() -> Result<bool> {
95    let mut config_dir: PathBuf = if cfg!(debug_assertions) {
96        ".".into() // current running dir
97    } else {
98        dirs::home_dir().context("Unable to reach user's home directory.")?
99    };
100    config_dir.push(".config/ssh-utils");
101
102    // check if the encrypted_data file exists
103    let config_file_path = config_dir.join("encrypted_data.bin");
104    if !config_file_path.exists() {
105        return Ok(false);
106    }
107
108    Ok(true)
109}
110
111/**
112    encrypt vault
113*/
114pub fn encrypt_vault(vault: &Vault, encryption_key: &[u8; 32]) -> Result<Vec<u8>> {
115    // Serialize the Vault object to a string.
116    let unencrypt_data = toml::to_string(vault).context("Unable to serialize vault to string.")?;
117
118    // Step 3: Generate a 16-byte IV (initialization vector).
119    let iv = generate_iv();
120
121    // Step 4: Encrypt the serialized Vault data.
122    let data = unencrypt_data.as_bytes();
123    let encrypted_data = aes_encrypt(encryption_key, &iv, data)?;
124
125    // Step 5: Compute HMAC for the IV and encrypted data
126    let mut mac = HmacSha256::new_from_slice(encryption_key)
127        .context("Failed to create HMAC instance")?;
128    mac.update(&iv);
129    mac.update(&encrypted_data);
130    let hmac = mac.finalize().into_bytes();
131
132    // Concatenate the IV, encrypted data, and HMAC and return the result.
133    let mut result = Vec::with_capacity(iv.len() + encrypted_data.len() + hmac.len());
134    result.extend_from_slice(&iv);
135    result.extend_from_slice(&encrypted_data);
136    result.extend_from_slice(&hmac);
137
138    Ok(result)
139}
140
141/**
142    decrypt vault
143*/
144pub fn decrypt_vault(vault: &[u8], encryption_key: &[u8; 32]) -> Result<Vault> {
145    // Extract the IV, encrypted data, and HMAC.
146    let (iv, rest) = vault.split_at(16);
147    let (encrypted_data, hmac) = rest.split_at(rest.len() - 32);
148
149    // Verify HMAC
150    let mut mac = HmacSha256::new_from_slice(encryption_key)
151        .context("Failed to create HMAC instance")?;
152    mac.update(iv);
153    mac.update(encrypted_data);
154    mac.verify_slice(hmac).context("HMAC verification failed")?;
155
156    // Decrypt the data.
157    let decrypted_data = aes_decrypt(encryption_key, iv, encrypted_data)?;
158
159    // Convert the decrypted data to a string and parse it into a Vault object.
160    let decrypted_str =
161        String::from_utf8(decrypted_data).context("Failed to convert decrypted data to string")?;
162    
163    // If decrypted_str is blank, return a default Vault.
164    if decrypted_str.trim().is_empty() {
165        return Ok(Vault::default());
166    }
167
168    let vault: Vault =
169        toml::from_str(&decrypted_str).context("Failed to parse decrypted data as Vault")?;
170
171    Ok(vault)
172}
173
174fn derive_iv_from_id(id: &str) -> [u8; 16] {
175    // Step 1: Hash the id using SHA-256.
176    let mut hasher = Sha256::new();
177    hasher.update(id);
178    let result = hasher.finalize();
179
180    // Step 2: Take the first 16 bytes of the hash as the IV.
181    let mut iv = [0u8; 16];
182    iv.copy_from_slice(&result[..16]);
183    iv
184}
185
186/**
187    encrypt password to string
188*/
189pub fn encrypt_password(id: &str, password: &str, encryption_key: &[u8; 32]) -> Result<String> {
190    // Derive IV from id.
191    let iv = derive_iv_from_id(id);
192
193    // Encrypt the password using the provided aes_encrypt function.
194    let encrypted_data = aes_encrypt(encryption_key, &iv, password.as_bytes())?;
195
196    // Encode the result as a hex string.
197    let encrypted_hex = hex::encode(encrypted_data);
198
199    Ok(encrypted_hex)
200}
201
202/**
203    decrypt password to string
204*/
205pub fn decrypt_password(id: &str, encrypted_password: &str, encryption_key: &[u8; 32]) -> Result<String> {
206    // Derive IV from id.
207    let iv = derive_iv_from_id(id);
208
209    // Decode the encrypted password from hex string.
210    let encrypted_data = hex::decode(encrypted_password)
211        .context("Failed to decode hex string")?;
212
213    // Decrypt the password using the provided aes_decrypt function.
214    let decrypted_data = aes_decrypt(encryption_key, &iv, &encrypted_data)?;
215
216    // Convert the decrypted data to a string.
217    let decrypted_password = String::from_utf8(decrypted_data)
218        .context("Failed to convert decrypted data to string")?;
219
220    Ok(decrypted_password)
221}
222
223#[test]
224/**
225    test encrypt_password and decrypt_password func
226*/
227fn test_encryption_decryption_password() -> Result<()> {
228    let id = "550e8400-e29b-41d4-a716-446655440000";
229    let password = "my_secure_password";
230    let encryption_key = derive_key_from_password("123")?;
231
232    let encrypted_password = encrypt_password(id, password, &encryption_key)?;
233    println!("Encrypted password: {}", encrypted_password);
234
235    let decrypted_password = decrypt_password(id, &encrypted_password, &encryption_key)?;
236    println!("Decrypted password: {}", decrypted_password);
237
238    assert_eq!(password,decrypted_password.as_str());
239    Ok(())
240}
241
242/**
243    test encrypt_vault and decrypt_vault func
244*/
245#[test]
246fn test_encryption_decryption_vault() -> Result<()> {
247    let pass_data = r#"
248[[servers]]
249id = "server1"
250password = "secret_password1"
251
252[[servers]]
253id = "server2"
254password = "secret_password2"
255    "#;
256    let origin_vault: Vault = toml::from_str(pass_data)?;
257    let encryption_key = derive_key_from_password("123")?;
258    let encrypt_data = encrypt_vault(&origin_vault, &encryption_key)?;
259    let decrypt_vault = match decrypt_vault(&encrypt_data, &encryption_key) {
260        Err(e) => {
261            if let Some(_) = e.downcast_ref::<hmac::digest::MacError>() {
262                println!("wrong password");
263                return Err(e);
264            } else {
265                println!("An unexpected error occurred: {}", e);
266                return Err(e);
267            }
268        },
269        Ok(o) => o
270    };
271    assert_eq!(origin_vault, decrypt_vault);
272    Ok(())
273}