feat: Reorganize project structure (#1)
* feat: implement API key validation with retry logic and refactor key loading * feat: add configuration management and client builder for API key validation * feat: restructure project modules and implement validation service for API key processingmain
Yoo1tic
GitHub
parent
8f8f8a3dff
commit
99804c5cab
|
|
@ -26,4 +26,5 @@ target
|
|||
/target
|
||||
|
||||
*.txt
|
||||
*.bak
|
||||
*.bak
|
||||
config.toml
|
||||
|
|
@ -104,6 +104,15 @@ dependencies = [
|
|||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "atomic"
|
||||
version = "0.6.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a89cbf775b137e9b968e67227ef7f775587cde3fd31b0d8599dbd0f598a48340"
|
||||
dependencies = [
|
||||
"bytemuck",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "atomic-waker"
|
||||
version = "1.1.2"
|
||||
|
|
@ -154,6 +163,12 @@ version = "3.19.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "46c5e41b57b8bba42a04676d81cb89e9ee8e859a1a66f80a5a72e1cb76b34d43"
|
||||
|
||||
[[package]]
|
||||
name = "bytemuck"
|
||||
version = "1.23.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5c76a5792e44e4abe34d3abf15636779261d45a7450612059293d1d2cfc63422"
|
||||
|
||||
[[package]]
|
||||
name = "bytes"
|
||||
version = "1.10.1"
|
||||
|
|
@ -279,6 +294,20 @@ version = "2.3.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be"
|
||||
|
||||
[[package]]
|
||||
name = "figment"
|
||||
version = "0.10.19"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8cb01cd46b0cf372153850f4c6c272d9cbea2da513e07538405148f95bd789f3"
|
||||
dependencies = [
|
||||
"atomic",
|
||||
"pear",
|
||||
"serde",
|
||||
"toml",
|
||||
"uncased",
|
||||
"version_check",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fnv"
|
||||
version = "1.0.7"
|
||||
|
|
@ -406,11 +435,14 @@ dependencies = [
|
|||
"async-stream",
|
||||
"backon",
|
||||
"clap",
|
||||
"figment",
|
||||
"futures",
|
||||
"regex",
|
||||
"reqwest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
"toml",
|
||||
"url",
|
||||
]
|
||||
|
||||
|
|
@ -721,6 +753,12 @@ dependencies = [
|
|||
"hashbrown",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "inlinable_string"
|
||||
version = "0.1.15"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c8fae54786f62fb2918dcfae3d568594e50eb9b5c25bf04371af6fe7516452fb"
|
||||
|
||||
[[package]]
|
||||
name = "io-uring"
|
||||
version = "0.7.8"
|
||||
|
|
@ -908,6 +946,29 @@ dependencies = [
|
|||
"vcpkg",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pear"
|
||||
version = "0.2.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bdeeaa00ce488657faba8ebf44ab9361f9365a97bd39ffb8a60663f57ff4b467"
|
||||
dependencies = [
|
||||
"inlinable_string",
|
||||
"pear_codegen",
|
||||
"yansi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pear_codegen"
|
||||
version = "0.2.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4bab5b985dc082b345f812b7df84e1bef27e7207b39e448439ba8bd69c93f147"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"proc-macro2-diagnostics",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "percent-encoding"
|
||||
version = "2.3.1"
|
||||
|
|
@ -950,6 +1011,19 @@ dependencies = [
|
|||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "proc-macro2-diagnostics"
|
||||
version = "0.10.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "af066a9c399a26e020ada66a034357a868728e72cd426f3adcd35f80d88d88c8"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
"version_check",
|
||||
"yansi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "quote"
|
||||
version = "1.0.40"
|
||||
|
|
@ -1176,6 +1250,15 @@ dependencies = [
|
|||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_spanned"
|
||||
version = "0.6.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bf41e0cfaf7226dca15e8197172c295a782857fcb97fad1808a166870dee75a3"
|
||||
dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_urlencoded"
|
||||
version = "0.7.1"
|
||||
|
|
@ -1371,6 +1454,47 @@ dependencies = [
|
|||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "toml"
|
||||
version = "0.8.23"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dc1beb996b9d83529a9e75c17a1686767d148d70663143c7854d8b4a09ced362"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"serde_spanned",
|
||||
"toml_datetime",
|
||||
"toml_edit",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "toml_datetime"
|
||||
version = "0.6.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "22cddaf88f4fbc13c51aebbf5f8eceb5c7c5a9da2ac40a13519eb5b0a0e8f11c"
|
||||
dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "toml_edit"
|
||||
version = "0.22.27"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "41fe8c660ae4257887cf66394862d21dbca4a6ddd26f04a3560410406a2f819a"
|
||||
dependencies = [
|
||||
"indexmap",
|
||||
"serde",
|
||||
"serde_spanned",
|
||||
"toml_datetime",
|
||||
"toml_write",
|
||||
"winnow",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "toml_write"
|
||||
version = "0.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5d99f8c9a7727884afe522e9bd5edbfc91a3312b36a77b5fb8926e4c31a41801"
|
||||
|
||||
[[package]]
|
||||
name = "tower"
|
||||
version = "0.5.2"
|
||||
|
|
@ -1441,6 +1565,15 @@ version = "0.2.5"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b"
|
||||
|
||||
[[package]]
|
||||
name = "uncased"
|
||||
version = "0.9.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e1b88fcfe09e89d3866a5c11019378088af2d24c3fbd4f0543f96b479ec90697"
|
||||
dependencies = [
|
||||
"version_check",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unicode-ident"
|
||||
version = "1.0.18"
|
||||
|
|
@ -1462,6 +1595,7 @@ dependencies = [
|
|||
"form_urlencoded",
|
||||
"idna",
|
||||
"percent-encoding",
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -1482,6 +1616,12 @@ version = "0.2.15"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426"
|
||||
|
||||
[[package]]
|
||||
name = "version_check"
|
||||
version = "0.9.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
|
||||
|
||||
[[package]]
|
||||
name = "want"
|
||||
version = "0.3.1"
|
||||
|
|
@ -1777,6 +1917,15 @@ version = "0.53.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "271414315aff87387382ec3d271b52d7ae78726f5d44ac98b4f4030c91880486"
|
||||
|
||||
[[package]]
|
||||
name = "winnow"
|
||||
version = "0.7.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f3edebf492c8125044983378ecb5766203ad3b4c2f7a922bd7dd207f6d443e95"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wit-bindgen-rt"
|
||||
version = "0.39.0"
|
||||
|
|
@ -1792,6 +1941,12 @@ version = "0.6.1"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ea2f10b9bb0928dfb1b42b65e1f9e36f7f54dbdf08457afefb38afcdec4fa2bb"
|
||||
|
||||
[[package]]
|
||||
name = "yansi"
|
||||
version = "1.0.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cfe53a6657fd280eaa890a3bc59152892ffa3e30101319d168b781ed6529b049"
|
||||
|
||||
[[package]]
|
||||
name = "yoke"
|
||||
version = "0.8.0"
|
||||
|
|
|
|||
|
|
@ -12,5 +12,8 @@ regex = "1.11.1"
|
|||
reqwest = { version = "0.12.22", features = ["json"] }
|
||||
serde_json = "1.0.140"
|
||||
tokio = { version = "1.46", features = ["macros", "rt-multi-thread", "time"] }
|
||||
url = "2.5.4"
|
||||
async-stream = "0.3"
|
||||
url = { version = "2.5.4", features = ["serde"] }
|
||||
async-stream = "0.3"
|
||||
figment = { version = "0.10.19", features = ["env", "toml"] }
|
||||
serde = { version = "1.0.219", features = ["derive"] }
|
||||
toml = "0.8"
|
||||
|
|
|
|||
|
|
@ -0,0 +1,46 @@
|
|||
use anyhow::Result;
|
||||
use std::{
|
||||
collections::HashSet,
|
||||
fs,
|
||||
path::Path,
|
||||
str::FromStr,
|
||||
};
|
||||
use crate::types::ApiKey;
|
||||
|
||||
/// Load and validate API keys from a file
|
||||
/// Returns a vector of unique, valid API keys
|
||||
pub fn load_keys(path: &Path) -> Result<Vec<ApiKey>> {
|
||||
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())
|
||||
.filter(|line| !line.is_empty())
|
||||
.collect();
|
||||
|
||||
let mut keys = Vec::new();
|
||||
let mut valid_keys_for_backup = Vec::new();
|
||||
|
||||
for key_str in unique_keys_set {
|
||||
match ApiKey::from_str(key_str) {
|
||||
Ok(api_key) => {
|
||||
keys.push(api_key.clone());
|
||||
valid_keys_for_backup.push(api_key.as_str().to_string());
|
||||
}
|
||||
Err(e) => eprintln!("Skipping invalid key: {}", e),
|
||||
}
|
||||
}
|
||||
|
||||
// Write validated keys to backup.txt
|
||||
let backup_content = valid_keys_for_backup.join("\n");
|
||||
if let Err(e) = fs::write("backup.txt", backup_content) {
|
||||
eprintln!("Failed to write backup file: {}", e);
|
||||
} else {
|
||||
println!(
|
||||
"Backup file created with {} valid keys",
|
||||
valid_keys_for_backup.len()
|
||||
);
|
||||
}
|
||||
|
||||
Ok(keys)
|
||||
}
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
pub mod local;
|
||||
|
||||
pub use local::*;
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
pub mod input;
|
||||
pub mod output;
|
||||
|
||||
pub use input::load_keys;
|
||||
pub use output::write_keys_txt_file;
|
||||
|
|
@ -0,0 +1,24 @@
|
|||
use crate::types::ApiKey;
|
||||
use anyhow::Result;
|
||||
use std::{fs, io::Write};
|
||||
use toml::Value;
|
||||
|
||||
// Write valid key to output file
|
||||
pub fn write_keys_txt_file(file: &mut fs::File, key: &ApiKey) -> Result<()> {
|
||||
writeln!(file, "{}", key.as_str())?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// Write valid key to output file in Clewdr format
|
||||
pub fn write_keys_clewdr_format(file: &mut fs::File, key: &ApiKey) -> Result<()> {
|
||||
let mut table = toml::value::Table::new();
|
||||
table.insert("key".to_string(), Value::String(key.as_str().to_string()));
|
||||
|
||||
let gemini_keys = Value::Array(vec![Value::Table(table)]);
|
||||
let mut root = toml::value::Table::new();
|
||||
root.insert("gemini_keys".to_string(), gemini_keys);
|
||||
|
||||
let toml_string = toml::to_string(&Value::Table(root))?;
|
||||
write!(file, "{}", toml_string)?;
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
pub mod local;
|
||||
|
||||
pub use local::*;
|
||||
|
|
@ -0,0 +1,15 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use reqwest::Client;
|
||||
|
||||
use crate::config::KeyCheckerConfig;
|
||||
|
||||
pub fn client_builder(config: &KeyCheckerConfig) -> Result<Client, reqwest::Error> {
|
||||
let mut builder = Client::builder().timeout(Duration::from_secs(config.timeout_sec));
|
||||
|
||||
if let Some(ref proxy_url) = config.proxy {
|
||||
builder = builder.proxy(reqwest::Proxy::all(proxy_url.clone())?);
|
||||
}
|
||||
|
||||
builder.build()
|
||||
}
|
||||
|
|
@ -0,0 +1,94 @@
|
|||
use anyhow::{Ok, Result};
|
||||
use figment::{
|
||||
Figment,
|
||||
providers::{Env, Format, Toml},
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::LazyLock;
|
||||
use url::Url;
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct KeyCheckerConfig {
|
||||
// Input file path containing API keys to check.
|
||||
#[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)]
|
||||
pub backup_path: PathBuf,
|
||||
|
||||
// API host URL for key validation.
|
||||
#[serde(default = "default_api_host")]
|
||||
pub api_host: Url,
|
||||
|
||||
// Request timeout in seconds.
|
||||
#[serde(default)]
|
||||
pub timeout_sec: u64,
|
||||
|
||||
// Maximum number of concurrent requests.
|
||||
#[serde(default)]
|
||||
pub concurrency: usize,
|
||||
|
||||
// Optional proxy URL for HTTP requests (e.g., --proxy http://user:pass@host:port).
|
||||
#[serde(default)]
|
||||
pub proxy: Option<Url>,
|
||||
}
|
||||
|
||||
impl Default for KeyCheckerConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
input_path: default_input_path(),
|
||||
output_path: default_output_path(),
|
||||
backup_path: default_backup_path(),
|
||||
api_host: default_api_host(),
|
||||
timeout_sec: default_timeout(),
|
||||
concurrency: default_concurrency(),
|
||||
proxy: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
impl KeyCheckerConfig {
|
||||
pub fn load_config() -> Result<Self> {
|
||||
// Define the path to the configuration file
|
||||
static CONFIG_PATH: LazyLock<PathBuf> = LazyLock::new(|| "config.toml".into());
|
||||
|
||||
// Check if config.toml exists, if not create it with default values
|
||||
if !CONFIG_PATH.exists() {
|
||||
let default_config = Self::default();
|
||||
let toml_content = toml::to_string_pretty(&default_config)?;
|
||||
fs::write(CONFIG_PATH.as_path(), toml_content)?;
|
||||
}
|
||||
|
||||
// Load configuration from config.toml, environment variables, and defaults
|
||||
let config = Figment::new()
|
||||
.merge(Toml::file(CONFIG_PATH.as_path()))
|
||||
.merge(Env::prefixed("KEYCHECKER_"))
|
||||
.extract()?;
|
||||
Ok(config)
|
||||
}
|
||||
}
|
||||
|
||||
fn default_input_path() -> PathBuf {
|
||||
"keys.txt".into()
|
||||
}
|
||||
fn default_output_path() -> PathBuf {
|
||||
"output_keys.txt".into()
|
||||
}
|
||||
fn default_backup_path() -> PathBuf {
|
||||
"backup_keys.txt".into()
|
||||
}
|
||||
fn default_api_host() -> Url {
|
||||
Url::parse("https://generativelanguage.googleapis.com/").unwrap()
|
||||
}
|
||||
fn default_timeout() -> u64 {
|
||||
20
|
||||
}
|
||||
fn default_concurrency() -> usize {
|
||||
30
|
||||
}
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
mod basic_config;
|
||||
mod basic_client;
|
||||
|
||||
pub use basic_config::KeyCheckerConfig;
|
||||
pub use basic_client::client_builder;
|
||||
|
|
@ -0,0 +1,80 @@
|
|||
use anyhow::Result;
|
||||
use backon::{ExponentialBuilder, Retryable};
|
||||
use reqwest::{Client, StatusCode};
|
||||
use serde_json;
|
||||
use tokio::time::Duration;
|
||||
use url::Url;
|
||||
|
||||
use crate::types::{KeyStatus, ApiKey};
|
||||
|
||||
pub async fn validate_key_with_retry(client: &Client, api_host: &Url, key: ApiKey) -> Option<ApiKey> {
|
||||
let retry_policy = ExponentialBuilder::default()
|
||||
.with_max_times(3)
|
||||
.with_min_delay(Duration::from_secs(3))
|
||||
.with_max_delay(Duration::from_secs(5));
|
||||
|
||||
let result = (|| async {
|
||||
match keytest(&client, &api_host, &key).await {
|
||||
Ok(KeyStatus::Valid) => {
|
||||
println!("Key: {}... -> SUCCESS", &key.as_str()[..10]);
|
||||
Ok(Some(key.clone()))
|
||||
}
|
||||
Ok(KeyStatus::Invalid) => {
|
||||
println!("Key: {}... -> INVALID (Forbidden)", &key.as_str()[..10]);
|
||||
Ok(None)
|
||||
}
|
||||
Ok(KeyStatus::Retryable(reason)) => {
|
||||
eprintln!("Key: {}... -> RETRYABLE (Reason: {})", &key.as_str()[..10], reason);
|
||||
Err(anyhow::anyhow!("Retryable error: {}", reason))
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("Key: {}... -> NETWORK ERROR (Reason: {})", &key.as_str()[..10], e);
|
||||
Err(e)
|
||||
}
|
||||
}
|
||||
})
|
||||
.retry(retry_policy)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(key_result) => key_result,
|
||||
Err(_) => {
|
||||
eprintln!("Key: {}... -> FAILED after all retries.", &key.as_str()[..10]);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn keytest(client: &Client, api_host: &Url, key: &ApiKey) -> Result<KeyStatus> {
|
||||
const API_PATH: &str = "v1beta/models/gemini-2.0-flash-exp:generateContent";
|
||||
let full_url = api_host.join(API_PATH)?;
|
||||
|
||||
let request_body = serde_json::json!({
|
||||
"contents": [
|
||||
{
|
||||
"parts": [
|
||||
{
|
||||
"text": "Hi"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
});
|
||||
|
||||
let response = client
|
||||
.post(full_url)
|
||||
.header("Content-Type", "application/json")
|
||||
.header("X-goog-api-key", key.as_str())
|
||||
.json(&request_body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let status = response.status();
|
||||
|
||||
let key_status = match status {
|
||||
StatusCode::OK => KeyStatus::Valid,
|
||||
StatusCode::FORBIDDEN | StatusCode::UNAUTHORIZED => KeyStatus::Invalid,
|
||||
other => KeyStatus::Retryable(format!("Received status {}, will retry.", other)),
|
||||
};
|
||||
Ok(key_status)
|
||||
}
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
pub mod config;
|
||||
pub mod types;
|
||||
pub mod key_validator;
|
||||
pub mod adapters;
|
||||
pub mod validation;
|
||||
229
src/main.rs
229
src/main.rs
|
|
@ -1,230 +1,17 @@
|
|||
use anyhow::Result;
|
||||
use async_stream::stream;
|
||||
use backon::{ExponentialBuilder, Retryable};
|
||||
use clap::Parser;
|
||||
use futures::{pin_mut, stream::StreamExt};
|
||||
use regex::Regex;
|
||||
use reqwest::{Client, StatusCode};
|
||||
use std::{
|
||||
collections::HashSet,
|
||||
fs,
|
||||
io::Write,
|
||||
path::{Path, PathBuf},
|
||||
sync::LazyLock,
|
||||
thread::spawn,
|
||||
time::Instant,
|
||||
};
|
||||
use tokio::time::Duration;
|
||||
use url::Url;
|
||||
|
||||
// Regex pattern for validating Google API keys (AIzaSy followed by 33 characters)
|
||||
static API_KEY_REGEX: LazyLock<Regex> = 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,
|
||||
|
||||
/// 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 (supports http://user:pass@host:port)
|
||||
#[arg(long, short = 'x')]
|
||||
proxy: Option<Url>,
|
||||
}
|
||||
/// 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<Vec<String>> {
|
||||
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())
|
||||
.filter(|line| !line.is_empty())
|
||||
.filter(|line| API_KEY_REGEX.is_match(line))
|
||||
.collect();
|
||||
let keys: Vec<String> = unique_keys_set.into_iter().map(String::from).collect();
|
||||
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<String> {
|
||||
// Configure exponential backoff retry policy
|
||||
let retry_policy = ExponentialBuilder::default()
|
||||
.with_max_times(3)
|
||||
.with_min_delay(Duration::from_secs(5))
|
||||
.with_max_delay(Duration::from_secs(10));
|
||||
|
||||
let result = (|| async {
|
||||
match keytest(&client, &api_host, &key).await {
|
||||
Ok(KeyStatus::Valid) => {
|
||||
println!("Key: {}... -> SUCCESS", &key[..10]);
|
||||
Ok(Some(key.clone()))
|
||||
}
|
||||
Ok(KeyStatus::Invalid) => {
|
||||
println!("Key: {}... -> INVALID (Forbidden)", &key[..10]);
|
||||
Ok(None)
|
||||
}
|
||||
Ok(KeyStatus::Retryable(reason)) => {
|
||||
eprintln!("Key: {}... -> RETRYABLE (Reason: {})", &key[..10], reason);
|
||||
Err(anyhow::anyhow!("Retryable error: {}", reason))
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("Key: {}... -> NETWORK ERROR (Reason: {})", &key[..10], e);
|
||||
Err(e)
|
||||
}
|
||||
}
|
||||
})
|
||||
.retry(retry_policy)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(key_result) => key_result,
|
||||
Err(_) => {
|
||||
eprintln!("Key: {}... -> FAILED after all retries.", &key[..10]);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 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<KeyStatus> {
|
||||
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": [
|
||||
{
|
||||
"parts": [
|
||||
{
|
||||
"text": "Hi"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
});
|
||||
|
||||
let response = client
|
||||
.post(full_url)
|
||||
.header("Content-Type", "application/json")
|
||||
.header("X-goog-api-key", keys)
|
||||
.json(&request_body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let status = response.status();
|
||||
|
||||
let key_status = match status {
|
||||
// 200 OK - Key is valid
|
||||
StatusCode::OK => KeyStatus::Valid,
|
||||
|
||||
// 403 & 401 - Key is invalid or unauthorized
|
||||
StatusCode::FORBIDDEN | StatusCode::UNAUTHORIZED => KeyStatus::Invalid,
|
||||
|
||||
// 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<Client> {
|
||||
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 {
|
||||
let mut proxy = reqwest::Proxy::all(proxy_url.clone())?;
|
||||
|
||||
// Extract username and password from URL if present
|
||||
if !proxy_url.username().is_empty() {
|
||||
let username = proxy_url.username();
|
||||
let password = proxy_url.password().unwrap_or("");
|
||||
proxy = proxy.basic_auth(username, password);
|
||||
}
|
||||
|
||||
client_builder = client_builder.proxy(proxy);
|
||||
}
|
||||
|
||||
client_builder.build().map_err(Into::into)
|
||||
}
|
||||
use gemini_keychecker::config::{KeyCheckerConfig, client_builder};
|
||||
use gemini_keychecker::adapters::load_keys;
|
||||
use gemini_keychecker::validation::ValidationService;
|
||||
|
||||
/// Main function - orchestrates the key validation process
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
let start_time = Instant::now();
|
||||
let config = KeyCheckerConfig::parse();
|
||||
let config = KeyCheckerConfig::load_config().unwrap();
|
||||
let keys = load_keys(&config.input_path)?;
|
||||
let client = build_client(&config)?;
|
||||
let client = client_builder(&config)?;
|
||||
|
||||
// Create channel for streaming keys from producer to consumer
|
||||
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<String>();
|
||||
let stream = stream! {
|
||||
while let Some(item) = rx.recv().await {
|
||||
yield item;
|
||||
}
|
||||
};
|
||||
let validation_service = ValidationService::new(config, client);
|
||||
validation_service.validate_keys(keys).await?;
|
||||
|
||||
// Spawn producer thread to send keys through channel
|
||||
spawn(move || {
|
||||
for key in keys {
|
||||
if API_KEY_REGEX.is_match(&key) {
|
||||
if let Err(e) = tx.send(key) {
|
||||
eprintln!("Failed to send key: {}", e);
|
||||
}
|
||||
} else {
|
||||
eprintln!("Invalid key format: {}", key);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// 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 for writing valid keys
|
||||
let mut output_file = fs::File::create(&config.output_path)?;
|
||||
|
||||
// Process validated keys and write to output file
|
||||
while let Some(valid_key) = valid_keys_stream.next().await {
|
||||
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(())
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,56 @@
|
|||
use regex::Regex;
|
||||
use std::str::FromStr;
|
||||
use std::sync::LazyLock;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum KeyStatus {
|
||||
Valid,
|
||||
Invalid,
|
||||
Retryable(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ApiKey {
|
||||
inner: String,
|
||||
}
|
||||
|
||||
impl ApiKey {
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.inner
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum KeyValidationError {
|
||||
InvalidFormat(String),
|
||||
}
|
||||
|
||||
impl std::fmt::Display for KeyValidationError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
KeyValidationError::InvalidFormat(msg) => write!(f, "Invalid key format: {}", msg),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for KeyValidationError {}
|
||||
|
||||
impl FromStr for ApiKey {
|
||||
type Err = KeyValidationError;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
static RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"^AIzaSy.{33}$").unwrap());
|
||||
|
||||
let cleaned = s.trim();
|
||||
|
||||
if RE.is_match(cleaned) {
|
||||
Ok(Self {
|
||||
inner: cleaned.to_string(),
|
||||
})
|
||||
} else {
|
||||
Err(KeyValidationError::InvalidFormat(
|
||||
"Google API key must start with 'AIzaSy' followed by 33 characters".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,64 @@
|
|||
use anyhow::Result;
|
||||
use async_stream::stream;
|
||||
use futures::{pin_mut, stream::StreamExt};
|
||||
use reqwest::Client;
|
||||
use std::{fs, time::Instant};
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::config::KeyCheckerConfig;
|
||||
use crate::adapters::write_keys_txt_file;
|
||||
use crate::key_validator::validate_key_with_retry;
|
||||
use crate::types::ApiKey;
|
||||
|
||||
pub struct ValidationService {
|
||||
config: KeyCheckerConfig,
|
||||
client: Client,
|
||||
}
|
||||
|
||||
impl ValidationService {
|
||||
pub fn new(config: KeyCheckerConfig, client: Client) -> Self {
|
||||
Self { config, client }
|
||||
}
|
||||
|
||||
pub async fn validate_keys(&self, keys: Vec<ApiKey>) -> Result<()> {
|
||||
let start_time = Instant::now();
|
||||
|
||||
// Create channel for streaming keys from producer to consumer
|
||||
let (tx, mut rx) = mpsc::unbounded_channel::<ApiKey>();
|
||||
let stream = stream! {
|
||||
while let Some(item) = rx.recv().await {
|
||||
yield item;
|
||||
}
|
||||
};
|
||||
|
||||
// Spawn producer task to send keys through channel
|
||||
tokio::spawn(async move {
|
||||
for key in keys {
|
||||
if let Err(e) = tx.send(key) {
|
||||
eprintln!("Failed to send key: {}", e);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Create stream to validate keys concurrently
|
||||
let valid_keys_stream = stream
|
||||
.map(|key| validate_key_with_retry(&self.client, &self.config.api_host, key))
|
||||
.buffer_unordered(self.config.concurrency)
|
||||
.filter_map(|r| async { r });
|
||||
pin_mut!(valid_keys_stream);
|
||||
|
||||
// Open output file for writing valid keys
|
||||
let mut output_file = fs::File::create(&self.config.output_path)?;
|
||||
|
||||
// Process validated keys and write to output file
|
||||
while let Some(valid_key) = valid_keys_stream.next().await {
|
||||
println!("Valid key found: {}", valid_key.as_str());
|
||||
if let Err(e) = write_keys_txt_file(&mut output_file, &valid_key) {
|
||||
eprintln!("Failed to write key to output file: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
println!("Total Elapsed Time: {:?}", start_time.elapsed());
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue