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

refactor: reorganize project structure by moving service modules to validation and utils

main
Yoo1tic 2025-08-02 21:42:08 +08:00
parent 5f1b30772b
commit 772280da20
13 changed files with 244 additions and 109 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
src/config/config.rs
View File

@ -5,7 +5,6 @@ use figment::{
providers::{Env, Format, Serialized, Toml},
};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::fmt::{self, Display};
use std::fs;
use std::path::PathBuf;
@ -168,26 +167,6 @@ static DEFAULT_CONFIG: LazyLock<KeyCheckerConfig> = LazyLock::new(|| KeyCheckerC
log_level: "info".to_string(),
});
// LazyLock for the test message body used in API key validation
pub static TEST_MESSAGE_BODY: LazyLock<Value> = LazyLock::new(|| {
serde_json::json!({
"contents": [
{
"parts": [
{
"text": "Hi"
}
]
}
],
"generationConfig": {
"thinkingConfig": {
"thinkingBudget": 0
}
}
})
});
fn default_api_host() -> Url {
DEFAULT_CONFIG.api_host.clone()
}

2
src/config/mod.rs
View File

@ -1,3 +1,3 @@
mod config;
pub use config::{KeyCheckerConfig, TEST_MESSAGE_BODY};
pub use config::KeyCheckerConfig;

16
src/error.rs
View File

@ -3,13 +3,13 @@ use thiserror::Error;
#[derive(Error, Debug)]
pub enum ValidatorError {
#[error("HTTP error: {0}")]
ReqwestError(#[from] reqwest::Error),
ReqwestError(Box<reqwest::Error>),
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("Configuration error: {0}")]
ConfigError(#[from] figment::Error),
ConfigError(Box<figment::Error>),
#[error("TOML serialization error: {0}")]
TomlSer(#[from] toml::ser::Error),
@ -30,4 +30,16 @@ pub enum ValidatorError {
KeyFormatInvalid(String),
}
impl From<reqwest::Error> for ValidatorError {
fn from(err: reqwest::Error) -> Self {
ValidatorError::ReqwestError(Box::new(err))
}
}
impl From<figment::Error> for ValidatorError {
fn from(err: figment::Error) -> Self {
ValidatorError::ConfigError(Box::new(err))
}
}
pub type Result<T> = std::result::Result<T, ValidatorError>;

3
src/lib.rs
View File

@ -2,7 +2,8 @@ pub mod adapters;
pub mod config;
pub mod error;
pub mod types;
pub mod service;
pub mod validation;
pub mod utils;
// ASCII art for Gemini
pub const BANNER: &str = r#"

2
src/main.rs
View File

@ -1,5 +1,5 @@
use gemini_keychecker::error::ValidatorError;
use gemini_keychecker::{BANNER, config::KeyCheckerConfig, service::start_validation};
use gemini_keychecker::{BANNER, config::KeyCheckerConfig, validation::start_validation};
use mimalloc::MiMalloc;
use tracing::info;

74
src/service/key_tester.rs
View File

@ -1,74 +0,0 @@
use backon::{ExponentialBuilder, Retryable};
use reqwest::{Client, IntoUrl, StatusCode};
use tokio::time::Duration;
use tracing::{debug, error, info, warn};
use url::Url;
use crate::config::{KeyCheckerConfig, TEST_MESSAGE_BODY};
use crate::error::ValidatorError;
use crate::types::GeminiKey;
pub async fn validate_key(
client: Client,
api_endpoint: impl IntoUrl,
api_key: GeminiKey,
config: KeyCheckerConfig,
) -> Result<GeminiKey, ValidatorError> {
let api_endpoint = api_endpoint.into_url()?;
match send_test_request(client, &api_endpoint, api_key.clone(), config.max_retries).await {
Ok(_) => {
info!("SUCCESS - {}... - Valid key found", &api_key.as_ref()[..10]);
Ok(api_key)
}
Err(e) => match e.status() {
Some(
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN | StatusCode::TOO_MANY_REQUESTS,
) => {
warn!(
"INVALID - {}... - {}",
&api_key.as_ref()[..10],
ValidatorError::KeyInvalid
);
Err(ValidatorError::KeyInvalid)
}
_ => {
let req_error = ValidatorError::ReqwestError(e);
error!("ERROR- {}... - {}", &api_key.as_ref()[..10], req_error);
Err(req_error)
}
},
}
}
async fn send_test_request(
client: Client,
api_endpoint: &Url,
key: GeminiKey,
max_retries: usize,
) -> Result<reqwest::Response, reqwest::Error> {
let retry_policy = ExponentialBuilder::default()
.with_max_times(max_retries)
.with_min_delay(Duration::from_secs(1))
.with_max_delay(Duration::from_secs(2));
(async || {
let response = client
.post(api_endpoint.clone())
.header("Content-Type", "application/json")
.header("X-goog-api-key", key.as_ref())
.json(&*TEST_MESSAGE_BODY)
.send()
.await?;
debug!("Response for key {}: {:?}", key.as_ref(), response.status());
response.error_for_status()
})
.retry(&retry_policy)
.when(|e: &reqwest::Error| {
!matches!(
e.status(),
Some(StatusCode::FORBIDDEN | StatusCode::UNAUTHORIZED)
)
})
.await
}

7
src/service/mod.rs
View File

@ -1,7 +0,0 @@
pub mod http_client;
pub mod key_tester;
pub mod validation;
pub use http_client::client_builder;
pub use key_tester::validate_key;
pub use validation::{ValidationService, start_validation};

0
src/service/http_client.rs → src/utils/http_client.rs
View File

5
src/utils/mod.rs Normal file
View File

@ -0,0 +1,5 @@
pub mod http_client;
pub mod request;
pub use http_client::client_builder;
pub use request::send_request;

41
src/utils/request.rs Normal file
View File

@ -0,0 +1,41 @@
use backon::{ExponentialBuilder, Retryable};
use reqwest::{Client, StatusCode};
use serde_json::Value;
use tokio::time::Duration;
use tracing::debug;
use url::Url;
use crate::types::GeminiKey;
pub async fn send_request(
client: Client,
api_endpoint: &Url,
key: GeminiKey,
payload: &Value,
max_retries: usize,
) -> Result<reqwest::Response, reqwest::Error> {
let retry_policy = ExponentialBuilder::default()
.with_max_times(max_retries)
.with_min_delay(Duration::from_secs(1))
.with_max_delay(Duration::from_secs(2));
(async || {
let response = client
.post(api_endpoint.clone())
.header("Content-Type", "application/json")
.header("X-goog-api-key", key.as_ref())
.json(payload)
.send()
.await?;
debug!("Response for key {}: {:?}", key.as_ref(), response.status());
response.error_for_status()
})
.retry(&retry_policy)
.when(|e: &reqwest::Error| {
!matches!(
e.status(),
Some(StatusCode::FORBIDDEN | StatusCode::UNAUTHORIZED)
)
})
.await
}

97
src/validation/key_validator.rs Normal file
View File

@ -0,0 +1,97 @@
use reqwest::{Client, IntoUrl, StatusCode};
use tracing::{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::utils::send_request;
pub async fn test_generate_content_api(
client: Client,
api_endpoint: impl IntoUrl,
api_key: GeminiKey,
config: KeyCheckerConfig,
) -> Result<GeminiKey, ValidatorError> {
let api_endpoint = api_endpoint.into_url()?;
match send_request(
client,
&api_endpoint,
api_key.clone(),
&GENERATE_CONTENT_TEST_BODY,
config.max_retries,
)
.await
{
Ok(_) => {
info!("SUCCESS - {}... - Valid key found", &api_key.as_ref()[..10]);
Ok(api_key)
}
Err(e) => match e.status() {
Some(
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN | StatusCode::TOO_MANY_REQUESTS,
) => {
warn!(
"INVALID - {}... - {}",
&api_key.as_ref()[..10],
ValidatorError::KeyInvalid
);
Err(ValidatorError::KeyInvalid)
}
_ => {
let req_error = ValidatorError::from(e);
error!("ERROR- {}... - {}", &api_key.as_ref()[..10], req_error);
Err(req_error)
}
},
}
}
pub async fn test_cache_content_api(
client: Client,
api_endpoint: impl IntoUrl,
api_key: GeminiKey,
config: KeyCheckerConfig,
) -> Result<GeminiKey, ValidatorError> {
let api_endpoint = api_endpoint.into_url()?;
match send_request(
client,
&api_endpoint,
api_key.clone(),
&CACHE_CONTENT_TEST_BODY,
config.max_retries,
)
.await
{
Ok(_) => {
info!(
"CACHE SUCCESS - {}... - Valid key for cache API",
&api_key.as_ref()[..10]
);
Ok(api_key)
}
Err(e) => match e.status() {
Some(
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN | StatusCode::TOO_MANY_REQUESTS,
) => {
warn!(
"CACHE INVALID - {}... - {}",
&api_key.as_ref()[..10],
ValidatorError::KeyInvalid
);
Err(ValidatorError::KeyInvalid)
}
_ => {
let req_error = ValidatorError::from(e);
error!(
"CACHE ERROR - {}... - {}",
&api_key.as_ref()[..10],
req_error
);
Err(req_error)
}
},
}
}

80
src/validation/mod.rs Normal file
View File

@ -0,0 +1,80 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::sync::LazyLock;
pub mod key_validator;
pub mod validation_service;
pub use key_validator::{test_cache_content_api, test_generate_content_api};
pub use validation_service::{ValidationService, start_validation};
#[derive(Serialize, Deserialize, Debug)]
pub struct TextPart {
pub text: String,
}
#[derive(Serialize, Deserialize, Debug)]
pub struct ContentPart {
pub parts: Vec<TextPart>,
#[serde(skip_serializing_if = "Option::is_none")]
pub role: Option<String>,
}
#[derive(Serialize, Deserialize, Debug)]
pub struct ThinkingConfig {
#[serde(rename = "thinkingBudget")]
pub thinking_budget: u32,
}
#[derive(Serialize, Deserialize, Debug)]
pub struct GenerationConfig {
#[serde(skip_serializing_if = "Option::is_none")]
#[serde(rename = "thinkingConfig")]
pub thinking_config: Option<ThinkingConfig>,
}
#[derive(Serialize, Deserialize, Debug)]
pub struct GeminiRequest {
#[serde(skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
pub contents: Vec<ContentPart>,
#[serde(skip_serializing_if = "Option::is_none")]
#[serde(rename = "generationConfig")]
pub generation_config: Option<GenerationConfig>,
#[serde(skip_serializing_if = "Option::is_none")]
pub ttl: Option<String>,
}
// LazyLock for the test message body used in API key validation
pub static GENERATE_CONTENT_TEST_BODY: LazyLock<Value> = LazyLock::new(|| {
let generate_request = GeminiRequest {
model: None,
contents: vec![ContentPart {
parts: vec![TextPart {
text: "Hi".to_string(),
}],
role: None,
}],
generation_config: Some(GenerationConfig {
thinking_config: Some(ThinkingConfig { thinking_budget: 0 }),
}),
ttl: None,
};
serde_json::to_value(generate_request).unwrap()
});
// LazyLock for the cached content test body used in cache API validation
pub static CACHE_CONTENT_TEST_BODY: LazyLock<Value> = LazyLock::new(|| {
// Generate random text content to meet the minimum 1024 tokens requirement for cache API
let long_text = "You are an expert at analyzing transcripts.".repeat(50);
let cache_request = GeminiRequest {
model: Some("models/gemini-2.5-flash-lite".to_string()),
contents: vec![ContentPart {
parts: vec![TextPart { text: long_text }],
role: Some("user".to_string()),
}],
generation_config: None,
ttl: Some("300s".to_string()),
};
serde_json::to_value(cache_request).unwrap()
});

5
src/service/validation.rs → src/validation/validation_service.rs
View File

@ -5,10 +5,11 @@ use reqwest::Client;
use std::time::Instant;
use tokio::{fs, io::AsyncWriteExt, sync::mpsc};
use super::{http_client::client_builder, key_tester::validate_key};
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 {
config: KeyCheckerConfig,
@ -49,7 +50,7 @@ impl ValidationService {
// Create stream to validate keys concurrently
let valid_keys_stream = stream
.map(|key| {
validate_key(
test_generate_content_api(
self.client.clone(),
self.full_url.clone(),
key,