diff --git a/nyash_server/proto/nyash.proto b/nyash_server/proto/nyash.proto index 81202c2..f7af140 100644 --- a/nyash_server/proto/nyash.proto +++ b/nyash_server/proto/nyash.proto @@ -1,5 +1,5 @@ syntax = "proto3"; -package nyash_server; +package nyash_proto; service NyashLuks { // Requesting work, gettin in response work data, error, or message that there is no work currently awaylable. @@ -74,5 +74,6 @@ message ProgressRequest { message ProgressReply { // Brutforce progress as ratio from 0 to 1(completed) + bool key_found = 1; double progress = 2; } \ No newline at end of file diff --git a/nyash_server/src/server.rs b/nyash_server/src/server.rs index 1775d6c..1568880 100644 --- a/nyash_server/src/server.rs +++ b/nyash_server/src/server.rs @@ -1,17 +1,32 @@ +use std::ops::DerefMut; +use std::str::FromStr; +use std::sync::atomic::AtomicBool; +use std::{path::PathBuf}; + use redb; +use clap::Parser; -use tonic::{transport::Server, Request, Response, Status}; +use std::sync::{Arc, Mutex}; +use tokio::sync::RwLock as AsyncRwLock; +use tokio::time::{interval, Duration}; -use nyash_server::nyash_luks_server::{NyashLuks, NyashLuksServer}; -use nyash_server::{WorkRequest, WorkData, WorkReply, work_reply, WorkCommit, CommitReply, ProgressRequest, ProgressReply}; +use tonic::{Request, Response, Status, transport::Server}; -pub mod nyash_server { - tonic::include_proto!("nyash_server"); // The string specified here must match the proto package name + +pub mod nyash_proto { + tonic::include_proto!("nyash_proto"); // The string specified here must match the proto package name } +use nyash_proto::nyash_luks_server::{NyashLuks, NyashLuksServer}; +use nyash_proto::{ + CommitReply, KeyData, ProgressReply, ProgressRequest, WorkCommit, WorkData, WorkReply, WorkRequest, work_commit, + work_reply, +}; + mod database; mod num_utils; +mod config; fn work_rec_to_work_data(work_rec: database::JobRecord) -> WorkData { let tw_key = num_utils::u128_to_u64arr(work_rec.tweak_key); @@ -28,64 +43,187 @@ fn work_rec_to_work_data(work_rec: database::JobRecord) -> WorkData { } } -#[derive(Debug, Default)] +// Shared state +// #[derive(Clone, Debug)] +// pub struct ServiceState { +// pub db: Arc, +// pub key_found: bool, +// pub progress: f64, +// } + +// Background task to update state +// async fn update_service_state(state: Arc>) { +// loop { +// tokio::time::sleep(tokio::time::Duration::from_secs(5)).await; +// let mut guard = state.write().await; +// match database::db_get_progress(&guard.db) { +// Ok(p) => { +// guard.progress = p; +// }, +// Err(_) => println!("Error getting progress from DB") +// } +// } +// } + +// Background task to update progress from DB +async fn update_service_progress(progress: Arc>, db: Arc) { + loop { + tokio::time::sleep(tokio::time::Duration::from_secs(5)).await; + match database::db_get_progress(&db) { + Ok(p) => { + let mut guard = progress.write().await; + *guard = p; + }, + Err(_) => println!("Error getting progress from DB") + } + } +} + +#[derive(Debug)] pub struct NyashService { - db: redb::Database, - progress: f64, + //state: Arc>, + db: Arc, + key_found: Arc>, + progress: Arc>, } #[tonic::async_trait] impl NyashLuks for NyashService { - async fn request_work(&self, request: Request) -> Result, Status> { + async fn request_work(&self, request: Request) -> Result, Status> { println!("Got a request: {:?}", request); let pref_work_size: u64 = request.into_inner().pref_work_size; + match database::db_create_job(&self.db, pref_work_size) { Ok(Some(work_rec)) => { let reply = WorkReply { result: Some(work_reply::Result::WorkData(work_rec_to_work_data(work_rec))), }; Ok(Response::new(reply)) // Send back our formatted greeting - }, + } Ok(None) => { let reply = WorkReply { result: Some(work_reply::Result::NoWork(true)), }; Ok(Response::new(reply)) - }, + } Err(_ex) => { let reply = WorkReply { result: Some(work_reply::Result::Error("DB Error!".to_string())), }; Ok(Response::new(reply)) } - } + } } async fn commit_work(&self, request: Request) -> Result, Status> { println!("Got a request: {:?}", request); - let work_id: u64 = request.into_inner().work_id; - match database::db_commit_job(&self.db, work_id) { - Ok(true) => Ok(Response::new(CommitReply {status_code: 0, progress: self.progress})), - Ok(false) => Ok(Response::new(CommitReply {status_code: 1, progress: self.progress})), - Err(_) => Err(Status::internal("DB error")) + let work_commit = request.into_inner(); + + let put_k_no_err: bool = match work_commit.result { + Some(work_commit::Result::FoundKey(key_data)) => { + { + let mut guard = self.key_found.write().await; + *guard = true; + } + let tw_k = num_utils::u64arr_to_u128((key_data.tweak_key0, key_data.tweak_key1)); + let en_k = num_utils::u64arr_to_u128((key_data.start_key0, key_data.start_key1)); + match database::db_put_found_key(&self.db, tw_k, en_k) { + Ok(_) => true, + Err(_) => false, + } + } + _ => true, + }; + + let comm_no_err: bool = match database::db_commit_job(&self.db, work_commit.work_id) { + Ok(true) => true, + Ok(false) => false, + Err(_) => false, + }; + + if (put_k_no_err && comm_no_err) == true { + Ok(Response::new(CommitReply { + status_code: 0, + progress: *self.progress.read().await + })) + } else { + Err(Status::internal("DB error")) } } async fn request_progress(&self, _request: Request) -> Result, Status> { - Ok(Response::new(ProgressReply {progress: self.progress})) + Ok(Response::new(ProgressReply { + key_found: *self.key_found.read().await, + progress: *self.progress.read().await, + })) } } +/// Server for AES-XTS distributed bruteforce +#[derive(clap::Parser, Debug)] +#[command(version, about, long_about = None)] +struct ProgArgs { + /// Config file name + #[arg(short, long, default_value = "/etc/nyash_aes_d.conf")] + config: PathBuf, +} + + + #[tokio::main] async fn main() -> Result<(), Box> { - let addr = "[::1]:50051".parse()?; - let greeter = MyGreeter::default(); + + let args = ProgArgs::parse(); + log::error!("Arguments parsed: {:?}", args.config); + println!("Path: {:?}", args.config); + let server_config: config::ServerConfig = config::read_config(&args.config).unwrap(); + + let db_dir = server_config.db_dir.unwrap(); + let db_file = format!("{}data.rdb",db_dir); + let db_path = PathBuf::from_str(db_file.as_str()).unwrap(); + + let db = database::db_open(&db_path).expect("Error Opening database"); + + // Get the address to bind to + let addr = format!("{}:{}",server_config.bind_addr, server_config.listen_port); + let addr: std::net::SocketAddr = addr.parse().expect("Invalid address"); + + println!("Listening on: {}", addr); + + + let found_keys = database::db_get_found_keys(&db).expect("Error getting found keys from DB"); + if found_keys.len() >0 { + println!("Key already found!"); + for key in found_keys { + println!("Key found: {:?}", key); + } + println!("Exit!"); + return Ok(()); + } + + + let progress = database::db_get_progress(&db).expect("Error getting progress from DB"); + + let shared_progress = Arc::new(AsyncRwLock::new(progress)); + let shared_db = Arc::new(db); + let shared_key_found = Arc::new(AsyncRwLock::new(false)); + + // let service_state = Arc::new(AsyncRwLock::new(ServiceState { + // db: Arc::new(db), + // key_found: false, + // progress: progress, + // })); + + // Spawn the periodic updater + tokio::spawn(update_service_progress(shared_progress.clone(), shared_db.clone())); + + let nyash_service = NyashService{db:shared_db, progress:shared_progress, key_found:shared_key_found}; + Server::builder() - .add_service(GreeterServer::new(greeter)) + .add_service(NyashLuksServer::new(nyash_service)) .serve(addr) .await?; - Ok(()) -} \ No newline at end of file +}