use serde::{Deserialize, Serialize};
use std::fs;
use std::path::Path;
const MAGIC: [u8; 4] = [0x4D, 0x4B, 0x45, 0x59]; const VERSION: u8 = 1;
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct KeyStore {
pub magic: [u8; 4],
pub version: u8,
pub entries: Vec<KeyEntry>,
}
impl Default for KeyStore {
fn default() -> Self {
Self {
magic: MAGIC,
version: VERSION,
entries: Vec::new(),
}
}
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct KeyEntry {
pub service: String,
pub key: String,
pub priority: u8,
pub quota_remaining: Option<f64>,
#[serde(default)]
pub last_error: Option<String>,
#[serde(default)]
pub disabled: bool,
pub created_at: i64,
}
impl KeyStore {
pub fn load<P: AsRef<Path>>(path: P) -> Result<Self, Box<dyn std::error::Error>> {
if !path.as_ref().exists() {
return Ok(Self {
magic: MAGIC,
version: VERSION,
entries: Vec::new(),
});
}
let bytes = fs::read(path.as_ref())?;
let store: KeyStore = bincode::deserialize(&bytes)?;
if store.magic != MAGIC {
return Err("Invalid keys.bin magic bytes".into());
}
if store.version != VERSION {
return Err(format!("Unsupported keys.bin version: {}", store.version).into());
}
Ok(store)
}
pub fn save<P: AsRef<Path>>(&self, path: P) -> Result<(), Box<dyn std::error::Error>> {
let bytes = bincode::serialize(self)?;
let tmp = path.as_ref().with_extension("bin.tmp");
fs::write(&tmp, &bytes)?;
fs::rename(&tmp, path.as_ref())?;
Ok(())
}
pub fn get(&self, service: &str) -> Option<&KeyEntry> {
self.entries
.iter()
.filter(|e| e.service == service && !e.disabled)
.min_by_key(|e| e.priority)
}
pub fn get_all(&self, service: &str) -> Vec<&KeyEntry> {
let mut v: Vec<_> = self
.entries
.iter()
.filter(|e| e.service == service && !e.disabled)
.collect();
v.sort_by_key(|e| e.priority);
v
}
pub fn set(&mut self, service: &str, key: String, priority: u8) {
if let Some(existing) = self
.entries
.iter_mut()
.find(|e| e.service == service && e.priority == priority)
{
existing.key = key;
existing.last_error = None;
existing.disabled = false;
existing.created_at = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs() as i64;
return;
}
self.entries.push(KeyEntry {
service: service.to_string(),
key,
priority,
quota_remaining: None,
last_error: None,
disabled: false,
created_at: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs() as i64,
});
}
pub fn remove(&mut self, service: &str, priority: u8) {
self.entries
.retain(|e| !(e.service == service && e.priority == priority));
}
pub fn record_error(&mut self, service: &str, priority: u8, error: &str) {
if let Some(entry) = self
.entries
.iter_mut()
.find(|e| e.service == service && e.priority == priority)
{
entry.last_error = Some(error.to_string());
if error.contains("429") || error.contains("401") || error.contains("403") {
entry.disabled = true;
}
}
}
pub fn record_success(&mut self, service: &str, priority: u8) {
if let Some(entry) = self
.entries
.iter_mut()
.find(|e| e.service == service && e.priority == priority)
{
entry.last_error = None;
entry.disabled = false;
}
}
pub fn reset_all(&mut self) {
for entry in &mut self.entries {
entry.disabled = false;
entry.last_error = None;
}
}
pub fn list(&self) -> Vec<KeyEntryInfo> {
self.entries
.iter()
.map(|e| KeyEntryInfo {
service: e.service.clone(),
priority: e.priority,
quota_remaining: e.quota_remaining,
last_error: e.last_error.clone(),
disabled: e.disabled,
created_at: e.created_at,
key_preview: if e.key.len() > 8 {
format!("{}...", &e.key[..8])
} else {
"***".to_string()
},
})
.collect()
}
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct KeyEntryInfo {
pub service: String,
pub priority: u8,
pub quota_remaining: Option<f64>,
pub last_error: Option<String>,
pub disabled: bool,
pub created_at: i64,
pub key_preview: String,
}
pub fn default_keys_path(output_dir: &str) -> String {
format!("{}/keys.bin", output_dir)
}