diff --git a/src/adapters/mod.rs b/src/adapters/mod.rs index b642a9c..7d10128 100644 --- a/src/adapters/mod.rs +++ b/src/adapters/mod.rs @@ -2,4 +2,4 @@ pub mod input; pub mod output; pub use input::load_keys; -pub use output::write_keys_txt_file; \ No newline at end of file +pub use output::write_validated_key_to_tier_files; \ No newline at end of file diff --git a/src/adapters/output/local.rs b/src/adapters/output/local.rs index 8a26c48..85fc849 100644 --- a/src/adapters/output/local.rs +++ b/src/adapters/output/local.rs @@ -1,16 +1,24 @@ -use crate::types::GeminiKey; +use crate::types::{GeminiKey, ValidatedKey, KeyTier}; use crate::error::Result; use std::{fs, io::Write}; use tokio::io::{AsyncWriteExt, BufWriter}; use toml::Value; use tracing::info; -// Write valid key to output file -pub async fn write_keys_txt_file( - file: &mut BufWriter, - key: &GeminiKey, +// Write valid key to appropriate tier file +pub async fn write_validated_key_to_tier_files( + free_file: &mut BufWriter, + paid_file: &mut BufWriter, + validated_key: &ValidatedKey, ) -> Result<()> { - file.write_all(format!("{}\n", key.as_ref()).as_bytes()).await?; + match validated_key.tier { + KeyTier::Free => { + free_file.write_all(format!("{}\n", validated_key.key.as_ref()).as_bytes()).await?; + } + KeyTier::Paid => { + paid_file.write_all(format!("{}\n", validated_key.key.as_ref()).as_bytes()).await?; + } + } Ok(()) } diff --git a/src/config/config.rs b/src/config/config.rs index 7672238..455deeb 100644 --- a/src/config/config.rs +++ b/src/config/config.rs @@ -18,9 +18,6 @@ struct Cli { #[serde(skip_serializing_if = "Option::is_none")] input_path: Option, - #[arg(short = 'o', long)] - #[serde(skip_serializing_if = "Option::is_none")] - output_path: Option, #[arg(short = 'b', long)] #[serde(skip_serializing_if = "Option::is_none")] @@ -53,9 +50,6 @@ pub struct KeyCheckerConfig { #[serde(default)] pub input_path: PathBuf, - // Output file path for valid API keys. - #[serde(default)] - pub output_path: PathBuf, // Backup file path for all API keys. #[serde(default)] @@ -123,6 +117,13 @@ impl KeyCheckerConfig { .join("v1beta/models/gemini-2.5-flash-lite:generateContent") .expect("Failed to join API URL") } + + /// Returns the complete Gemini API URL for cachedContents endpoint + pub fn cache_api_url(&self) -> Url { + self.api_host + .join("v1beta/cachedContents") + .expect("Failed to join cache API URL") + } } impl Display for KeyCheckerConfig { @@ -140,14 +141,13 @@ impl Display for KeyCheckerConfig { write!( f, - "Host={}, Proxy={}, Protocol={}, Timeout={}s, Concurrency={}, Input={}, Output={}, Backup={}", + "Host={}, Proxy={}, Protocol={}, Timeout={}s, Concurrency={}, Input={}, Backup={}", self.api_host, proxy_status, protocol_status, self.timeout_sec, self.concurrency, self.input_path.display(), - self.output_path.display(), self.backup_path.display() ) } @@ -156,7 +156,6 @@ impl Display for KeyCheckerConfig { // Single LazyLock for entire default configuration static DEFAULT_CONFIG: LazyLock = LazyLock::new(|| KeyCheckerConfig { input_path: "keys.txt".into(), - output_path: "output_keys.txt".into(), backup_path: "backup_keys.txt".into(), api_host: Url::parse("https://generativelanguage.googleapis.com/").unwrap(), timeout_sec: 15, diff --git a/src/types.rs b/src/types.rs index e81cf4b..017729a 100644 --- a/src/types.rs +++ b/src/types.rs @@ -31,3 +31,29 @@ impl FromStr for GeminiKey { } } } + +#[derive(Debug, Clone, PartialEq)] +pub enum KeyTier { + Free, + Paid, +} + +#[derive(Debug, Clone)] +pub struct ValidatedKey { + pub key: GeminiKey, + pub tier: KeyTier, +} + +impl ValidatedKey { + pub fn new(key: GeminiKey) -> Self { + Self { + key, + tier: KeyTier::Free, + } + } + + pub fn with_paid_tier(mut self) -> Self { + self.tier = KeyTier::Paid; + self + } +} diff --git a/src/validation/key_validator.rs b/src/validation/key_validator.rs index 3616ffe..fac9795 100644 --- a/src/validation/key_validator.rs +++ b/src/validation/key_validator.rs @@ -1,10 +1,10 @@ use reqwest::{Client, IntoUrl}; -use tracing::{error, info, warn}; +use tracing::{debug, error, info, warn}; use super::{CACHE_CONTENT_TEST_BODY, GENERATE_CONTENT_TEST_BODY}; use crate::config::KeyCheckerConfig; use crate::error::ValidatorError; -use crate::types::GeminiKey; +use crate::types::{GeminiKey, ValidatedKey}; use crate::utils::send_request; pub async fn test_generate_content_api( @@ -12,8 +12,8 @@ pub async fn test_generate_content_api( api_endpoint: impl IntoUrl, api_key: GeminiKey, config: KeyCheckerConfig, -) -> Result { - let api_endpoint = api_endpoint.into_url()?; +) -> Result { + let api_endpoint = api_endpoint.into_url().unwrap(); match send_request( client, @@ -25,8 +25,8 @@ pub async fn test_generate_content_api( .await { Ok(_) => { - info!("SUCCESS - {}... - Valid key found", &api_key.as_ref()[..10]); - Ok(api_key) + info!("BASIC API VALID - {}... - Passed generate content API test", &api_key.as_ref()[..10]); + Ok(ValidatedKey::new(api_key)) } Err(e) => match &e { ValidatorError::HttpBadRequest { .. } @@ -51,15 +51,15 @@ pub async fn test_generate_content_api( pub async fn test_cache_content_api( client: Client, api_endpoint: impl IntoUrl, - api_key: GeminiKey, + validated_key: ValidatedKey, config: KeyCheckerConfig, -) -> Result { - let api_endpoint = api_endpoint.into_url()?; +) -> ValidatedKey { + let api_endpoint = api_endpoint.into_url().unwrap(); match send_request( client, &api_endpoint, - api_key.clone(), + validated_key.key.clone(), &CACHE_CONTENT_TEST_BODY, config.max_retries, ) @@ -67,27 +67,28 @@ pub async fn test_cache_content_api( { Ok(_) => { info!( - "CACHE SUCCESS - {}... - Valid key for cache API", - &api_key.as_ref()[..10] + "PAID KEY DETECTED - {}... - Cache API accessible", + &validated_key.key.as_ref()[..10] ); - Ok(api_key) + validated_key.with_paid_tier() } Err(e) => match &e { - ValidatorError::HttpBadRequest { .. } - | ValidatorError::HttpUnauthorized { .. } - | ValidatorError::HttpForbidden { .. } - | ValidatorError::HttpTooManyRequests { .. } => { - warn!( - "CACHE INVALID - {}... - {}", - &api_key.as_ref()[..10], - ValidatorError::KeyInvalid + ValidatorError::HttpTooManyRequests { .. } => { + debug!( + "FREE KEY DETECTED - {}... - Cache API not accessible", + &validated_key.key.as_ref()[..10] ); - Err(ValidatorError::KeyInvalid) + validated_key } _ => { - error!("CACHE ERROR - {}... - {}", &api_key.as_ref()[..10], e); - Err(e) + error!( + "CACHE API ERROR - {}... - {}", + &validated_key.key.as_ref()[..10], + e + ); + validated_key } - }, + } } } + diff --git a/src/validation/validation_service.rs b/src/validation/validation_service.rs index ec27f42..2a9fca6 100644 --- a/src/validation/validation_service.rs +++ b/src/validation/validation_service.rs @@ -1,15 +1,15 @@ +use super::key_validator::{test_cache_content_api, test_generate_content_api}; +use crate::adapters::{load_keys, write_validated_key_to_tier_files}; +use crate::config::KeyCheckerConfig; use crate::error::ValidatorError; +use crate::types::GeminiKey; +use crate::utils::client_builder; use async_stream::stream; use futures::{pin_mut, stream::StreamExt}; use reqwest::Client; use std::time::Instant; use tokio::{fs, io::AsyncWriteExt, sync::mpsc}; - -use super::key_validator::test_generate_content_api; -use crate::adapters::{load_keys, write_keys_txt_file}; -use crate::config::KeyCheckerConfig; -use crate::types::GeminiKey; -use crate::utils::client_builder; +use tracing::error; pub struct ValidationService { config: KeyCheckerConfig, @@ -47,7 +47,8 @@ impl ValidationService { } }); - // Create stream to validate keys concurrently + // Create stream to validate keys concurrently (two-stage pipeline) + let cache_api_url = self.config.cache_api_url(); let valid_keys_stream = stream .map(|key| { test_generate_content_api( @@ -58,22 +59,44 @@ impl ValidationService { ) }) .buffer_unordered(self.config.concurrency) - .filter_map(|result| async { result.ok() }); + .filter_map(|result| async { result.ok() }) + .map(|validated_key| { + test_cache_content_api( + self.client.clone(), + cache_api_url.clone(), + validated_key, + self.config.clone(), + ) + }) + .buffer_unordered(self.config.concurrency); pin_mut!(valid_keys_stream); - // Open output file for writing valid keys - let output_file = fs::File::create(&self.config.output_path).await?; - let mut buffer_writer = tokio::io::BufWriter::new(output_file); + // Open output files for writing keys by tier (fixed filenames) + let free_keys_path = "freekey.txt"; + let paid_keys_path = "paidkey.txt"; - // Process validated keys and write to output file + let free_file = fs::File::create(&free_keys_path).await?; + let paid_file = fs::File::create(&paid_keys_path).await?; + + let mut free_buffer_writer = tokio::io::BufWriter::new(free_file); + let mut paid_buffer_writer = tokio::io::BufWriter::new(paid_file); + + // Process validated keys and write to appropriate tier files while let Some(valid_key) = valid_keys_stream.next().await { - if let Err(e) = write_keys_txt_file(&mut buffer_writer, &valid_key).await { - eprintln!("Failed to write key to output file: {}", e); + if let Err(e) = write_validated_key_to_tier_files( + &mut free_buffer_writer, + &mut paid_buffer_writer, + &valid_key, + ) + .await + { + error!("Failed to write key to output file: {e}"); } } - // Flush the buffer to ensure all data is written to the file - buffer_writer.flush().await?; + // Flush both buffers to ensure all data is written to files + free_buffer_writer.flush().await?; + paid_buffer_writer.flush().await?; println!("Total Elapsed Time: {:?}", start_time.elapsed()); Ok(())