Initial commit

6 files changed, 642 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..d3611fa
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,3 @@
+/target
+/.idea
+Cargo.lock
\ No newline at end of file
diff --git a/Cargo.toml b/Cargo.toml
new file mode 100644
index 0000000..19c9061
--- /dev/null
+++ b/Cargo.toml
@@ -0,0 +1,17 @@
+[package]
+name = "kps"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+axum = "0.8.9"
+base64 = "0.22.1"
+chrono = { version = "0.4.45", features = ["clock"] }
+clap = { version = "4.6.1", features = ["derive"] }
+ed25519-dalek = "3.0.0-rc.1"
+rand = "0.10.2"
+serde = { version = "1.0.228", features = ["derive"] }
+serde_json = "1.0.150"
+sha2 = "0.11.0"
+tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread"] }
+ureq = "3.3.0"
diff --git a/src/client.rs b/src/client.rs
new file mode 100644
index 0000000..85a1221
--- /dev/null
+++ b/src/client.rs
@@ -0,0 +1,181 @@
+use base64::Engine;
+use base64::prelude::BASE64_STANDARD;
+use chrono::{DateTime, Utc};
+use ed25519_dalek::{Signer, SigningKey};
+use rand::TryRng;
+use rand::distr::{Alphanumeric, SampleString};
+use rand::rngs::SysRng;
+use std::fs::File;
+use std::io;
+use std::io::{BufReader, Read, Write};
+use ureq::Error::StatusCode;
+use ureq::SendBody;
+use crate::config::{read_config, ClientConfig};
+
+struct AuthExchange {
+    nonce: String,
+    timestamp: i64,
+    signature: String,
+}
+
+fn generate_nonce() -> String {
+    Alphanumeric.sample_string(&mut rand::rng(), 32)
+}
+
+fn create_auth(config: &ClientConfig) -> AuthExchange {
+    let mut buffer = [0u8; 44];
+    File::open(config.private_key.clone())
+        .expect("Failed to open private key file")
+        .read_exact(&mut buffer)
+        .expect("Failed to read private key file");
+    let decoded = BASE64_STANDARD
+        .decode(buffer)
+        .expect("Malformed private key");
+    let signing_key = SigningKey::from_bytes(&decoded.try_into().expect("Malformed private key"));
+
+    let nonce = generate_nonce();
+    let timestamp = Utc::now().timestamp();
+    let signature = BASE64_STANDARD.encode(
+        signing_key
+            .sign(format!("{}{}", nonce, timestamp).as_bytes())
+            .to_bytes(),
+    );
+
+    AuthExchange {
+        nonce,
+        timestamp,
+        signature,
+    }
+}
+
+pub fn gen_pubkey() {
+    let mut buffer = [0u8; 44];
+    io::stdin().read_exact(&mut buffer).unwrap();
+
+    let decoded = BASE64_STANDARD
+        .decode(buffer)
+        .expect("Malformed private key");
+
+    let signing_key = SigningKey::from_bytes(&decoded.try_into().expect("Malformed private key"));
+    let verifying_key = signing_key.verifying_key();
+    let encoded = BASE64_STANDARD.encode(verifying_key);
+    println!("{}", encoded);
+}
+
+pub fn gen_privkey() {
+    let mut seed = [0u8; 32];
+    SysRng
+        .try_fill_bytes(&mut seed)
+        .expect("Failed to generate random seed");
+    let signing_key = SigningKey::from_bytes(&seed);
+    let encoded = BASE64_STANDARD.encode(signing_key.to_bytes());
+    println!("{}", encoded);
+}
+
+pub fn upload_vault() {
+    let config: ClientConfig = read_config().expect("Failed to read config");
+    let server_url = &config.server_url;
+    let file = File::open(&config.db_path).expect("Failed to open database file");
+    let reader = BufReader::new(file);
+
+    let auth = create_auth(&config);
+
+    let code = ureq::post(format!("{}/vaults", server_url))
+        .content_type("application/x-keepass2")
+        .header("X-Signature", auth.signature)
+        .header("X-Nonce", auth.nonce)
+        .header("X-Timestamp", auth.timestamp)
+        .send(SendBody::from_owned_reader(reader))
+        .expect("Failed to send file")
+        .status();
+    if code == 200 {
+        println!("Successfully uploaded vault file");
+    } else {
+        println!("Failed to upload vault file: {}", code);
+    }
+}
+
+pub fn ls_vaults(long: bool) {
+    let config: ClientConfig = read_config().expect("Failed to read config");
+    let server_url = &config.server_url;
+
+    let auth = create_auth(&config);
+
+    let result = ureq::get(format!("{}/vaults", server_url))
+        .header("X-Signature", auth.signature)
+        .header("X-Nonce", auth.nonce)
+        .header("X-Timestamp", auth.timestamp)
+        .call()
+        .expect("Failed to get from server")
+        .body_mut()
+        .read_to_string()
+        .expect("Failed to read server response");
+
+    let files = result.split("\n");
+    for file in files {
+        if file.is_empty() {
+            break;
+        }
+        let mut parts = file.split("_");
+
+        let vault_id = parts.next().unwrap();
+
+        let timestamp: i64 = parts
+            .next()
+            .expect("Malformed response")
+            .parse()
+            .expect("Malformed response");
+        let dt = DateTime::from_timestamp(timestamp, 0).expect("Failed to parse timestamp");
+
+        if long {
+            println!(
+                "{} \t [Created: {}]",
+                vault_id,
+                dt.format("%m-%d-%Y %H:%M:%S")
+            );
+        } else {
+            println!(
+                "{} \t [Created: {}]",
+                &vault_id[0..8],
+                dt.format("%m-%d-%Y %H:%M:%S")
+            );
+        }
+    }
+}
+
+pub fn download_vault(id: String, install: bool, file: String) {
+    let config: ClientConfig = read_config().expect("Failed to read config");
+    let server_url = &config.server_url;
+
+    let auth = create_auth(&config);
+    let response = ureq::get(format!("{}/vaults/{}", server_url, id))
+        .header("X-Signature", auth.signature)
+        .header("X-Nonce", auth.nonce)
+        .header("X-Timestamp", auth.timestamp)
+        .call();
+    match response {
+        Ok(mut response) => {
+            let result = response
+                .body_mut()
+                .read_to_vec()
+                .expect("Failed to read server response");
+
+            let mut file = File::create(if install { config.db_path } else { file })
+                .expect("Failed to open database file");
+
+            file.write_all(&result)
+                .expect("Failed to write to database file");
+        }
+        Err(error) => match error {
+            StatusCode(stat) => {
+                if stat == 300 {
+                    println!("There are multiple candidates available for download.");
+                    println!("** Rerun the command with a long ID.")
+                } else {
+                    println!("An internal server error occurred!");
+                }
+            }
+            _ => eprintln!("An error occurred: {}", error),
+        },
+    }
+}
diff --git a/src/config.rs b/src/config.rs
new file mode 100644
index 0000000..85dbbe2
--- /dev/null
+++ b/src/config.rs
@@ -0,0 +1,32 @@
+use std::fs;
+use serde::{Deserialize, Serialize};
+use serde::de::DeserializeOwned;
+
+const CONFIG_LOCATIONS: [&str; 3] = [
+    "config.json",
+    "/etc/kps/config.json",
+    "~/.config/kps/config.json",
+];
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct ClientConfig {
+    pub private_key: String,
+    pub db_path: String,
+    pub server_url: String,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct ServerConfig {
+    pub public_key: String,
+    pub vaults_directory: String
+}
+
+pub fn read_config<T: DeserializeOwned>() -> Option<T> {
+    for file in CONFIG_LOCATIONS {
+        let data = fs::read_to_string(file);
+        if let Ok(data) = data {
+            return Some(serde_json::from_str(data.as_str()).expect("Malformed config file"));
+        }
+    }
+    None
+}
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
new file mode 100644
index 0000000..a26cfc8
--- /dev/null
+++ b/src/main.rs
@@ -0,0 +1,60 @@
+pub mod client;
+pub mod config;
+pub mod server;
+
+use crate::client::{download_vault, gen_privkey, gen_pubkey, ls_vaults, upload_vault};
+use crate::server::server_main;
+use clap::{Parser, Subcommand};
+
+#[derive(Parser)]
+#[command(
+    name = "kps",
+    version,
+    about = "A simple personal KeePassXC vault sync manager"
+)]
+struct CliArgs {
+    #[command(subcommand)]
+    command: Commands,
+}
+#[derive(Subcommand)]
+enum Commands {
+    #[command(
+        alias = "pubkey",
+        about = "Generates an ed25519 public key from the given private key"
+    )]
+    PublicKey,
+    #[command(alias = "privkey", about = "Generates an ed25519 private key")]
+    PrivateKey,
+    #[command(about = "Starts the kps server")]
+    Server,
+    #[command(about = "Uploads your vault to the kps server")]
+    Upload,
+    #[command(alias = "ls", about = "Lists the available vaults")]
+    Vaults {
+        #[arg(long, default_value = "false")]
+        long: bool,
+    },
+    #[command(about = "Downloads the specified vault file")]
+    Download {
+        #[arg(default_value = "latest")]
+        id: String,
+        #[arg(long, default_value = "false")]
+        install: bool,
+        #[arg(long, default_value = "vault.kdbx")]
+        file: String,
+    },
+}
+
+#[tokio::main]
+async fn main() {
+    let args = CliArgs::parse();
+
+    match args.command {
+        Commands::PublicKey => gen_pubkey(),
+        Commands::PrivateKey => gen_privkey(),
+        Commands::Server => server_main().await,
+        Commands::Upload => upload_vault(),
+        Commands::Vaults { long } => ls_vaults(long),
+        Commands::Download { id, install, file } => download_vault(id, install, file),
+    }
+}
diff --git a/src/server.rs b/src/server.rs
new file mode 100644
index 0000000..ced51d1
--- /dev/null
+++ b/src/server.rs
@@ -0,0 +1,349 @@
+use std::fs;
+use axum::body::{Body, HttpBody};
+use axum::http::{HeaderMap, Request, Response, StatusCode, header};
+use axum::middleware::Next;
+use axum::routing::{get, post};
+use axum::{Router, middleware};
+use base64::Engine;
+use base64::prelude::BASE64_STANDARD;
+use chrono::Utc;
+use ed25519_dalek::{Signature, VerifyingKey};
+use sha2::{Digest, Sha256};
+use std::fs::{File, ReadDir, rename};
+use std::future::poll_fn;
+use std::io::{BufWriter, Read, Write};
+use std::path::Path;
+use std::pin::Pin;
+use std::sync::OnceLock;
+use tokio::net::TcpListener;
+use crate::config::{read_config, ServerConfig};
+
+static PUBLIC_KEY: OnceLock<VerifyingKey> = OnceLock::new();
+static SERVER_CONFIG: OnceLock<ServerConfig> = OnceLock::new();
+
+enum ChallengeStatus {
+    BadTimestamp,
+    BadNonce,
+    BadSignature,
+    Success,
+}
+
+fn parse_header<T: std::str::FromStr>(headers: &HeaderMap, key: &str) -> Result<T, StatusCode> {
+    headers
+        .get(key)
+        .and_then(|value| value.to_str().ok())
+        .and_then(|string| string.parse().ok())
+        .ok_or(StatusCode::BAD_REQUEST)
+}
+
+fn log_challenge(signature: &String, status: ChallengeStatus) {
+    let short_signature = signature.chars().take(10).collect::<String>();
+    match status {
+        ChallengeStatus::BadTimestamp => println!(
+            "[*] Challenge {}... was invalid due to bad timestamp",
+            short_signature
+        ),
+        ChallengeStatus::BadNonce => println!(
+            "[*] Challenge {}... was invalid due to malformed nonce",
+            short_signature
+        ),
+        ChallengeStatus::BadSignature => println!(
+            "[*] Challenge {}... was invalid due to an invalid signature",
+            short_signature
+        ),
+        ChallengeStatus::Success => println!("[*] Challenge {}... succeeded", short_signature),
+    }
+}
+
+fn create_symlink(original: &Path, link: &Path, _is_dir: bool) -> std::io::Result<()> {
+    #[cfg(unix)]
+    {
+        let full_target = std::path::absolute(original)?;
+        std::os::unix::fs::symlink(full_target, link)
+    }
+
+    #[cfg(windows)]
+    {
+        if _is_dir {
+            std::os::windows::fs::symlink_dir(original, link)
+        } else {
+            std::os::windows::fs::symlink_file(original, link)
+        }
+    }
+}
+
+fn perform_download(file: String) -> Result<Response<Body>, StatusCode> {
+    let file_bytes = std::fs::read(file).map_err(|_| StatusCode::NOT_FOUND)?;
+    let body = Body::from(file_bytes);
+    let response = Response::builder()
+        .status(StatusCode::OK)
+        .header(header::CONTENT_TYPE, "application/x-keepass2")
+        .body(body)
+        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+    Ok(response)
+}
+
+fn find_short_id_candidates(path: &String, vaults: ReadDir) -> Vec<String> {
+    let candidates: Vec<_> = vaults
+        .filter_map(|entry| entry.ok())
+        .map(|entry| entry.path())
+        .filter(|entry| entry.is_file())
+        .map(|entry| entry.to_string_lossy().into_owned())
+        .filter(|entry| entry[7..].starts_with(path.as_str()))
+        .collect();
+    candidates
+}
+
+async fn auth_middleware(req: Request<Body>, next: Next) -> Result<Response<Body>, StatusCode> {
+    let headers = req.headers();
+    let signature: String = parse_header(headers, "X-Signature")?;
+    let nonce: String = parse_header(headers, "X-Nonce")?;
+    let timestamp: i64 = parse_header(headers, "X-Timestamp")?;
+    if !(0..=60).contains(&(Utc::now().timestamp() - timestamp)) {
+        log_challenge(&signature, ChallengeStatus::BadTimestamp);
+        return Err(StatusCode::UNAUTHORIZED);
+    }
+    if nonce.len() != 32 {
+        log_challenge(&signature, ChallengeStatus::BadNonce);
+        println!("[*] A challenge was failed for a malformed nonce");
+        return Err(StatusCode::UNAUTHORIZED);
+    }
+    let message = format!("{}{}", nonce, timestamp);
+    let verifier_signature = Signature::from_bytes(
+        &BASE64_STANDARD
+            .decode(signature.as_bytes())
+            .map_err(|_| StatusCode::UNAUTHORIZED)?
+            .try_into()
+            .map_err(|_| StatusCode::UNAUTHORIZED)?,
+    );
+
+    match PUBLIC_KEY
+        .get()
+        .unwrap()
+        .verify_strict(message.as_bytes(), &verifier_signature)
+    {
+        Ok(_) => {
+            log_challenge(&signature, ChallengeStatus::Success);
+            Ok(next.run(req).await)
+        }
+        Err(_) => {
+            log_challenge(&signature, ChallengeStatus::BadSignature);
+            Err(StatusCode::UNAUTHORIZED)
+        }
+    }
+}
+
+async fn upload(mut body: Body) -> Result<StatusCode, StatusCode> {
+    let temp_filename = "vaults/temporary.kdbx";
+    let file = File::create(temp_filename).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+    let mut writer = BufWriter::new(file);
+
+    let mut hasher = Sha256::new();
+
+    loop {
+        let frame_result = poll_fn(|cx| Pin::new(&mut body).poll_frame(cx)).await;
+
+        match frame_result {
+            Some(Ok(chunk)) => {
+                if let Some(chunk) = chunk.data_ref() {
+                    Digest::update(&mut hasher, chunk);
+                    writer
+                        .write_all(chunk)
+                        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+                }
+            }
+            Some(Err(_)) => return Err(StatusCode::INTERNAL_SERVER_ERROR),
+            None => break,
+        }
+    }
+
+    writer
+        .flush()
+        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+
+    let file_hash: String = hasher
+        .finalize()
+        .iter()
+        .map(|b| format!("{:02x}", b))
+        .collect();
+
+    println!(
+        "[*] Uploading file with hash {}...",
+        file_hash.chars().take(8).collect::<String>()
+    );
+
+    let vaults = std::fs::read_dir("vaults").map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+    for vault in vaults {
+        let entry = vault.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+        let path = entry.path();
+
+        if path.is_file()
+            && let Some(file_name) = path.file_name()
+        {
+            if file_name.to_string_lossy().starts_with(file_hash.as_str()) {
+                println!("[*] File already exists, doing nothing...");
+                std::fs::remove_file(temp_filename)
+                    .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+                return Ok(StatusCode::NOT_MODIFIED);
+            }
+        }
+    }
+
+    let final_filename = format!("vaults/{}_{}.kdbx", file_hash, Utc::now().timestamp());
+
+    rename(temp_filename, &final_filename).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+
+    // A temporary symlink is created, because creating a symlink is not atomic whereas
+    // renaming a file is atomic in Unix-like systems. This pattern is to prevent data races
+    // from concurrent calls to upload or concurrent calls to get the latest vault.
+    let tmp_symlink_name = format!("vaults/latest-{}.kdbx", file_hash);
+    let target = Path::new(&final_filename);
+    let tmp_symlink = Path::new(&tmp_symlink_name);
+
+    create_symlink(target, tmp_symlink, false).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+    rename(tmp_symlink, "vaults/latest.kdbx").map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+
+    println!("[*] File uploaded and latest symlink reassigned");
+
+    Ok(StatusCode::OK)
+}
+
+async fn download_vault(
+    axum::extract::Path(path): axum::extract::Path<String>,
+) -> Result<Response<Body>, StatusCode> {
+    if path == "latest" {
+        println!("[*] Downloading latest vault");
+        return perform_download(String::from("vaults/latest.kdbx"));
+    }
+
+    if path.len() != 8 && path.len() != 64 {
+        return Err(StatusCode::BAD_REQUEST);
+    }
+
+    println!("[*] Received download request for path: {}", path);
+
+    let vaults = std::fs::read_dir("vaults").map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+
+    if path.len() == 8 {
+        let candidates = find_short_id_candidates(&path, vaults);
+
+        if candidates.is_empty() {
+            println!("[*] No candidates found for the given short ID");
+            return Err(StatusCode::NOT_FOUND);
+        }
+        if candidates.len() != 1 {
+            println!("[*] Multiple candidates found for the given short ID");
+            return Err(StatusCode::MULTIPLE_CHOICES);
+        }
+        println!("[*] File found found for the given short ID");
+        perform_download(candidates[0].clone())
+    } else {
+        for vault in vaults {
+            let entry = match vault {
+                Ok(v) => v,
+                Err(_) => continue,
+            };
+            let file_name = entry.path();
+
+            if file_name.to_string_lossy().contains(path.as_str()) {
+                println!("[*] File found for the given long ID");
+                return perform_download(file_name.to_string_lossy().into_owned());
+            }
+        }
+        Err(StatusCode::NOT_FOUND)
+    }
+}
+
+async fn list_vaults() -> Result<String, StatusCode> {
+    let mut s = String::new();
+    let path = "vaults";
+    let mut counter = 0;
+    for entry in std::fs::read_dir(path).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)? {
+        let entry = entry.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?.path();
+        if entry.is_file()
+            && !entry.is_symlink()
+            && let Some(path) = entry.as_path().file_stem()
+            && let Some(path) = path.to_str()
+        {
+            s += path;
+            s += "\n";
+            counter += 1;
+        }
+    }
+
+    println!("[*] Listed {} vaults", counter);
+    Ok(s)
+}
+
+async fn delete_vault(
+    axum::extract::Path(path): axum::extract::Path<String>,
+) -> Result<StatusCode, StatusCode> {
+    if path.len() != 8 && path.len() != 64 {
+        return Err(StatusCode::BAD_REQUEST);
+    }
+    println!("[*] Received deletion request for path: {}", path);
+
+    let vaults = std::fs::read_dir("vaults").map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+
+    if path.len() == 8 {
+        let candidates = find_short_id_candidates(&path, vaults);
+
+        if candidates.is_empty() {
+            println!("[*] No candidates found for the given short ID");
+            return Err(StatusCode::NOT_FOUND);
+        }
+        if candidates.len() != 1 {
+            println!("[*] Multiple candidates found for the given short ID");
+            return Err(StatusCode::MULTIPLE_CHOICES);
+        }
+        println!("[*] File found found for the given short ID");
+        std::fs::remove_file(candidates[0].clone())
+            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+        Ok(StatusCode::OK)
+    } else {
+        for vault in vaults {
+            let entry = match vault {
+                Ok(v) => v,
+                Err(_) => continue,
+            };
+            let file_name = entry.path();
+
+            if file_name.to_string_lossy().contains(path.as_str()) {
+                println!("[*] File found for the given long ID");
+                std::fs::remove_file(file_name).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+                return Ok(StatusCode::OK);
+            }
+        }
+        Err(StatusCode::NOT_FOUND)
+    }
+}
+
+pub async fn server_main() {
+    SERVER_CONFIG.set(read_config().unwrap()).unwrap();
+
+    let mut buffer: [u8; 44] = [0u8; 44];
+    let pubkey_file = SERVER_CONFIG.get().unwrap().public_key.clone();
+    let file = File::open(pubkey_file);
+    if let Ok(mut file) = file {
+        file.read_exact(&mut buffer).expect("Failed to read file");
+    } else {
+        eprintln!("Failed to read public key!");
+        return;
+    }
+
+    let decoded = BASE64_STANDARD
+        .decode(buffer)
+        .expect("Malformed private key");
+    let verifying_key =
+        VerifyingKey::from_bytes(&decoded.try_into().expect("Malformed public key"))
+            .expect("Malformed public key");
+
+    PUBLIC_KEY.set(verifying_key).unwrap();
+
+    let app = Router::new()
+        .route("/vaults", post(upload).get(list_vaults))
+        .route("/vaults/{*path}", get(download_vault).delete(delete_vault))
+        .layer(middleware::from_fn(auth_middleware));
+    let listener = TcpListener::bind("0.0.0.0:3000").await.unwrap();
+    axum::serve(listener, app).await.unwrap();
+}