From e6d2309d9e2c2434b0a118d6ae40fdac87f6b17e Mon Sep 17 00:00:00 2001 From: Yoo1tic <137816438+Yoo1tic@users.noreply.github.com> Date: Fri, 18 Jul 2025 02:51:21 +0800 Subject: [PATCH] feat: refactor key handling to replace ApiKey with GeminiKey across the codebase --- src/adapters/input/local.rs | 26 ++++++++------------------ src/adapters/output/local.rs | 18 +++++++++++++----- src/key_validator.rs | 20 ++++++++++---------- src/types.rs | 34 +++++++++------------------------- src/validation.rs | 8 ++++---- 5 files changed, 44 insertions(+), 62 deletions(-) diff --git a/src/adapters/input/local.rs b/src/adapters/input/local.rs index 91f1a5f..1f52f9a 100644 --- a/src/adapters/input/local.rs +++ b/src/adapters/input/local.rs @@ -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> { +pub fn load_keys(path: &Path) -> Result> { 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> { 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) -} \ No newline at end of file +} diff --git a/src/adapters/output/local.rs b/src/adapters/output/local.rs index 017469b..bf9d62c 100644 --- a/src/adapters/output/local.rs +++ b/src/adapters/output/local.rs @@ -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, - 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(()) +} diff --git a/src/key_validator.rs b/src/key_validator.rs index 39220d2..41f56e7 100644 --- a/src/key_validator.rs +++ b/src/key_validator.rs @@ -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 { + key: GeminiKey, +) -> Option { 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 { +async fn keytest(client: Client, api_host: &Url, key: &GeminiKey) -> Result { 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 &str { +impl AsRef 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 { - static RE: LazyLock = LazyLock::new(|| Regex::new(r"^AIzaSy.{33}$").unwrap()); + static RE: LazyLock = + 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") } } } diff --git a/src/validation.rs b/src/validation.rs index 8a12473..07b7cdd 100644 --- a/src/validation.rs +++ b/src/validation.rs @@ -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) -> Result<()> { + pub async fn validate_keys(&self, keys: Vec) -> Result<()> { let start_time = Instant::now(); // Create channel for streaming keys from producer to consumer - let (tx, mut rx) = mpsc::unbounded_channel::(); + let (tx, mut rx) = mpsc::unbounded_channel::(); 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); }