maye
/
gemini-keychecker
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/adapters/mod.rs
View File

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

20
src/adapters/output/local.rs
View File

@ -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<tokio::fs::File>,
key: &GeminiKey,
// Write valid key to appropriate tier file
pub async fn write_validated_key_to_tier_files(
free_file: &mut BufWriter<tokio::fs::File>,
paid_file: &mut BufWriter<tokio::fs::File>,
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(())
}

17
src/config/config.rs
View File

@ -18,9 +18,6 @@ struct Cli {
#[serde(skip_serializing_if = "Option::is_none")]
input_path: Option<PathBuf>,
#[arg(short = 'o', long)]
#[serde(skip_serializing_if = "Option::is_none")]
output_path: Option<PathBuf>,
#[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<KeyCheckerConfig> = 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,

26
src/types.rs
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
}
}

51
src/validation/key_validator.rs
View File

@ -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<GeminiKey, ValidatorError> {
let api_endpoint = api_endpoint.into_url()?;
) -> Result<ValidatedKey, ValidatorError> {
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<GeminiKey, ValidatorError> {
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
}
},
}
}
}

55
src/validation/validation_service.rs
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::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(())