diff --git a/crates/web3_account_server/Cargo.toml b/crates/web3_account_server/Cargo.toml index 9b8d5926..0e9a13dc 100644 --- a/crates/web3_account_server/Cargo.toml +++ b/crates/web3_account_server/Cargo.toml @@ -5,7 +5,7 @@ edition = "2021" name = "web3_account_server" license = "MIT" repository = "https://github.com/scx1332/" -version = "0.4.11" +version = "0.4.12" [dependencies] actix-cors = { workspace = true } diff --git a/crates/web3_account_server/src/main.rs b/crates/web3_account_server/src/main.rs index 2f0d37e9..234c117a 100644 --- a/crates/web3_account_server/src/main.rs +++ b/crates/web3_account_server/src/main.rs @@ -5,7 +5,8 @@ use actix_web_httpauth::middleware::HttpAuthentication; use std::env; use std::fs::OpenOptions; use std::io::{BufReader, BufWriter}; -use std::sync::{Arc}; +use std::path::{Component, Path, PathBuf}; +use std::sync::Arc; use structopt::StructOpt; fn read_results(file_name: &str) -> Vec { @@ -119,6 +120,70 @@ async fn get_from_queue(data: web::Data) -> impl Responder { } } +fn get_file_name_from_filename_and_group(file_name: &str, group: &str) -> String { + let sanitized_group: String = group + .chars() + .filter(|c| c.is_alphanumeric() || *c == '_') // Allow only alphanumeric and underscores + .collect(); + + let sanitized_file_name = Path::new(file_name) + .components() + .filter(|comp| matches!(comp, Component::Normal(_))) // Allow only valid file names + .collect::(); + + format!("{}_{}", sanitized_group, sanitized_file_name.display()) +} + +async fn add_to_queue_group( + data: web::Data, + path: web::Path, + item: String, +) -> impl Responder { + let _lock = data.lock.lock().await; + let group = path.into_inner(); + let file_name = get_file_name_from_filename_and_group(&data.file_name, &group); + let Ok(private_key) = hex::decode(item.replace("0x", "")) else { + return HttpResponse::BadRequest().body("Invalid item type"); + }; + if private_key.len() != 32 { + return HttpResponse::BadRequest().body("Invalid item length"); + } + match add(hex::encode(private_key), &file_name) { + Ok(true) => HttpResponse::Ok().body("Added to the queue"), + Ok(false) => HttpResponse::Ok().body("Item already in the queue"), + Err(e) => { + log::error!("Error adding item: {}", e); + HttpResponse::InternalServerError().finish() + } + } +} + +async fn count_group(data: web::Data, path: web::Path) -> impl Responder { + let _lock = data.lock.lock().await; + let group = path.into_inner(); + let file_name = get_file_name_from_filename_and_group(&data.file_name, &group); + let results = read_results(&file_name); + HttpResponse::Ok().body(results.len().to_string()) +} + +async fn get_from_queue_group( + data: web::Data, + path: web::Path, +) -> impl Responder { + let _lock = data.lock.lock().await; + + let group = path.into_inner(); + let file_name = get_file_name_from_filename_and_group(&data.file_name, &group); + match get(&file_name) { + Ok(Some(item)) => HttpResponse::Ok().body(item), + Ok(None) => HttpResponse::BadRequest().body("Queue is empty"), + Err(e) => { + log::error!("Error getting item: {}", e); + HttpResponse::InternalServerError().finish() + } + } +} + fn get_env_access_token() -> String { env::var("BEARER_KEY").unwrap_or("change_me".to_string()) } @@ -167,6 +232,9 @@ async fn main() -> std::io::Result<()> { .route("/count", web::get().to(count)) .route("/add", web::post().to(add_to_queue)) .route("/get", web::get().to(get_from_queue)) + .route("/count/{group}", web::get().to(count_group)) + .route("/add/{group}", web::post().to(add_to_queue_group)) + .route("/get/{group}", web::get().to(get_from_queue_group)) }) .bind(format!("{}:{}", args.http_addr, args.http_port))? .workers(1)