From 0056b4fa7bff96a3649629bd72a26f129f50bfbb Mon Sep 17 00:00:00 2001 From: Xerxes-2 Date: Thu, 17 Jul 2025 19:26:41 +1000 Subject: [PATCH] feat: update KeyCheckerConfig to use clap for argument parsing and improve client builder with proper timeout handling --- src/config/basic_client.rs | 4 +-- src/config/basic_config.rs | 64 ++++++++++++++++++++++++++++---------- src/key_validator.rs | 61 ++++++++++++++++++++++-------------- src/main.rs | 4 +-- src/validation.rs | 10 +++--- 5 files changed, 95 insertions(+), 48 deletions(-) diff --git a/src/config/basic_client.rs b/src/config/basic_client.rs index b3dcf29..46a15d4 100644 --- a/src/config/basic_client.rs +++ b/src/config/basic_client.rs @@ -5,9 +5,9 @@ use reqwest::Client; use crate::config::KeyCheckerConfig; pub fn client_builder(config: &KeyCheckerConfig) -> Result { - let mut builder = Client::builder().timeout(Duration::from_secs(config.timeout_sec)); + let mut builder = Client::builder().timeout(Duration::from_secs(config.timeout_sec())); - if let Some(ref proxy_url) = config.proxy { + if let Some(ref proxy_url) = config.proxy() { builder = builder.proxy(reqwest::Proxy::all(proxy_url.clone())?); } diff --git a/src/config/basic_config.rs b/src/config/basic_config.rs index dd0de0a..c4a1b54 100644 --- a/src/config/basic_config.rs +++ b/src/config/basic_config.rs @@ -1,7 +1,8 @@ use anyhow::{Ok, Result}; +use clap::Parser; use figment::{ Figment, - providers::{Env, Format, Toml}, + providers::{Env, Format, Serialized, Toml}, }; use serde::{Deserialize, Serialize}; use std::fs; @@ -9,46 +10,53 @@ use std::path::PathBuf; use std::sync::LazyLock; use url::Url; -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Parser)] pub struct KeyCheckerConfig { // Input file path containing API keys to check. #[serde(default)] - pub input_path: PathBuf, + #[arg(short, long)] + input_path: Option, // Output file path for valid API keys. #[serde(default)] - pub output_path: PathBuf, + #[arg(short, long)] + output_path: Option, // Backup file path for all API keys. #[serde(default)] - pub backup_path: PathBuf, + #[arg(short, long)] + backup_path: Option, // API host URL for key validation. - #[serde(default = "default_api_host")] - pub api_host: Url, + #[serde(default)] + #[arg(short, long)] + api_host: Option, // Request timeout in seconds. #[serde(default)] - pub timeout_sec: u64, + #[arg(short, long)] + timeout_sec: Option, // Maximum number of concurrent requests. #[serde(default)] - pub concurrency: usize, + #[arg(short, long)] + concurrency: Option, // Optional proxy URL for HTTP requests (e.g., --proxy http://user:pass@host:port). #[serde(default)] - pub proxy: Option, + #[arg(short, long)] + proxy: Option, } impl Default for KeyCheckerConfig { fn default() -> Self { Self { - input_path: default_input_path(), - output_path: default_output_path(), - backup_path: default_backup_path(), - api_host: default_api_host(), - timeout_sec: default_timeout(), - concurrency: default_concurrency(), + input_path: Some(default_input_path()), + output_path: Some(default_output_path()), + backup_path: Some(default_backup_path()), + api_host: Some(default_api_host()), + timeout_sec: Some(default_timeout()), + concurrency: Some(default_concurrency()), proxy: None, } } @@ -67,16 +75,40 @@ impl KeyCheckerConfig { // Load configuration from config.toml, environment variables, and defaults let config = Figment::new() + .merge(Serialized::defaults(Self::default())) .merge(Toml::file(CONFIG_PATH.as_path())) .merge(Env::prefixed("KEYCHECKER_")) + .merge(Serialized::defaults(Self::parse())) .extract()?; Ok(config) } + pub fn input_path(&self) -> PathBuf { + self.input_path.clone().unwrap_or_else(default_input_path) + } + pub fn output_path(&self) -> PathBuf { + self.output_path.clone().unwrap_or_else(default_output_path) + } + pub fn backup_path(&self) -> PathBuf { + self.backup_path.clone().unwrap_or_else(default_backup_path) + } + pub fn api_host(&self) -> Url { + self.api_host.clone().unwrap_or_else(default_api_host) + } + pub fn timeout_sec(&self) -> u64 { + self.timeout_sec.unwrap_or_else(default_timeout) + } + pub fn concurrency(&self) -> usize { + self.concurrency.unwrap_or_else(default_concurrency) + } + pub fn proxy(&self) -> Option { + self.proxy.clone() + } } fn default_input_path() -> PathBuf { "keys.txt".into() } + fn default_output_path() -> PathBuf { "output_keys.txt".into() } diff --git a/src/key_validator.rs b/src/key_validator.rs index b02443e..39220d2 100644 --- a/src/key_validator.rs +++ b/src/key_validator.rs @@ -5,32 +5,42 @@ use serde_json; use tokio::time::Duration; use url::Url; -use crate::types::{KeyStatus, ApiKey}; +use crate::types::{ApiKey, KeyStatus}; -pub async fn validate_key_with_retry(client: &Client, api_host: &Url, key: ApiKey) -> Option { +pub async fn validate_key_with_retry( + client: Client, + api_host: Url, + key: ApiKey, +) -> Option { let retry_policy = ExponentialBuilder::default() .with_max_times(3) .with_min_delay(Duration::from_secs(3)) .with_max_delay(Duration::from_secs(5)); - let result = (|| async { - match keytest(&client, &api_host, &key).await { - Ok(KeyStatus::Valid) => { - println!("Key: {}... -> SUCCESS", &key.as_str()[..10]); - Ok(Some(key.clone())) - } - Ok(KeyStatus::Invalid) => { - println!("Key: {}... -> INVALID (Forbidden)", &key.as_str()[..10]); - Ok(None) - } - Ok(KeyStatus::Retryable(reason)) => { - eprintln!("Key: {}... -> RETRYABLE (Reason: {})", &key.as_str()[..10], reason); - Err(anyhow::anyhow!("Retryable error: {}", reason)) - } - Err(e) => { - eprintln!("Key: {}... -> NETWORK ERROR (Reason: {})", &key.as_str()[..10], e); - Err(e) - } + let result = (async || match keytest(client.to_owned(), &api_host, &key).await { + Ok(KeyStatus::Valid) => { + println!("Key: {}... -> SUCCESS", &key.as_str()[..10]); + Ok(Some(key.clone())) + } + Ok(KeyStatus::Invalid) => { + println!("Key: {}... -> INVALID (Forbidden)", &key.as_str()[..10]); + Ok(None) + } + Ok(KeyStatus::Retryable(reason)) => { + eprintln!( + "Key: {}... -> RETRYABLE (Reason: {})", + &key.as_str()[..10], + reason + ); + Err(anyhow::anyhow!("Retryable error: {}", reason)) + } + Err(e) => { + eprintln!( + "Key: {}... -> NETWORK ERROR (Reason: {})", + &key.as_str()[..10], + e + ); + Err(e) } }) .retry(retry_policy) @@ -39,16 +49,19 @@ pub async fn validate_key_with_retry(client: &Client, api_host: &Url, key: ApiKe match result { Ok(key_result) => key_result, Err(_) => { - eprintln!("Key: {}... -> FAILED after all retries.", &key.as_str()[..10]); + eprintln!( + "Key: {}... -> FAILED after all retries.", + &key.as_str()[..10] + ); None } } } -async fn keytest(client: &Client, api_host: &Url, key: &ApiKey) -> Result { +async fn keytest(client: Client, api_host: &Url, key: &ApiKey) -> Result { const API_PATH: &str = "v1beta/models/gemini-2.0-flash-exp:generateContent"; let full_url = api_host.join(API_PATH)?; - + let request_body = serde_json::json!({ "contents": [ { @@ -77,4 +90,4 @@ async fn keytest(client: &Client, api_host: &Url, key: &ApiKey) -> Result KeyStatus::Retryable(format!("Received status {}, will retry.", other)), }; Ok(key_status) -} \ No newline at end of file +} diff --git a/src/main.rs b/src/main.rs index 25b2de9..a20d2d6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,13 +1,13 @@ use anyhow::Result; -use gemini_keychecker::config::{KeyCheckerConfig, client_builder}; use gemini_keychecker::adapters::load_keys; +use gemini_keychecker::config::{KeyCheckerConfig, client_builder}; use gemini_keychecker::validation::ValidationService; /// Main function - orchestrates the key validation process #[tokio::main] async fn main() -> Result<()> { let config = KeyCheckerConfig::load_config().unwrap(); - let keys = load_keys(&config.input_path)?; + let keys = load_keys(config.input_path().as_path())?; let client = client_builder(&config)?; let validation_service = ValidationService::new(config, client); diff --git a/src/validation.rs b/src/validation.rs index 65c4460..54ce76c 100644 --- a/src/validation.rs +++ b/src/validation.rs @@ -5,8 +5,8 @@ use reqwest::Client; use std::{fs, time::Instant}; use tokio::sync::mpsc; -use crate::config::KeyCheckerConfig; use crate::adapters::write_keys_txt_file; +use crate::config::KeyCheckerConfig; use crate::key_validator::validate_key_with_retry; use crate::types::ApiKey; @@ -42,13 +42,15 @@ impl ValidationService { // Create stream to validate keys concurrently let valid_keys_stream = stream - .map(|key| validate_key_with_retry(&self.client, &self.config.api_host, key)) - .buffer_unordered(self.config.concurrency) + .map(|key| { + validate_key_with_retry(self.client.to_owned(), self.config.api_host(), key) + }) + .buffer_unordered(self.config.concurrency()) .filter_map(|r| async { r }); pin_mut!(valid_keys_stream); // Open output file for writing valid keys - let mut output_file = fs::File::create(&self.config.output_path)?; + let mut output_file = fs::File::create(&self.config.output_path())?; // Process validated keys and write to output file while let Some(valid_key) = valid_keys_stream.next().await {