use crate::models::{Account, TokenData};
use crate::modules::auth::account;
use crate::utils::protobuf;
use base64::{engine::general_purpose, Engine as _};
use serde_json::Value;
use std::fs;
use std::path::PathBuf;
pub async fn import_from_v1() -> Result<Vec<Account>, String> {
use crate::modules::auth::oauth;
let home = dirs::home_dir().ok_or("Failed to get home directory")?;
let v1_dir = home.join(".antigravity-agent");
let mut imported_accounts = Vec::new();
let index_files = vec!["antigravity_accounts.json", "accounts.json"];
let mut found_index = false;
for index_filename in index_files {
let v1_accounts_path = v1_dir.join(index_filename);
if !v1_accounts_path.exists() {
continue;
}
found_index = true;
crate::modules::system::logger::log_info(&format!(
"V1 data discovered: {:?}",
v1_accounts_path
));
let content = match fs::read_to_string(&v1_accounts_path) {
Ok(c) => c,
Err(e) => {
crate::modules::system::logger::log_warn(&format!("Failed to read index: {}", e));
continue;
}
};
let v1_index: Value = match serde_json::from_str(&content) {
Ok(v) => v,
Err(e) => {
crate::modules::system::logger::log_warn(&format!(
"Failed to parse index JSON: {}",
e
));
continue;
}
};
let accounts_map = if let Some(map) = v1_index.as_object() {
if let Some(accounts) = map.get("accounts").and_then(|v| v.as_object()) {
accounts
} else {
map
}
} else {
continue;
};
for (id, acc_info) in accounts_map {
let email_placeholder = acc_info
.get("email")
.and_then(|v| v.as_str())
.unwrap_or("Unknown")
.to_string();
if !acc_info.is_object() {
continue;
}
let backup_file_str = acc_info.get("backup_file").and_then(|v| v.as_str());
let data_file_str = acc_info.get("data_file").and_then(|v| v.as_str());
let target_file = backup_file_str.or(data_file_str);
if target_file.is_none() {
crate::modules::system::logger::log_warn(&format!(
"Account {} ({}) missing data file path",
id, email_placeholder
));
continue;
}
let mut backup_path = PathBuf::from(target_file.unwrap());
if !backup_path.exists() {
backup_path = v1_dir.join(backup_path.file_name().unwrap_or_default());
}
if !backup_path.exists() {
let file_name = backup_path.file_name().unwrap_or_default();
let try_backups = v1_dir.join("backups").join(file_name);
if try_backups.exists() {
backup_path = try_backups;
} else {
let try_accounts = v1_dir.join("accounts").join(file_name);
if try_accounts.exists() {
backup_path = try_accounts;
}
}
}
if !backup_path.exists() {
crate::modules::system::logger::log_warn(&format!(
"Account {} ({}) backup file not found: {:?}",
id, email_placeholder, backup_path
));
continue;
}
if let Ok(backup_content) = fs::read_to_string(&backup_path) {
if let Ok(backup_json) = serde_json::from_str::<Value>(&backup_content) {
let mut refresh_token_opt = None;
if let Some(token_data) = backup_json.get("token") {
if let Some(rt) = token_data.get("refresh_token").and_then(|v| v.as_str()) {
refresh_token_opt = Some(rt.to_string());
}
}
if refresh_token_opt.is_none() {
if let Some(state_b64) = backup_json
.get("jetskiStateSync.agentManagerInitState")
.and_then(|v| v.as_str())
{
if let Ok(blob) = general_purpose::STANDARD.decode(state_b64) {
if let Ok(Some(oauth_data)) = protobuf::find_field(&blob, 6) {
if let Ok(Some(refresh_bytes)) =
protobuf::find_field(&oauth_data, 3)
{
if let Ok(rt) = String::from_utf8(refresh_bytes) {
refresh_token_opt = Some(rt);
}
}
}
}
}
}
if let Some(refresh_token) = refresh_token_opt {
crate::modules::system::logger::log_info(&format!(
"Importing account: {}",
email_placeholder
));
let (token_resp, identity) = match oauth::refresh_and_verify_identity(
&refresh_token,
None,
)
.await
{
Ok(result) => result,
Err(e) => {
crate::modules::system::logger::log_warn(&format!(
"Skipping import for {}: refresh/identity verification failed: {}",
email_placeholder, e
));
continue;
}
};
let crate::modules::auth::oauth::VerifiedIdentity {
email,
name,
google_sub,
..
} = identity;
let token_data = TokenData::new(
token_resp.access_token,
refresh_token,
token_resp.expires_in,
Some(email.clone()),
None,
None,
);
match account::upsert_account(email.clone(), name, token_data, google_sub) {
Ok(acc) => {
crate::modules::system::logger::log_info(&format!(
"Import successful: {}",
email
));
imported_accounts.push(acc);
}
Err(e) => crate::modules::system::logger::log_error(&format!(
"Import save failed {}: {}",
email, e
)),
}
} else {
crate::modules::system::logger::log_warn(&format!(
"Account {} data file missing Refresh Token",
email_placeholder
));
}
}
}
}
}
if !found_index {
return Err("V1 account data file not found".to_string());
}
Ok(imported_accounts)
}
pub async fn import_from_custom_db_path(path_str: String) -> Result<Account, String> {
use crate::modules::auth::oauth;
let path = PathBuf::from(path_str);
if !path.exists() {
return Err(format!("File does not exist: {:?}", path));
}
let refresh_token = extract_refresh_token_from_file(&path)?;
crate::modules::system::logger::log_info(
"Refreshing token and verifying Google identity for import...",
);
let (token_resp, identity) = oauth::refresh_and_verify_identity(&refresh_token, None).await?;
let crate::modules::auth::oauth::VerifiedIdentity {
email,
name,
google_sub,
..
} = identity;
crate::modules::system::logger::log_info(&format!(
"Successfully retrieved account info: {}",
email
));
let token_data = TokenData::new(
token_resp.access_token,
refresh_token,
token_resp.expires_in,
Some(email.clone()),
None,
None,
);
account::upsert_account(email.clone(), name, token_data, google_sub)
}
pub fn extract_refresh_token_from_file(db_path: &PathBuf) -> Result<String, String> {
if !db_path.exists() {
return Err(format!("Database file not found: {:?}", db_path));
}
let conn = rusqlite::Connection::open(db_path)
.map_err(|e| format!("Failed to open database: {}", e))?;
let current_data: String = conn
.query_row(
"SELECT value FROM ItemTable WHERE key = ?",
["jetskiStateSync.agentManagerInitState"],
|row| row.get(0),
)
.map_err(|_| {
"Login state data not found (jetskiStateSync.agentManagerInitState)".to_string()
})?;
let blob = general_purpose::STANDARD
.decode(¤t_data)
.map_err(|e| format!("Base64 decoding failed: {}", e))?;
let oauth_data = protobuf::find_field(&blob, 6)
.map_err(|e| format!("Protobuf parsing failed: {}", e))?
.ok_or("OAuth data not found (Field 6)")?;
let refresh_bytes = protobuf::find_field(&oauth_data, 3)
.map_err(|e| format!("OAuth data parsing failed: {}", e))?
.ok_or("Refresh Token not included in data (Field 3)")?;
String::from_utf8(refresh_bytes).map_err(|_| "Refresh Token is not UTF-8 encoded".to_string())
}