feat: implement tier-based key validation with separate output files

main
Yoo1tic 2025-08-03 02:04:04 +08:00
parent babbcb8e03
commit 218e8421c7
6 changed files with 114 additions and 57 deletions

View File

@ -2,4 +2,4 @@ pub mod input;
pub mod output; pub mod output;
pub use input::load_keys; pub use input::load_keys;
pub use output::write_keys_txt_file; pub use output::write_validated_key_to_tier_files;

View File

@ -1,16 +1,24 @@
use crate::types::GeminiKey; use crate::types::{GeminiKey, ValidatedKey, KeyTier};
use crate::error::Result; use crate::error::Result;
use std::{fs, io::Write}; use std::{fs, io::Write};
use tokio::io::{AsyncWriteExt, BufWriter}; use tokio::io::{AsyncWriteExt, BufWriter};
use toml::Value; use toml::Value;
use tracing::info; use tracing::info;
// Write valid key to output file // Write valid key to appropriate tier file
pub async fn write_keys_txt_file( pub async fn write_validated_key_to_tier_files(
file: &mut BufWriter<tokio::fs::File>, free_file: &mut BufWriter<tokio::fs::File>,
key: &GeminiKey, paid_file: &mut BufWriter<tokio::fs::File>,
validated_key: &ValidatedKey,
) -> Result<()> { ) -> 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(()) Ok(())
} }

View File

@ -18,9 +18,6 @@ struct Cli {
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
input_path: Option<PathBuf>, input_path: Option<PathBuf>,
#[arg(short = 'o', long)]
#[serde(skip_serializing_if = "Option::is_none")]
output_path: Option<PathBuf>,
#[arg(short = 'b', long)] #[arg(short = 'b', long)]
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
@ -53,9 +50,6 @@ pub struct KeyCheckerConfig {
#[serde(default)] #[serde(default)]
pub input_path: PathBuf, pub input_path: PathBuf,
// Output file path for valid API keys.
#[serde(default)]
pub output_path: PathBuf,
// Backup file path for all API keys. // Backup file path for all API keys.
#[serde(default)] #[serde(default)]
@ -123,6 +117,13 @@ impl KeyCheckerConfig {
.join("v1beta/models/gemini-2.5-flash-lite:generateContent") .join("v1beta/models/gemini-2.5-flash-lite:generateContent")
.expect("Failed to join API URL") .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 { impl Display for KeyCheckerConfig {
@ -140,14 +141,13 @@ impl Display for KeyCheckerConfig {
write!( write!(
f, f,
"Host={}, Proxy={}, Protocol={}, Timeout={}s, Concurrency={}, Input={}, Output={}, Backup={}", "Host={}, Proxy={}, Protocol={}, Timeout={}s, Concurrency={}, Input={}, Backup={}",
self.api_host, self.api_host,
proxy_status, proxy_status,
protocol_status, protocol_status,
self.timeout_sec, self.timeout_sec,
self.concurrency, self.concurrency,
self.input_path.display(), self.input_path.display(),
self.output_path.display(),
self.backup_path.display() self.backup_path.display()
) )
} }
@ -156,7 +156,6 @@ impl Display for KeyCheckerConfig {
// Single LazyLock for entire default configuration // Single LazyLock for entire default configuration
static DEFAULT_CONFIG: LazyLock<KeyCheckerConfig> = LazyLock::new(|| KeyCheckerConfig { static DEFAULT_CONFIG: LazyLock<KeyCheckerConfig> = LazyLock::new(|| KeyCheckerConfig {
input_path: "keys.txt".into(), input_path: "keys.txt".into(),
output_path: "output_keys.txt".into(),
backup_path: "backup_keys.txt".into(), backup_path: "backup_keys.txt".into(),
api_host: Url::parse("https://generativelanguage.googleapis.com/").unwrap(), api_host: Url::parse("https://generativelanguage.googleapis.com/").unwrap(),
timeout_sec: 15, timeout_sec: 15,

View File

@ -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
}
}

View File

@ -1,10 +1,10 @@
use reqwest::{Client, IntoUrl}; 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 super::{CACHE_CONTENT_TEST_BODY, GENERATE_CONTENT_TEST_BODY};
use crate::config::KeyCheckerConfig; use crate::config::KeyCheckerConfig;
use crate::error::ValidatorError; use crate::error::ValidatorError;
use crate::types::GeminiKey; use crate::types::{GeminiKey, ValidatedKey};
use crate::utils::send_request; use crate::utils::send_request;
pub async fn test_generate_content_api( pub async fn test_generate_content_api(
@ -12,8 +12,8 @@ pub async fn test_generate_content_api(
api_endpoint: impl IntoUrl, api_endpoint: impl IntoUrl,
api_key: GeminiKey, api_key: GeminiKey,
config: KeyCheckerConfig, config: KeyCheckerConfig,
) -> Result<GeminiKey, ValidatorError> { ) -> Result<ValidatedKey, ValidatorError> {
let api_endpoint = api_endpoint.into_url()?; let api_endpoint = api_endpoint.into_url().unwrap();
match send_request( match send_request(
client, client,
@ -25,8 +25,8 @@ pub async fn test_generate_content_api(
.await .await
{ {
Ok(_) => { Ok(_) => {
info!("SUCCESS - {}... - Valid key found", &api_key.as_ref()[..10]); info!("BASIC API VALID - {}... - Passed generate content API test", &api_key.as_ref()[..10]);
Ok(api_key) Ok(ValidatedKey::new(api_key))
} }
Err(e) => match &e { Err(e) => match &e {
ValidatorError::HttpBadRequest { .. } ValidatorError::HttpBadRequest { .. }
@ -51,15 +51,15 @@ pub async fn test_generate_content_api(
pub async fn test_cache_content_api( pub async fn test_cache_content_api(
client: Client, client: Client,
api_endpoint: impl IntoUrl, api_endpoint: impl IntoUrl,
api_key: GeminiKey, validated_key: ValidatedKey,
config: KeyCheckerConfig, config: KeyCheckerConfig,
) -> Result<GeminiKey, ValidatorError> { ) -> ValidatedKey {
let api_endpoint = api_endpoint.into_url()?; let api_endpoint = api_endpoint.into_url().unwrap();
match send_request( match send_request(
client, client,
&api_endpoint, &api_endpoint,
api_key.clone(), validated_key.key.clone(),
&CACHE_CONTENT_TEST_BODY, &CACHE_CONTENT_TEST_BODY,
config.max_retries, config.max_retries,
) )
@ -67,27 +67,28 @@ pub async fn test_cache_content_api(
{ {
Ok(_) => { Ok(_) => {
info!( info!(
"CACHE SUCCESS - {}... - Valid key for cache API", "PAID KEY DETECTED - {}... - Cache API accessible",
&api_key.as_ref()[..10] &validated_key.key.as_ref()[..10]
); );
Ok(api_key) validated_key.with_paid_tier()
} }
Err(e) => match &e { Err(e) => match &e {
ValidatorError::HttpBadRequest { .. } ValidatorError::HttpTooManyRequests { .. } => {
| ValidatorError::HttpUnauthorized { .. } debug!(
| ValidatorError::HttpForbidden { .. } "FREE KEY DETECTED - {}... - Cache API not accessible",
| ValidatorError::HttpTooManyRequests { .. } => { &validated_key.key.as_ref()[..10]
warn!(
"CACHE INVALID - {}... - {}",
&api_key.as_ref()[..10],
ValidatorError::KeyInvalid
); );
Err(ValidatorError::KeyInvalid) validated_key
} }
_ => { _ => {
error!("CACHE ERROR - {}... - {}", &api_key.as_ref()[..10], e); error!(
Err(e) "CACHE API ERROR - {}... - {}",
&validated_key.key.as_ref()[..10],
e
);
validated_key
} }
}, }
} }
} }

View File

@ -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::error::ValidatorError;
use crate::types::GeminiKey;
use crate::utils::client_builder;
use async_stream::stream; use async_stream::stream;
use futures::{pin_mut, stream::StreamExt}; use futures::{pin_mut, stream::StreamExt};
use reqwest::Client; use reqwest::Client;
use std::time::Instant; use std::time::Instant;
use tokio::{fs, io::AsyncWriteExt, sync::mpsc}; use tokio::{fs, io::AsyncWriteExt, sync::mpsc};
use tracing::error;
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;
pub struct ValidationService { pub struct ValidationService {
config: KeyCheckerConfig, 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 let valid_keys_stream = stream
.map(|key| { .map(|key| {
test_generate_content_api( test_generate_content_api(
@ -58,22 +59,44 @@ impl ValidationService {
) )
}) })
.buffer_unordered(self.config.concurrency) .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); pin_mut!(valid_keys_stream);
// Open output file for writing valid keys // Open output files for writing keys by tier (fixed filenames)
let output_file = fs::File::create(&self.config.output_path).await?; let free_keys_path = "freekey.txt";
let mut buffer_writer = tokio::io::BufWriter::new(output_file); 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 { while let Some(valid_key) = valid_keys_stream.next().await {
if let Err(e) = write_keys_txt_file(&mut buffer_writer, &valid_key).await { if let Err(e) = write_validated_key_to_tier_files(
eprintln!("Failed to write key to output file: {}", e); &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 // Flush both buffers to ensure all data is written to files
buffer_writer.flush().await?; free_buffer_writer.flush().await?;
paid_buffer_writer.flush().await?;
println!("Total Elapsed Time: {:?}", start_time.elapsed()); println!("Total Elapsed Time: {:?}", start_time.elapsed());
Ok(()) Ok(())