From 25e89d8610e5640723c909c8d201f304b13fc3f1 Mon Sep 17 00:00:00 2001 From: Yoo1tic <137816438+Yoo1tic@users.noreply.github.com> Date: Mon, 7 Jul 2025 00:39:12 +0800 Subject: [PATCH] feat: enhance API key validation tool with proxy support and improved timeout settings --- src/main.rs | 72 ++++++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 58 insertions(+), 14 deletions(-) diff --git a/src/main.rs b/src/main.rs index adefc77..ad9b076 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,33 +16,52 @@ use std::{ }; use tokio::time::Duration; use url::Url; + +// Regex pattern for validating Google API keys (AIzaSy followed by 33 characters) static API_KEY_REGEX: LazyLock = LazyLock::new(|| Regex::new(r"^AIzaSy.{33}$").unwrap()); +/// Configuration structure for the key checker tool #[derive(Parser, Debug)] #[command(version, about = "A tool to check and backup API keys", long_about = None)] struct KeyCheckerConfig { + /// Input file path containing API keys to check #[arg(long, short = 'i', default_value = "keys.txt")] input_path: PathBuf, + /// Output file path for valid API keys #[arg(long, short = 'o', default_value = "output_keys.txt")] output_path: PathBuf, + /// API host URL for key validation #[arg(long, short = 'u', default_value = "https://generativelanguage.googleapis.com/")] api_host: Url, - #[arg(long, short = 't', default_value_t = 5000)] - timeout_ms: u64, + /// Request timeout in seconds + #[arg(long, short = 't', default_value_t = 60)] + timeout_sec: u64, + /// Maximum number of concurrent requests #[arg(long, short = 'c', default_value_t = 30)] concurrency: usize, + + /// Optional proxy URL for HTTP requests + #[arg(long, short = 'x')] + proxy: Option, } +/// Status of API key validation #[derive(Debug)] enum KeyStatus { + /// Key is valid and working Valid, + /// Key is invalid or unauthorized Invalid, + /// Temporary error, key validation should be retried Retryable(String), } +/// Load and validate API keys from a file +/// Returns a vector of unique, valid API keys fn load_keys(path: &Path) -> Result> { let keys_txt = fs::read_to_string(path)?; + // Use HashSet to automatically deduplicate keys let unique_keys_set: HashSet<&str> = keys_txt .lines() .map(|line| line.trim()) @@ -53,11 +72,14 @@ fn load_keys(path: &Path) -> Result> { Ok(keys) } +/// Validate an API key with exponential backoff retry logic +/// Returns Some(key) if valid, None if invalid or failed after all retries async fn validate_key_with_retry(client: &Client, api_host: &Url, key: String) -> Option { + // Configure exponential backoff retry policy let retry_policy = ExponentialBuilder::default() .with_max_times(3) - .with_min_delay(Duration::from_secs(3)) - .with_max_delay(Duration::from_secs(8)); + .with_min_delay(Duration::from_secs(5)) + .with_max_delay(Duration::from_secs(10)); let result = (|| async { match keytest(&client, &api_host, &key).await { @@ -91,9 +113,13 @@ async fn validate_key_with_retry(client: &Client, api_host: &Url, key: String) - } } +/// Test a single API key by making a request to the Gemini API +/// Returns the validation status based on the HTTP response async fn keytest(client: &Client, api_host: &Url, keys: &str) -> Result { - const API_PATH: &str = "v1beta/models/gemini-2.0-flash:generateContent"; + const API_PATH: &str = "v1beta/models/gemini-2.0-flash-exp:generateContent"; let full_url = api_host.join(API_PATH)?; + + // Simple test request body let request_body = serde_json::json!({ "contents": [ { @@ -117,27 +143,41 @@ async fn keytest(client: &Client, api_host: &Url, keys: &str) -> Result KeyStatus::Valid, - // 403 & 401 + // 403 & 401 - Key is invalid or unauthorized StatusCode::FORBIDDEN | StatusCode::UNAUTHORIZED => KeyStatus::Invalid, - // Other Status Code + // Other status codes - Temporary error, retry other => KeyStatus::Retryable(format!("Received status {}, will retry.", other)), }; Ok(key_status) } +/// Build HTTP client with optional proxy configuration +/// Returns a configured reqwest Client +fn build_client(config: &KeyCheckerConfig) -> Result { + let mut client_builder = Client::builder() + .timeout(Duration::from_secs(config.timeout_sec)); + + // Add proxy configuration if specified + if let Some(proxy_url) = &config.proxy { + client_builder = client_builder.proxy(reqwest::Proxy::all(proxy_url.clone())?); + } + + client_builder.build().map_err(Into::into) +} + +/// Main function - orchestrates the key validation process #[tokio::main] async fn main() -> Result<()> { let start_time = Instant::now(); let config = KeyCheckerConfig::parse(); let keys = load_keys(&config.input_path)?; - let client = Client::builder() - .timeout(Duration::from_millis(config.timeout_ms)) - .build()?; + let client = build_client(&config)?; + // Create channel for streaming keys from producer to consumer let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::(); let stream = stream! { while let Some(item) = rx.recv().await { @@ -145,6 +185,7 @@ async fn main() -> Result<()> { } }; + // Spawn producer thread to send keys through channel spawn(move || { for key in keys { if API_KEY_REGEX.is_match(&key) { @@ -157,21 +198,24 @@ async fn main() -> Result<()> { } }); + // Create stream to validate keys concurrently let valid_keys_stream = stream .map(|key| validate_key_with_retry(&client, &config.api_host, key)) .buffer_unordered(config.concurrency) .filter_map(|r| async { r }); pin_mut!(valid_keys_stream); - // open output file + + // Open output file for writing valid keys let mut output_file = fs::File::create(&config.output_path)?; - // Write valid keys to output file + + // Process validated keys and write to output file while let Some(valid_key) = valid_keys_stream.next().await { - // Collect valid keys println!("Valid key found: {}", valid_key); if let Err(e) = writeln!(output_file, "{}", valid_key) { eprintln!("Failed to write key to output file: {}", e); } } + println!("Total Elapsed Time: {:?}", start_time.elapsed()); Ok(()) } \ No newline at end of file