mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 16:06:08 +00:00
feat: add admin account export
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
Codex
@@ -215,6 +215,19 @@ impl AccountRepository {
|
||||
Ok(Page { items, next_cursor })
|
||||
}
|
||||
|
||||
pub async fn list_all(&self) -> AccountRepositoryResult<Vec<StoredAccount>> {
|
||||
let rows = sqlx::query(
|
||||
"select id, email, account_id, user_id, label, plan_type, access_token_cipher, refresh_token_cipher, access_token_expires_at, status, added_at, updated_at from accounts order by added_at desc, id desc",
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
let mut accounts = Vec::with_capacity(rows.len());
|
||||
for row in rows {
|
||||
accounts.push(self.account_from_row(&row)?);
|
||||
}
|
||||
Ok(accounts)
|
||||
}
|
||||
|
||||
pub async fn list_metadata(
|
||||
&self,
|
||||
cursor: Option<String>,
|
||||
|
||||
+5
-3
@@ -10,9 +10,10 @@ use crate::{
|
||||
admin::{
|
||||
accounts, api_keys, batch_delete_accounts, batch_delete_api_keys,
|
||||
batch_update_account_status, create_api_key, delete_account, delete_account_cookies,
|
||||
delete_api_key, get_account_cookies, import_accounts, login, logs, refresh_models,
|
||||
set_account_cookies, settings, update_account_label, update_account_status,
|
||||
update_api_key_label, update_api_key_status, usage_stats, usage_stats_summary,
|
||||
delete_api_key, export_accounts, get_account_cookies, import_accounts, login, logs,
|
||||
refresh_models, set_account_cookies, settings, update_account_label,
|
||||
update_account_status, update_api_key_label, update_api_key_status, usage_stats,
|
||||
usage_stats_summary,
|
||||
},
|
||||
health::health,
|
||||
middleware::attach_request_id,
|
||||
@@ -46,6 +47,7 @@ pub fn build_router(state: AppState) -> Router {
|
||||
"/admin/accounts/batch-status",
|
||||
post(batch_update_account_status),
|
||||
)
|
||||
.route("/admin/accounts/export", get(export_accounts))
|
||||
.route(
|
||||
"/admin/accounts/{account_id}/cookies",
|
||||
get(get_account_cookies)
|
||||
|
||||
+197
-2
@@ -7,10 +7,10 @@ use axum::{
|
||||
response::{IntoResponse, Response},
|
||||
Extension, Json,
|
||||
};
|
||||
use chrono::{Duration, Utc};
|
||||
use chrono::{Duration, SecondsFormat, Utc};
|
||||
use secrecy::{ExposeSecret, SecretString};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use serde_json::{json, Map, Value};
|
||||
use sqlx::Row;
|
||||
use uuid::Uuid;
|
||||
|
||||
@@ -138,6 +138,13 @@ pub struct AccountsQuery {
|
||||
pub limit: Option<u32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountExportQuery {
|
||||
pub ids: Option<String>,
|
||||
pub format: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ApiKeysQuery {
|
||||
@@ -229,6 +236,12 @@ pub struct AccountImportData {
|
||||
pub source_format: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum AccountExportFormat {
|
||||
Native,
|
||||
Sub2Api,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct UpdateAccountLabelRequest {
|
||||
@@ -992,6 +1005,84 @@ pub async fn accounts(
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn export_accounts(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
Query(query): Query<AccountExportQuery>,
|
||||
) -> Response {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
let format = match parse_account_export_format(query.format.as_deref()) {
|
||||
Ok(format) => format,
|
||||
Err(message) => {
|
||||
return AdminResponse::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
AdminEnvelope::new(40001, message, (), request_id),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
};
|
||||
let Some(pool) = state.db() else {
|
||||
return AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(50001, "Database is not initialized", (), request_id),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
if let Err(response) = require_admin_session(pool, &headers, &request_id).await {
|
||||
return response;
|
||||
}
|
||||
let Some(repo) = state.account_repository() else {
|
||||
return AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(
|
||||
50001,
|
||||
"Account repository is not initialized",
|
||||
(),
|
||||
request_id,
|
||||
),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
|
||||
let ids = account_export_ids(query.ids.as_deref());
|
||||
let accounts = if ids.is_empty() {
|
||||
match repo.list_all().await {
|
||||
Ok(accounts) => accounts,
|
||||
Err(_) => {
|
||||
return AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(50001, "Failed to export accounts", (), request_id),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let mut accounts = Vec::with_capacity(ids.len());
|
||||
for id in ids {
|
||||
match repo.get(&id).await {
|
||||
Ok(Some(account)) => accounts.push(account),
|
||||
Ok(None) => {}
|
||||
Err(_) => {
|
||||
return AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(50001, "Failed to export accounts", (), request_id),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
}
|
||||
}
|
||||
accounts
|
||||
};
|
||||
|
||||
// 账号导出会返回可重新导入的 OAuth token；只允许 admin session 访问，不写入日志。
|
||||
let data = match format {
|
||||
AccountExportFormat::Native => native_account_export(accounts),
|
||||
AccountExportFormat::Sub2Api => sub2api_account_export(accounts),
|
||||
};
|
||||
AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data, request_id)).into_response()
|
||||
}
|
||||
|
||||
pub async fn update_account_label(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
@@ -2226,6 +2317,110 @@ fn normalized_account_id(id: Option<String>) -> String {
|
||||
.unwrap_or_else(|| format!("acct_{}", Uuid::new_v4().simple()))
|
||||
}
|
||||
|
||||
fn account_export_ids(value: Option<&str>) -> Vec<String> {
|
||||
value
|
||||
.into_iter()
|
||||
.flat_map(|ids| ids.split(','))
|
||||
.filter_map(|id| {
|
||||
let id = id.trim();
|
||||
(!id.is_empty()).then(|| id.to_string())
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn parse_account_export_format(value: Option<&str>) -> Result<AccountExportFormat, &'static str> {
|
||||
match value.unwrap_or("native") {
|
||||
"" | "native" | "full" => Ok(AccountExportFormat::Native),
|
||||
"sub2api" => Ok(AccountExportFormat::Sub2Api),
|
||||
_ => Err("Unsupported account export format"),
|
||||
}
|
||||
}
|
||||
|
||||
fn native_account_export(accounts: Vec<StoredAccount>) -> Value {
|
||||
json!({
|
||||
"sourceFormat": "native",
|
||||
"accounts": accounts.into_iter().map(native_export_account).collect::<Vec<_>>(),
|
||||
})
|
||||
}
|
||||
|
||||
fn native_export_account(account: StoredAccount) -> Value {
|
||||
json!({
|
||||
"id": account.id,
|
||||
"email": account.email,
|
||||
"accountId": account.account_id,
|
||||
"userId": account.user_id,
|
||||
"label": account.label,
|
||||
"planType": account.plan_type,
|
||||
"token": account.access_token.expose_secret(),
|
||||
"refreshToken": account
|
||||
.refresh_token
|
||||
.as_ref()
|
||||
.map(|token| token.expose_secret().to_string()),
|
||||
"status": account_status_value(account.status),
|
||||
"accessTokenExpiresAt": account.access_token_expires_at.map(|value| value.to_rfc3339()),
|
||||
"addedAt": account.added_at.to_rfc3339(),
|
||||
"updatedAt": account.updated_at.to_rfc3339(),
|
||||
})
|
||||
}
|
||||
|
||||
fn sub2api_account_export(accounts: Vec<StoredAccount>) -> Value {
|
||||
json!({
|
||||
"exported_at": Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true),
|
||||
"proxies": [],
|
||||
"accounts": accounts.into_iter().map(sub2api_export_account).collect::<Vec<_>>(),
|
||||
"type": "sub2api-data",
|
||||
"version": 1,
|
||||
})
|
||||
}
|
||||
|
||||
fn sub2api_export_account(account: StoredAccount) -> Value {
|
||||
let mut credentials = Map::new();
|
||||
credentials.insert(
|
||||
"access_token".to_string(),
|
||||
Value::String(account.access_token.expose_secret().to_string()),
|
||||
);
|
||||
insert_optional_string(
|
||||
&mut credentials,
|
||||
"refresh_token",
|
||||
account
|
||||
.refresh_token
|
||||
.as_ref()
|
||||
.map(|token| token.expose_secret().to_string()),
|
||||
);
|
||||
insert_optional_string(&mut credentials, "email", account.email.clone());
|
||||
insert_optional_string(&mut credentials, "chatgpt_account_id", account.account_id);
|
||||
insert_optional_string(&mut credentials, "chatgpt_user_id", account.user_id);
|
||||
insert_optional_string(&mut credentials, "plan_type", account.plan_type);
|
||||
insert_optional_string(
|
||||
&mut credentials,
|
||||
"expires_at",
|
||||
account
|
||||
.access_token_expires_at
|
||||
.map(|value| value.to_rfc3339()),
|
||||
);
|
||||
|
||||
json!({
|
||||
"name": account
|
||||
.label
|
||||
.as_deref()
|
||||
.filter(|label| !label.trim().is_empty())
|
||||
.map(str::to_string)
|
||||
.or_else(|| account.email.clone())
|
||||
.unwrap_or_else(|| account.id.clone()),
|
||||
"platform": "openai",
|
||||
"type": "oauth",
|
||||
"credentials": credentials,
|
||||
"concurrency": 0,
|
||||
"priority": 0,
|
||||
})
|
||||
}
|
||||
|
||||
fn insert_optional_string(map: &mut Map<String, Value>, key: &str, value: Option<String>) {
|
||||
if let Some(value) = value.filter(|value| !value.is_empty()) {
|
||||
map.insert(key.to_string(), Value::String(value));
|
||||
}
|
||||
}
|
||||
|
||||
fn distinct_active_plan_accounts(accounts: Vec<Account>) -> Vec<(String, Account)> {
|
||||
let mut by_plan = std::collections::BTreeMap::new();
|
||||
for account in accounts {
|
||||
|
||||
@@ -824,6 +824,112 @@ async fn admin_account_cookies_should_require_existing_account_and_non_empty_coo
|
||||
assert_eq!(invalid.status(), StatusCode::BAD_REQUEST);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_accounts_export_should_return_native_accounts_with_tokens_and_filter_ids() {
|
||||
let (app, _state, _pool, _dir) =
|
||||
admin_accounts_test_app("admin-account-export-native.sqlite", 24).await;
|
||||
import_test_account(&app, "session_1", "acct_export_a").await;
|
||||
import_test_account(&app, "session_1", "acct_export_b").await;
|
||||
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/admin/accounts/export?ids=acct_export_a")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = response_json(response).await;
|
||||
assert_eq!(body["data"]["sourceFormat"], "native");
|
||||
assert_eq!(body["data"]["accounts"].as_array().unwrap().len(), 1);
|
||||
assert_eq!(body["data"]["accounts"][0]["id"], "acct_export_a");
|
||||
assert_eq!(body["data"]["accounts"][0]["token"], "access-acct_export_a");
|
||||
assert_eq!(
|
||||
body["data"]["accounts"][0]["refreshToken"],
|
||||
"refresh-acct_export_a"
|
||||
);
|
||||
assert!(body["data"]["accounts"][0].get("proxyApiKey").is_none());
|
||||
|
||||
let invalid = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/admin/accounts/export?format=proxy")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(invalid.status(), StatusCode::BAD_REQUEST);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_accounts_export_should_return_sub2api_openai_oauth_payload_without_proxy_fields() {
|
||||
let (app, _state, _pool, _dir) =
|
||||
admin_accounts_test_app("admin-account-export-sub2api.sqlite", 25).await;
|
||||
import_test_account(&app, "session_1", "acct_export_sub2api").await;
|
||||
|
||||
let label_response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("PATCH")
|
||||
.uri("/admin/accounts/acct_export_sub2api/label")
|
||||
.header("content-type", "application/json")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.body(Body::from(r#"{"label":"Sub2API Export"}"#))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(label_response.status(), StatusCode::OK);
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/admin/accounts/export?format=sub2api")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = response_json(response).await;
|
||||
assert_eq!(body["data"]["type"], "sub2api-data");
|
||||
assert_eq!(body["data"]["version"], 1);
|
||||
assert_eq!(body["data"]["proxies"], json!([]));
|
||||
assert_eq!(body["data"]["accounts"].as_array().unwrap().len(), 1);
|
||||
assert_eq!(body["data"]["accounts"][0]["name"], "Sub2API Export");
|
||||
assert_eq!(body["data"]["accounts"][0]["platform"], "openai");
|
||||
assert_eq!(body["data"]["accounts"][0]["type"], "oauth");
|
||||
assert_eq!(
|
||||
body["data"]["accounts"][0]["credentials"]["access_token"],
|
||||
"access-acct_export_sub2api"
|
||||
);
|
||||
assert_eq!(
|
||||
body["data"]["accounts"][0]["credentials"]["refresh_token"],
|
||||
"refresh-acct_export_sub2api"
|
||||
);
|
||||
assert_eq!(
|
||||
body["data"]["accounts"][0]["credentials"]["email"],
|
||||
"acct_export_sub2api@example.com"
|
||||
);
|
||||
assert_eq!(
|
||||
body["data"]["accounts"][0]["credentials"]["plan_type"],
|
||||
"plus"
|
||||
);
|
||||
assert!(body["data"]["accounts"][0].get("proxy").is_none());
|
||||
assert!(body["data"]["accounts"][0].get("proxyUrl").is_none());
|
||||
}
|
||||
|
||||
async fn admin_accounts_test_app(
|
||||
db_name: &str,
|
||||
key_byte: u8,
|
||||
|
||||
Reference in New Issue
Block a user