feat: update KeyCheckerConfig to use clap for argument parsing and improve client builder with proper timeout handling
parent
99804c5cab
commit
0056b4fa7b
|
@ -5,9 +5,9 @@ use reqwest::Client;
|
||||||
use crate::config::KeyCheckerConfig;
|
use crate::config::KeyCheckerConfig;
|
||||||
|
|
||||||
pub fn client_builder(config: &KeyCheckerConfig) -> Result<Client, reqwest::Error> {
|
pub fn client_builder(config: &KeyCheckerConfig) -> Result<Client, reqwest::Error> {
|
||||||
let mut builder = Client::builder().timeout(Duration::from_secs(config.timeout_sec));
|
let mut builder = Client::builder().timeout(Duration::from_secs(config.timeout_sec()));
|
||||||
|
|
||||||
if let Some(ref proxy_url) = config.proxy {
|
if let Some(ref proxy_url) = config.proxy() {
|
||||||
builder = builder.proxy(reqwest::Proxy::all(proxy_url.clone())?);
|
builder = builder.proxy(reqwest::Proxy::all(proxy_url.clone())?);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
@ -1,7 +1,8 @@
|
||||||
use anyhow::{Ok, Result};
|
use anyhow::{Ok, Result};
|
||||||
|
use clap::Parser;
|
||||||
use figment::{
|
use figment::{
|
||||||
Figment,
|
Figment,
|
||||||
providers::{Env, Format, Toml},
|
providers::{Env, Format, Serialized, Toml},
|
||||||
};
|
};
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use std::fs;
|
use std::fs;
|
||||||
|
@ -9,46 +10,53 @@ use std::path::PathBuf;
|
||||||
use std::sync::LazyLock;
|
use std::sync::LazyLock;
|
||||||
use url::Url;
|
use url::Url;
|
||||||
|
|
||||||
#[derive(Debug, Serialize, Deserialize)]
|
#[derive(Debug, Serialize, Deserialize, Parser)]
|
||||||
pub struct KeyCheckerConfig {
|
pub struct KeyCheckerConfig {
|
||||||
// Input file path containing API keys to check.
|
// Input file path containing API keys to check.
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub input_path: PathBuf,
|
#[arg(short, long)]
|
||||||
|
input_path: Option<PathBuf>,
|
||||||
|
|
||||||
// Output file path for valid API keys.
|
// Output file path for valid API keys.
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub output_path: PathBuf,
|
#[arg(short, long)]
|
||||||
|
output_path: Option<PathBuf>,
|
||||||
|
|
||||||
// Backup file path for all API keys.
|
// Backup file path for all API keys.
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub backup_path: PathBuf,
|
#[arg(short, long)]
|
||||||
|
backup_path: Option<PathBuf>,
|
||||||
|
|
||||||
// API host URL for key validation.
|
// API host URL for key validation.
|
||||||
#[serde(default = "default_api_host")]
|
#[serde(default)]
|
||||||
pub api_host: Url,
|
#[arg(short, long)]
|
||||||
|
api_host: Option<Url>,
|
||||||
|
|
||||||
// Request timeout in seconds.
|
// Request timeout in seconds.
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub timeout_sec: u64,
|
#[arg(short, long)]
|
||||||
|
timeout_sec: Option<u64>,
|
||||||
|
|
||||||
// Maximum number of concurrent requests.
|
// Maximum number of concurrent requests.
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub concurrency: usize,
|
#[arg(short, long)]
|
||||||
|
concurrency: Option<usize>,
|
||||||
|
|
||||||
// Optional proxy URL for HTTP requests (e.g., --proxy http://user:pass@host:port).
|
// Optional proxy URL for HTTP requests (e.g., --proxy http://user:pass@host:port).
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
pub proxy: Option<Url>,
|
#[arg(short, long)]
|
||||||
|
proxy: Option<Url>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for KeyCheckerConfig {
|
impl Default for KeyCheckerConfig {
|
||||||
fn default() -> Self {
|
fn default() -> Self {
|
||||||
Self {
|
Self {
|
||||||
input_path: default_input_path(),
|
input_path: Some(default_input_path()),
|
||||||
output_path: default_output_path(),
|
output_path: Some(default_output_path()),
|
||||||
backup_path: default_backup_path(),
|
backup_path: Some(default_backup_path()),
|
||||||
api_host: default_api_host(),
|
api_host: Some(default_api_host()),
|
||||||
timeout_sec: default_timeout(),
|
timeout_sec: Some(default_timeout()),
|
||||||
concurrency: default_concurrency(),
|
concurrency: Some(default_concurrency()),
|
||||||
proxy: None,
|
proxy: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
@ -67,16 +75,40 @@ impl KeyCheckerConfig {
|
||||||
|
|
||||||
// Load configuration from config.toml, environment variables, and defaults
|
// Load configuration from config.toml, environment variables, and defaults
|
||||||
let config = Figment::new()
|
let config = Figment::new()
|
||||||
|
.merge(Serialized::defaults(Self::default()))
|
||||||
.merge(Toml::file(CONFIG_PATH.as_path()))
|
.merge(Toml::file(CONFIG_PATH.as_path()))
|
||||||
.merge(Env::prefixed("KEYCHECKER_"))
|
.merge(Env::prefixed("KEYCHECKER_"))
|
||||||
|
.merge(Serialized::defaults(Self::parse()))
|
||||||
.extract()?;
|
.extract()?;
|
||||||
Ok(config)
|
Ok(config)
|
||||||
}
|
}
|
||||||
|
pub fn input_path(&self) -> PathBuf {
|
||||||
|
self.input_path.clone().unwrap_or_else(default_input_path)
|
||||||
|
}
|
||||||
|
pub fn output_path(&self) -> PathBuf {
|
||||||
|
self.output_path.clone().unwrap_or_else(default_output_path)
|
||||||
|
}
|
||||||
|
pub fn backup_path(&self) -> PathBuf {
|
||||||
|
self.backup_path.clone().unwrap_or_else(default_backup_path)
|
||||||
|
}
|
||||||
|
pub fn api_host(&self) -> Url {
|
||||||
|
self.api_host.clone().unwrap_or_else(default_api_host)
|
||||||
|
}
|
||||||
|
pub fn timeout_sec(&self) -> u64 {
|
||||||
|
self.timeout_sec.unwrap_or_else(default_timeout)
|
||||||
|
}
|
||||||
|
pub fn concurrency(&self) -> usize {
|
||||||
|
self.concurrency.unwrap_or_else(default_concurrency)
|
||||||
|
}
|
||||||
|
pub fn proxy(&self) -> Option<Url> {
|
||||||
|
self.proxy.clone()
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn default_input_path() -> PathBuf {
|
fn default_input_path() -> PathBuf {
|
||||||
"keys.txt".into()
|
"keys.txt".into()
|
||||||
}
|
}
|
||||||
|
|
||||||
fn default_output_path() -> PathBuf {
|
fn default_output_path() -> PathBuf {
|
||||||
"output_keys.txt".into()
|
"output_keys.txt".into()
|
||||||
}
|
}
|
||||||
|
|
|
@ -5,32 +5,42 @@ use serde_json;
|
||||||
use tokio::time::Duration;
|
use tokio::time::Duration;
|
||||||
use url::Url;
|
use url::Url;
|
||||||
|
|
||||||
use crate::types::{KeyStatus, ApiKey};
|
use crate::types::{ApiKey, KeyStatus};
|
||||||
|
|
||||||
pub async fn validate_key_with_retry(client: &Client, api_host: &Url, key: ApiKey) -> Option<ApiKey> {
|
pub async fn validate_key_with_retry(
|
||||||
|
client: Client,
|
||||||
|
api_host: Url,
|
||||||
|
key: ApiKey,
|
||||||
|
) -> Option<ApiKey> {
|
||||||
let retry_policy = ExponentialBuilder::default()
|
let retry_policy = ExponentialBuilder::default()
|
||||||
.with_max_times(3)
|
.with_max_times(3)
|
||||||
.with_min_delay(Duration::from_secs(3))
|
.with_min_delay(Duration::from_secs(3))
|
||||||
.with_max_delay(Duration::from_secs(5));
|
.with_max_delay(Duration::from_secs(5));
|
||||||
|
|
||||||
let result = (|| async {
|
let result = (async || match keytest(client.to_owned(), &api_host, &key).await {
|
||||||
match keytest(&client, &api_host, &key).await {
|
Ok(KeyStatus::Valid) => {
|
||||||
Ok(KeyStatus::Valid) => {
|
println!("Key: {}... -> SUCCESS", &key.as_str()[..10]);
|
||||||
println!("Key: {}... -> SUCCESS", &key.as_str()[..10]);
|
Ok(Some(key.clone()))
|
||||||
Ok(Some(key.clone()))
|
}
|
||||||
}
|
Ok(KeyStatus::Invalid) => {
|
||||||
Ok(KeyStatus::Invalid) => {
|
println!("Key: {}... -> INVALID (Forbidden)", &key.as_str()[..10]);
|
||||||
println!("Key: {}... -> INVALID (Forbidden)", &key.as_str()[..10]);
|
Ok(None)
|
||||||
Ok(None)
|
}
|
||||||
}
|
Ok(KeyStatus::Retryable(reason)) => {
|
||||||
Ok(KeyStatus::Retryable(reason)) => {
|
eprintln!(
|
||||||
eprintln!("Key: {}... -> RETRYABLE (Reason: {})", &key.as_str()[..10], reason);
|
"Key: {}... -> RETRYABLE (Reason: {})",
|
||||||
Err(anyhow::anyhow!("Retryable error: {}", reason))
|
&key.as_str()[..10],
|
||||||
}
|
reason
|
||||||
Err(e) => {
|
);
|
||||||
eprintln!("Key: {}... -> NETWORK ERROR (Reason: {})", &key.as_str()[..10], e);
|
Err(anyhow::anyhow!("Retryable error: {}", reason))
|
||||||
Err(e)
|
}
|
||||||
}
|
Err(e) => {
|
||||||
|
eprintln!(
|
||||||
|
"Key: {}... -> NETWORK ERROR (Reason: {})",
|
||||||
|
&key.as_str()[..10],
|
||||||
|
e
|
||||||
|
);
|
||||||
|
Err(e)
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
.retry(retry_policy)
|
.retry(retry_policy)
|
||||||
|
@ -39,13 +49,16 @@ pub async fn validate_key_with_retry(client: &Client, api_host: &Url, key: ApiKe
|
||||||
match result {
|
match result {
|
||||||
Ok(key_result) => key_result,
|
Ok(key_result) => key_result,
|
||||||
Err(_) => {
|
Err(_) => {
|
||||||
eprintln!("Key: {}... -> FAILED after all retries.", &key.as_str()[..10]);
|
eprintln!(
|
||||||
|
"Key: {}... -> FAILED after all retries.",
|
||||||
|
&key.as_str()[..10]
|
||||||
|
);
|
||||||
None
|
None
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn keytest(client: &Client, api_host: &Url, key: &ApiKey) -> Result<KeyStatus> {
|
async fn keytest(client: Client, api_host: &Url, key: &ApiKey) -> Result<KeyStatus> {
|
||||||
const API_PATH: &str = "v1beta/models/gemini-2.0-flash-exp:generateContent";
|
const API_PATH: &str = "v1beta/models/gemini-2.0-flash-exp:generateContent";
|
||||||
let full_url = api_host.join(API_PATH)?;
|
let full_url = api_host.join(API_PATH)?;
|
||||||
|
|
||||||
|
|
|
@ -1,13 +1,13 @@
|
||||||
use anyhow::Result;
|
use anyhow::Result;
|
||||||
use gemini_keychecker::config::{KeyCheckerConfig, client_builder};
|
|
||||||
use gemini_keychecker::adapters::load_keys;
|
use gemini_keychecker::adapters::load_keys;
|
||||||
|
use gemini_keychecker::config::{KeyCheckerConfig, client_builder};
|
||||||
use gemini_keychecker::validation::ValidationService;
|
use gemini_keychecker::validation::ValidationService;
|
||||||
|
|
||||||
/// Main function - orchestrates the key validation process
|
/// Main function - orchestrates the key validation process
|
||||||
#[tokio::main]
|
#[tokio::main]
|
||||||
async fn main() -> Result<()> {
|
async fn main() -> Result<()> {
|
||||||
let config = KeyCheckerConfig::load_config().unwrap();
|
let config = KeyCheckerConfig::load_config().unwrap();
|
||||||
let keys = load_keys(&config.input_path)?;
|
let keys = load_keys(config.input_path().as_path())?;
|
||||||
let client = client_builder(&config)?;
|
let client = client_builder(&config)?;
|
||||||
|
|
||||||
let validation_service = ValidationService::new(config, client);
|
let validation_service = ValidationService::new(config, client);
|
||||||
|
|
|
@ -5,8 +5,8 @@ use reqwest::Client;
|
||||||
use std::{fs, time::Instant};
|
use std::{fs, time::Instant};
|
||||||
use tokio::sync::mpsc;
|
use tokio::sync::mpsc;
|
||||||
|
|
||||||
use crate::config::KeyCheckerConfig;
|
|
||||||
use crate::adapters::write_keys_txt_file;
|
use crate::adapters::write_keys_txt_file;
|
||||||
|
use crate::config::KeyCheckerConfig;
|
||||||
use crate::key_validator::validate_key_with_retry;
|
use crate::key_validator::validate_key_with_retry;
|
||||||
use crate::types::ApiKey;
|
use crate::types::ApiKey;
|
||||||
|
|
||||||
|
@ -42,13 +42,15 @@ impl ValidationService {
|
||||||
|
|
||||||
// Create stream to validate keys concurrently
|
// Create stream to validate keys concurrently
|
||||||
let valid_keys_stream = stream
|
let valid_keys_stream = stream
|
||||||
.map(|key| validate_key_with_retry(&self.client, &self.config.api_host, key))
|
.map(|key| {
|
||||||
.buffer_unordered(self.config.concurrency)
|
validate_key_with_retry(self.client.to_owned(), self.config.api_host(), key)
|
||||||
|
})
|
||||||
|
.buffer_unordered(self.config.concurrency())
|
||||||
.filter_map(|r| async { r });
|
.filter_map(|r| async { r });
|
||||||
pin_mut!(valid_keys_stream);
|
pin_mut!(valid_keys_stream);
|
||||||
|
|
||||||
// Open output file for writing valid keys
|
// Open output file for writing valid keys
|
||||||
let mut output_file = fs::File::create(&self.config.output_path)?;
|
let mut output_file = fs::File::create(&self.config.output_path())?;
|
||||||
|
|
||||||
// Process validated keys and write to output file
|
// Process validated keys and write to output file
|
||||||
while let Some(valid_key) = valid_keys_stream.next().await {
|
while let Some(valid_key) = valid_keys_stream.next().await {
|
||||||
|
|
Loading…
Reference in New Issue