feat: add Gemini-Keychecker tool for API key validation
parent
146e660c8e
commit
30c2610925
|
@ -19,3 +19,11 @@ target
|
||||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||||
#.idea/
|
#.idea/
|
||||||
|
|
||||||
|
|
||||||
|
# Added by cargo
|
||||||
|
|
||||||
|
/target
|
||||||
|
|
||||||
|
*.txt
|
||||||
|
*.bak
|
File diff suppressed because it is too large
Load Diff
|
@ -0,0 +1,13 @@
|
||||||
|
[package]
|
||||||
|
name = "Gemini-Keychecker"
|
||||||
|
version = "0.1.0"
|
||||||
|
edition = "2024"
|
||||||
|
|
||||||
|
[dependencies]
|
||||||
|
anyhow = "1.0.98"
|
||||||
|
clap = { version = "4.5.40", features = ["derive"] }
|
||||||
|
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"
|
10
README.md
10
README.md
|
@ -1 +1,11 @@
|
||||||
# Gemini-Keychecker
|
# Gemini-Keychecker
|
||||||
|
|
||||||
|
A tool to check and backup API keys
|
||||||
|
```
|
||||||
|
Options:
|
||||||
|
-i, --input-path <INPUT_PATH> [default: keys.txt]
|
||||||
|
-o, --output-path <OUTPUT_PATH> [default: output_keys.txt]
|
||||||
|
-u, --api-host <API_HOST> [default: https://generativelanguage.googleapis.com/]
|
||||||
|
-t, --timeout-ms <TIMEOUT_MS> [default: 5000]
|
||||||
|
-c, --concurrency <CONCURRENCY> [default: 30]
|
||||||
|
```
|
|
@ -0,0 +1,174 @@
|
||||||
|
use anyhow::Result;
|
||||||
|
use clap::Parser;
|
||||||
|
use regex::Regex;
|
||||||
|
use reqwest::{Client, StatusCode};
|
||||||
|
use std::{
|
||||||
|
collections::HashSet,
|
||||||
|
fs,
|
||||||
|
path::PathBuf,
|
||||||
|
sync::{Arc, LazyLock},
|
||||||
|
time::Instant,
|
||||||
|
};
|
||||||
|
use tokio::{
|
||||||
|
sync::Semaphore,
|
||||||
|
task::JoinSet,
|
||||||
|
time::{Duration, sleep},
|
||||||
|
};
|
||||||
|
use url::Url;
|
||||||
|
static API_KEY_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"^AIzaSy.{33}$").unwrap());
|
||||||
|
#[derive(Parser, Debug)]
|
||||||
|
#[command(version, about = "A tool to check and backup API keys", long_about = None)]
|
||||||
|
struct KeyCheckerConfig {
|
||||||
|
#[arg(long, short = 'i', default_value = "keys.txt")]
|
||||||
|
input_path: PathBuf,
|
||||||
|
|
||||||
|
#[arg(long, short = 'o', default_value = "output_keys.txt")]
|
||||||
|
output_path: PathBuf,
|
||||||
|
|
||||||
|
#[arg(long, short = 'u', default_value = "https://generativelanguage.googleapis.com/")]
|
||||||
|
api_host: Url,
|
||||||
|
|
||||||
|
#[arg(long, short = 't', default_value_t = 5000)]
|
||||||
|
timeout_ms: u64,
|
||||||
|
|
||||||
|
#[arg(long, short = 'c', default_value_t = 30)]
|
||||||
|
concurrency: usize,
|
||||||
|
}
|
||||||
|
#[derive(Debug)]
|
||||||
|
enum KeyStatus {
|
||||||
|
Valid,
|
||||||
|
Invalid,
|
||||||
|
Retryable(String),
|
||||||
|
}
|
||||||
|
fn load_keys(path: &PathBuf) -> Result<Vec<String>> {
|
||||||
|
let keys_txt = fs::read_to_string(path)?;
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn output_file_txt(keys: &[String], output_path: &PathBuf) -> Result<()> {
|
||||||
|
let content = keys.join("\n");
|
||||||
|
fs::write(output_path, content)?;
|
||||||
|
println!(
|
||||||
|
"Successfully wrote {} keys to {:?}",
|
||||||
|
keys.len(),
|
||||||
|
output_path
|
||||||
|
);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
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)?;
|
||||||
|
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
|
||||||
|
StatusCode::OK => KeyStatus::Valid,
|
||||||
|
|
||||||
|
// 403 & 401
|
||||||
|
StatusCode::FORBIDDEN | StatusCode::UNAUTHORIZED => KeyStatus::Invalid,
|
||||||
|
|
||||||
|
// Other Status Code
|
||||||
|
other => KeyStatus::Retryable(format!("Received status {}, will retry.", other)),
|
||||||
|
};
|
||||||
|
Ok(key_status)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::main]
|
||||||
|
async fn main() -> Result<()> {
|
||||||
|
let start_time = Instant::now();
|
||||||
|
let config = KeyCheckerConfig::parse();
|
||||||
|
let keys = load_keys(&config.input_path)?;
|
||||||
|
let client = Client::builder()
|
||||||
|
.timeout(Duration::from_millis(config.timeout_ms))
|
||||||
|
.build()?;
|
||||||
|
|
||||||
|
let semaphore = Arc::new(Semaphore::new(config.concurrency));
|
||||||
|
let mut set = JoinSet::new();
|
||||||
|
|
||||||
|
for key in keys {
|
||||||
|
let client_clone = client.clone();
|
||||||
|
let api_host_clone = config.api_host.clone();
|
||||||
|
let semaphore_clone = Arc::clone(&semaphore);
|
||||||
|
|
||||||
|
set.spawn(async move {
|
||||||
|
const MAX_RETRIES: u32 = 3;
|
||||||
|
let _permit = semaphore_clone.acquire().await.unwrap();
|
||||||
|
for attempt in 0..MAX_RETRIES {
|
||||||
|
match keytest(&client_clone, &api_host_clone, &key).await {
|
||||||
|
Ok(KeyStatus::Valid) => {
|
||||||
|
println!("Key: {}... -> SUCCESS", &key[..10]);
|
||||||
|
return Some(key);
|
||||||
|
}
|
||||||
|
Ok(KeyStatus::Invalid) => {
|
||||||
|
println!("Key: {}... -> INVALID (Forbidden)", &key[..10]);
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
Ok(KeyStatus::Retryable(reason)) => {
|
||||||
|
eprintln!(
|
||||||
|
"Key: {}... -> RETRYABLE (Attempt {}/{}, Reason: {})",
|
||||||
|
&key[..10],
|
||||||
|
attempt + 1,
|
||||||
|
MAX_RETRIES,
|
||||||
|
reason
|
||||||
|
);
|
||||||
|
if attempt < MAX_RETRIES - 1 {
|
||||||
|
sleep(Duration::from_secs(2_u64.pow(attempt))).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
eprintln!(
|
||||||
|
"Key: {}... -> NETWORK ERROR (Attempt {}/{}, Reason: {})",
|
||||||
|
&key[..10],
|
||||||
|
attempt + 1,
|
||||||
|
MAX_RETRIES,
|
||||||
|
e
|
||||||
|
);
|
||||||
|
if attempt < MAX_RETRIES - 1 {
|
||||||
|
sleep(Duration::from_secs(2_u64.pow(attempt))).await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
eprintln!("Key: {}... -> FAILED after all retries.", &key[..10]);
|
||||||
|
None
|
||||||
|
});
|
||||||
|
}
|
||||||
|
let mut valid_keys = Vec::new();
|
||||||
|
while let Some(res) = set.join_next().await {
|
||||||
|
if let Ok(Some(key)) = res {
|
||||||
|
valid_keys.push(key);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
output_file_txt(&valid_keys, &config.output_path)?;
|
||||||
|
println!("Total cost time:{:?}", start_time.elapsed());
|
||||||
|
Ok(())
|
||||||
|
}
|
Loading…
Reference in New Issue