feat: refactor key handling to replace ApiKey with GeminiKey across the codebase
parent
35b9d0b0b0
commit
e6d2309d9e
|
@ -1,15 +1,11 @@
|
|||
use crate::adapters::output::write_keys_to_file;
|
||||
use crate::types::GeminiKey;
|
||||
use anyhow::Result;
|
||||
use std::{
|
||||
collections::HashSet,
|
||||
fs,
|
||||
path::Path,
|
||||
str::FromStr,
|
||||
};
|
||||
use crate::types::ApiKey;
|
||||
use std::{collections::HashSet, fs, path::Path, str::FromStr};
|
||||
|
||||
/// 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>> {
|
||||
pub fn load_keys(path: &Path) -> Result<Vec<GeminiKey>> {
|
||||
let keys_txt = fs::read_to_string(path)?;
|
||||
// Use HashSet to automatically deduplicate keys
|
||||
let unique_keys_set: HashSet<&str> = keys_txt
|
||||
|
@ -22,25 +18,19 @@ pub fn load_keys(path: &Path) -> Result<Vec<ApiKey>> {
|
|||
let mut valid_keys_for_backup = Vec::new();
|
||||
|
||||
for key_str in unique_keys_set {
|
||||
match ApiKey::from_str(key_str) {
|
||||
match GeminiKey::from_str(key_str) {
|
||||
Ok(api_key) => {
|
||||
keys.push(api_key.clone());
|
||||
valid_keys_for_backup.push(api_key.as_str().to_string());
|
||||
valid_keys_for_backup.push(api_key.inner.clone());
|
||||
}
|
||||
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) {
|
||||
if let Err(e) = write_keys_to_file(&valid_keys_for_backup, "backup.txt") {
|
||||
eprintln!("Failed to write backup file: {}", e);
|
||||
} else {
|
||||
println!(
|
||||
"Backup file created with {} valid keys",
|
||||
valid_keys_for_backup.len()
|
||||
);
|
||||
}
|
||||
|
||||
Ok(keys)
|
||||
}
|
||||
}
|
||||
|
|
|
@ -1,4 +1,4 @@
|
|||
use crate::types::ApiKey;
|
||||
use crate::types::GeminiKey;
|
||||
use anyhow::Result;
|
||||
use std::{fs, io::Write};
|
||||
use tokio::io::{AsyncWriteExt, BufWriter};
|
||||
|
@ -7,16 +7,16 @@ use toml::Value;
|
|||
// Write valid key to output file
|
||||
pub async fn write_keys_txt_file(
|
||||
file: &mut BufWriter<tokio::fs::File>,
|
||||
key: &ApiKey,
|
||||
key: &GeminiKey,
|
||||
) -> Result<()> {
|
||||
file.write_all(format!("{}\n", key.as_str()).as_bytes()).await?;
|
||||
file.write_all(format!("{}\n", key.as_ref()).as_bytes()).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// Write valid key to output file in Clewdr format
|
||||
pub fn write_keys_clewdr_format(file: &mut fs::File, key: &ApiKey) -> Result<()> {
|
||||
pub fn write_keys_clewdr_format(file: &mut fs::File, key: &GeminiKey) -> Result<()> {
|
||||
let mut table = toml::value::Table::new();
|
||||
table.insert("key".to_string(), Value::String(key.as_str().to_string()));
|
||||
table.insert("key".to_string(), Value::String(key.as_ref().to_string()));
|
||||
|
||||
let gemini_keys = Value::Array(vec![Value::Table(table)]);
|
||||
let mut root = toml::value::Table::new();
|
||||
|
@ -26,3 +26,11 @@ pub fn write_keys_clewdr_format(file: &mut fs::File, key: &ApiKey) -> Result<()>
|
|||
write!(file, "{}", toml_string)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// Write keys to a text file with custom filename
|
||||
pub fn write_keys_to_file(keys: &[String], filename: &str) -> Result<()> {
|
||||
let content = keys.join("\n");
|
||||
fs::write(filename, content)?;
|
||||
println!("File '{}' created with {} keys", filename, keys.len());
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
@ -5,13 +5,13 @@ use serde_json;
|
|||
use tokio::time::Duration;
|
||||
use url::Url;
|
||||
|
||||
use crate::types::{ApiKey, KeyStatus};
|
||||
use crate::types::{GeminiKey, KeyStatus};
|
||||
|
||||
pub async fn validate_key_with_retry(
|
||||
client: Client,
|
||||
api_host: Url,
|
||||
key: ApiKey,
|
||||
) -> Option<ApiKey> {
|
||||
key: GeminiKey,
|
||||
) -> Option<GeminiKey> {
|
||||
let retry_policy = ExponentialBuilder::default()
|
||||
.with_max_times(3)
|
||||
.with_min_delay(Duration::from_secs(3))
|
||||
|
@ -19,17 +19,17 @@ pub async fn validate_key_with_retry(
|
|||
|
||||
let result = (async || match keytest(client.to_owned(), &api_host, &key).await {
|
||||
Ok(KeyStatus::Valid) => {
|
||||
println!("Key: {}... -> SUCCESS", &key.as_str()[..10]);
|
||||
println!("Key: {}... -> SUCCESS", &key.as_ref()[..10]);
|
||||
Ok(Some(key.clone()))
|
||||
}
|
||||
Ok(KeyStatus::Invalid) => {
|
||||
println!("Key: {}... -> INVALID (Forbidden)", &key.as_str()[..10]);
|
||||
println!("Key: {}... -> INVALID (Forbidden)", &key.as_ref()[..10]);
|
||||
Ok(None)
|
||||
}
|
||||
Ok(KeyStatus::Retryable(reason)) => {
|
||||
eprintln!(
|
||||
"Key: {}... -> RETRYABLE (Reason: {})",
|
||||
&key.as_str()[..10],
|
||||
&key.as_ref()[..10],
|
||||
reason
|
||||
);
|
||||
Err(anyhow::anyhow!("Retryable error: {}", reason))
|
||||
|
@ -37,7 +37,7 @@ pub async fn validate_key_with_retry(
|
|||
Err(e) => {
|
||||
eprintln!(
|
||||
"Key: {}... -> NETWORK ERROR (Reason: {})",
|
||||
&key.as_str()[..10],
|
||||
&key.as_ref()[..10],
|
||||
e
|
||||
);
|
||||
Err(e)
|
||||
|
@ -51,14 +51,14 @@ pub async fn validate_key_with_retry(
|
|||
Err(_) => {
|
||||
eprintln!(
|
||||
"Key: {}... -> FAILED after all retries.",
|
||||
&key.as_str()[..10]
|
||||
&key.as_ref()[..10]
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn keytest(client: Client, api_host: &Url, key: &ApiKey) -> Result<KeyStatus> {
|
||||
async fn keytest(client: Client, api_host: &Url, key: &GeminiKey) -> Result<KeyStatus> {
|
||||
const API_PATH: &str = "v1beta/models/gemini-2.0-flash-exp:generateContent";
|
||||
let full_url = api_host.join(API_PATH)?;
|
||||
|
||||
|
@ -77,7 +77,7 @@ async fn keytest(client: Client, api_host: &Url, key: &ApiKey) -> Result<KeyStat
|
|||
let response = client
|
||||
.post(full_url)
|
||||
.header("Content-Type", "application/json")
|
||||
.header("X-goog-api-key", key.as_str())
|
||||
.header("X-goog-api-key", key.as_ref())
|
||||
.json(&request_body)
|
||||
.send()
|
||||
.await?;
|
||||
|
|
34
src/types.rs
34
src/types.rs
|
@ -10,36 +10,22 @@ pub enum KeyStatus {
|
|||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ApiKey {
|
||||
inner: String,
|
||||
pub struct GeminiKey {
|
||||
pub inner: String,
|
||||
}
|
||||
|
||||
impl ApiKey {
|
||||
pub fn as_str(&self) -> &str {
|
||||
impl AsRef<str> for GeminiKey {
|
||||
fn as_ref(&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;
|
||||
impl FromStr for GeminiKey {
|
||||
type Err = &'static str;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
static RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"^AIzaSy.{33}$").unwrap());
|
||||
static RE: LazyLock<Regex> =
|
||||
LazyLock::new(|| Regex::new(r"^AIzaSy[A-Za-z0-9_-]{33}$").unwrap());
|
||||
|
||||
let cleaned = s.trim();
|
||||
|
||||
|
@ -48,9 +34,7 @@ impl FromStr for ApiKey {
|
|||
inner: cleaned.to_string(),
|
||||
})
|
||||
} else {
|
||||
Err(KeyValidationError::InvalidFormat(
|
||||
"Google API key must start with 'AIzaSy' followed by 33 characters".to_string(),
|
||||
))
|
||||
Err("Invalid Google API key format")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
@ -8,7 +8,7 @@ use tokio::{fs, io::AsyncWriteExt, sync::mpsc};
|
|||
use crate::adapters::write_keys_txt_file;
|
||||
use crate::config::KeyCheckerConfig;
|
||||
use crate::key_validator::validate_key_with_retry;
|
||||
use crate::types::ApiKey;
|
||||
use crate::types::GeminiKey;
|
||||
|
||||
pub struct ValidationService {
|
||||
config: KeyCheckerConfig,
|
||||
|
@ -20,11 +20,11 @@ impl ValidationService {
|
|||
Self { config, client }
|
||||
}
|
||||
|
||||
pub async fn validate_keys(&self, keys: Vec<ApiKey>) -> Result<()> {
|
||||
pub async fn validate_keys(&self, keys: Vec<GeminiKey>) -> Result<()> {
|
||||
let start_time = Instant::now();
|
||||
|
||||
// Create channel for streaming keys from producer to consumer
|
||||
let (tx, mut rx) = mpsc::unbounded_channel::<ApiKey>();
|
||||
let (tx, mut rx) = mpsc::unbounded_channel::<GeminiKey>();
|
||||
let stream = stream! {
|
||||
while let Some(item) = rx.recv().await {
|
||||
yield item;
|
||||
|
@ -53,7 +53,7 @@ impl ValidationService {
|
|||
|
||||
// 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());
|
||||
println!("Valid key found: {}", valid_key.as_ref());
|
||||
if let Err(e) = write_keys_txt_file(&mut buffer_writer, &valid_key).await {
|
||||
eprintln!("Failed to write key to output file: {}", e);
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue