feat: update dependencies and enhance key validation logic
Yoo1tic
parent
30c2610925
commit
63f0599c43
|
|
@ -7,7 +7,10 @@ name = "Gemini-Keychecker"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"async-stream",
|
||||
"backon",
|
||||
"clap",
|
||||
"futures",
|
||||
"regex",
|
||||
"reqwest",
|
||||
"serde_json",
|
||||
|
|
@ -95,12 +98,45 @@ version = "1.0.98"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e16d2d3311acee920a9eb8d33b8cbc1787ce4a264e85f964c2404b969bdcd487"
|
||||
|
||||
[[package]]
|
||||
name = "async-stream"
|
||||
version = "0.3.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476"
|
||||
dependencies = [
|
||||
"async-stream-impl",
|
||||
"futures-core",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-stream-impl"
|
||||
version = "0.3.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "atomic-waker"
|
||||
version = "1.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0"
|
||||
|
||||
[[package]]
|
||||
name = "backon"
|
||||
version = "1.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "302eaff5357a264a2c42f127ecb8bac761cf99749fc3dc95677e2743991f99e7"
|
||||
dependencies = [
|
||||
"fastrand",
|
||||
"gloo-timers",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "backtrace"
|
||||
version = "0.3.75"
|
||||
|
|
@ -289,6 +325,21 @@ dependencies = [
|
|||
"percent-encoding",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures"
|
||||
version = "0.3.31"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876"
|
||||
dependencies = [
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
"futures-executor",
|
||||
"futures-io",
|
||||
"futures-sink",
|
||||
"futures-task",
|
||||
"futures-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures-channel"
|
||||
version = "0.3.31"
|
||||
|
|
@ -296,6 +347,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"futures-sink",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -304,6 +356,34 @@ version = "0.3.31"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e"
|
||||
|
||||
[[package]]
|
||||
name = "futures-executor"
|
||||
version = "0.3.31"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"futures-task",
|
||||
"futures-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures-io"
|
||||
version = "0.3.31"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6"
|
||||
|
||||
[[package]]
|
||||
name = "futures-macro"
|
||||
version = "0.3.31"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures-sink"
|
||||
version = "0.3.31"
|
||||
|
|
@ -322,10 +402,16 @@ version = "0.3.31"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81"
|
||||
dependencies = [
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
"futures-io",
|
||||
"futures-macro",
|
||||
"futures-sink",
|
||||
"futures-task",
|
||||
"memchr",
|
||||
"pin-project-lite",
|
||||
"pin-utils",
|
||||
"slab",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -357,6 +443,18 @@ version = "0.31.1"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f"
|
||||
|
||||
[[package]]
|
||||
name = "gloo-timers"
|
||||
version = "0.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bbb143cf96099802033e0d4f4963b19fd2e0b728bcf076cd9cf7f6634f092994"
|
||||
dependencies = [
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
"js-sys",
|
||||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "h2"
|
||||
version = "0.4.11"
|
||||
|
|
|
|||
|
|
@ -5,9 +5,12 @@ edition = "2024"
|
|||
|
||||
[dependencies]
|
||||
anyhow = "1.0.98"
|
||||
backon = "1"
|
||||
clap = { version = "4.5.40", features = ["derive"] }
|
||||
futures = "0.3"
|
||||
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"
|
||||
url = "2.5.4"
|
||||
async-stream = "0.3"
|
||||
149
src/main.rs
149
src/main.rs
|
|
@ -1,19 +1,20 @@
|
|||
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,
|
||||
path::PathBuf,
|
||||
sync::{Arc, LazyLock},
|
||||
io::Write,
|
||||
path::{Path, PathBuf},
|
||||
sync::LazyLock,
|
||||
thread::spawn,
|
||||
time::Instant,
|
||||
};
|
||||
use tokio::{
|
||||
sync::Semaphore,
|
||||
task::JoinSet,
|
||||
time::{Duration, sleep},
|
||||
};
|
||||
use tokio::time::Duration;
|
||||
use url::Url;
|
||||
static API_KEY_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"^AIzaSy.{33}$").unwrap());
|
||||
#[derive(Parser, Debug)]
|
||||
|
|
@ -40,7 +41,7 @@ enum KeyStatus {
|
|||
Invalid,
|
||||
Retryable(String),
|
||||
}
|
||||
fn load_keys(path: &PathBuf) -> Result<Vec<String>> {
|
||||
fn load_keys(path: &Path) -> Result<Vec<String>> {
|
||||
let keys_txt = fs::read_to_string(path)?;
|
||||
let unique_keys_set: HashSet<&str> = keys_txt
|
||||
.lines()
|
||||
|
|
@ -52,19 +53,46 @@ fn load_keys(path: &PathBuf) -> Result<Vec<String>> {
|
|||
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 validate_key_with_retry(client: &Client, api_host: &Url, key: String) -> Option<String> {
|
||||
let retry_policy = ExponentialBuilder::default()
|
||||
.with_max_times(3)
|
||||
.with_min_delay(Duration::from_secs(3))
|
||||
.with_max_delay(Duration::from_secs(8));
|
||||
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn keytest(client: &Client, api_host: &Url, keys: &str) -> Result<KeyStatus> {
|
||||
const API_PATH: &str = "v1beta/models/gemini-2.0-flash-exp:generateContent";
|
||||
const API_PATH: &str = "v1beta/models/gemini-2.0-flash:generateContent";
|
||||
let full_url = api_host.join(API_PATH)?;
|
||||
let request_body = serde_json::json!({
|
||||
"contents": [
|
||||
|
|
@ -110,65 +138,40 @@ async fn main() -> Result<()> {
|
|||
.timeout(Duration::from_millis(config.timeout_ms))
|
||||
.build()?;
|
||||
|
||||
let semaphore = Arc::new(Semaphore::new(config.concurrency));
|
||||
let mut set = JoinSet::new();
|
||||
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<String>();
|
||||
let stream = stream! {
|
||||
while let Some(item) = rx.recv().await {
|
||||
yield item;
|
||||
}
|
||||
};
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
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);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
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);
|
||||
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
|
||||
let mut output_file = fs::File::create(&config.output_path)?;
|
||||
// Write valid keys to output file
|
||||
while let Some(valid_key) = valid_keys_stream.next().await {
|
||||
// Collect valid keys
|
||||
println!("Valid key found: {}", valid_key);
|
||||
if let Err(e) = writeln!(output_file, "{}", valid_key) {
|
||||
eprintln!("Failed to write key to output file: {}", e);
|
||||
}
|
||||
}
|
||||
output_file_txt(&valid_keys, &config.output_path)?;
|
||||
println!("Total cost time:{:?}", start_time.elapsed());
|
||||
println!("Total Elapsed Time: {:?}", start_time.elapsed());
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue