Team Ai
Apppublic

devin15/cursor2api-rust

sourceHugging Faceupdated 2y agoView on Hugging Face
0likes
token.rs276 linesDownload Raw Back to route
1use crate::{2    app::{3        constant::{4            AUTHORIZATION_BEARER_PREFIX, CONTENT_TYPE_TEXT_HTML_WITH_UTF8,5            CONTENT_TYPE_TEXT_PLAIN_WITH_UTF8, ROUTE_TOKENINFO_PATH,6        },7        lazy::{AUTH_TOKEN, TOKEN_FILE, TOKEN_LIST_FILE},8        model::{AppConfig, AppState, PageContent, TokenUpdateRequest},9    },10    common::{11        models::{ApiStatus, NormalResponseNoData},12        utils::{13            extract_time, extract_time_ks, extract_user_id, generate_checksum_with_default, generate_checksum_with_repair, generate_hash, generate_timestamp_header, load_tokens, validate_token_and_checksum14        },15    },16};17use axum::{18    extract::{Query, State},19    http::{20        header::{AUTHORIZATION, CONTENT_TYPE},21        HeaderMap,22    },23    response::{IntoResponse, Response},24    Json,25};26use reqwest::StatusCode;27use serde::{Deserialize, Serialize};28use std::sync::Arc;29use tokio::sync::Mutex;30 31pub async fn handle_get_hash() -> Response {32    let hash = generate_hash();33 34    let mut headers = HeaderMap::new();35    headers.insert(36        CONTENT_TYPE,37        CONTENT_TYPE_TEXT_PLAIN_WITH_UTF8.parse().unwrap(),38    );39 40    (headers, hash).into_response()41}42 43#[derive(Deserialize)]44pub struct ChecksumQuery {45    #[serde(default)]46    pub checksum: Option<String>,47}48 49pub async fn handle_get_checksum(Query(query): Query<ChecksumQuery>) -> Response {50    let checksum = match query.checksum {51        None => generate_checksum_with_default(),52        Some(checksum) => generate_checksum_with_repair(&checksum),53    };54 55    let mut headers = HeaderMap::new();56    headers.insert(57        CONTENT_TYPE,58        CONTENT_TYPE_TEXT_PLAIN_WITH_UTF8.parse().unwrap(),59    );60 61    (headers, checksum).into_response()62}63 64pub async fn handle_get_timestamp_header() -> Response {65    let timestamp_header = generate_timestamp_header();66 67    let mut headers = HeaderMap::new();68    headers.insert(69        CONTENT_TYPE,70        CONTENT_TYPE_TEXT_PLAIN_WITH_UTF8.parse().unwrap(),71    );72 73    (headers, timestamp_header).into_response()74}75 76// 更新 TokenInfo 处理77pub async fn handle_update_tokeninfo(78    State(state): State<Arc<Mutex<AppState>>>,79) -> Json<NormalResponseNoData> {80    // 重新加载 tokens81    let token_infos = load_tokens();82 83    // 更新应用状态84    {85        let mut state = state.lock().await;86        state.token_infos = token_infos;87    }88 89    Json(NormalResponseNoData {90        status: ApiStatus::Success,91        message: Some("Token list has been reloaded".to_string()),92    })93}94 95// 获取 TokenInfo 处理96pub async fn handle_get_tokeninfo(97    headers: HeaderMap,98) -> Result<Json<TokenInfoResponse>, StatusCode> {99    // 验证 AUTH_TOKEN100    let auth_header = headers101        .get(AUTHORIZATION)102        .and_then(|h| h.to_str().ok())103        .and_then(|h| h.strip_prefix(AUTHORIZATION_BEARER_PREFIX))104        .ok_or(StatusCode::UNAUTHORIZED)?;105 106    if auth_header != AUTH_TOKEN.as_str() {107        return Err(StatusCode::UNAUTHORIZED);108    }109 110    let token_file = TOKEN_FILE.as_str();111    let token_list_file = TOKEN_LIST_FILE.as_str();112 113    // 读取文件内容114    let tokens = std::fs::read_to_string(&token_file).unwrap_or_else(|_| String::new());115    let token_list = std::fs::read_to_string(&token_list_file).unwrap_or_else(|_| String::new());116 117    // 获取 tokens_count118    let tokens_count = {119        let token_file_content = std::fs::read_to_string(&token_file).unwrap_or_else(|_| String::new());120        token_file_content.lines().count()121    };122 123    Ok(Json(TokenInfoResponse {124        status: ApiStatus::Success,125        token_file: token_file.to_string(),126        token_list_file: token_list_file.to_string(),127        tokens: Some(tokens),128        tokens_count: Some(tokens_count),129        token_list: Some(token_list),130        message: None,131    }))132}133 134#[derive(Serialize)]135pub struct TokenInfoResponse {136    pub status: ApiStatus,137    pub token_file: String,138    pub token_list_file: String,139    #[serde(skip_serializing_if = "Option::is_none")]140    pub tokens: Option<String>,141    #[serde(skip_serializing_if = "Option::is_none")]142    pub tokens_count: Option<usize>,143    #[serde(skip_serializing_if = "Option::is_none")]144    pub token_list: Option<String>,145    #[serde(skip_serializing_if = "Option::is_none")]146    pub message: Option<String>,147}148 149pub async fn handle_update_tokeninfo_post(150    State(state): State<Arc<Mutex<AppState>>>,151    headers: HeaderMap,152    Json(request): Json<TokenUpdateRequest>,153) -> Result<Json<TokenInfoResponse>, StatusCode> {154    // 验证 AUTH_TOKEN155    let auth_header = headers156        .get(AUTHORIZATION)157        .and_then(|h| h.to_str().ok())158        .and_then(|h| h.strip_prefix(AUTHORIZATION_BEARER_PREFIX))159        .ok_or(StatusCode::UNAUTHORIZED)?;160 161    if auth_header != AUTH_TOKEN.as_str() {162        return Err(StatusCode::UNAUTHORIZED);163    }164 165    let token_file = TOKEN_FILE.as_str();166    let token_list_file = TOKEN_LIST_FILE.as_str();167 168    // 写入文件169    std::fs::write(&token_file, &request.tokens).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;170 171    if let Some(token_list) = &request.token_list {172        std::fs::write(&token_list_file, token_list)173            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;174    }175 176    // 重新加载 tokens177    let token_infos = load_tokens();178    let token_infos_len = token_infos.len();179 180    // 更新应用状态181    {182        let mut state = state.lock().await;183        state.token_infos = token_infos;184    }185 186    Ok(Json(TokenInfoResponse {187        status: ApiStatus::Success,188        token_file: token_file.to_string(),189        token_list_file: token_list_file.to_string(),190        tokens: None,191        tokens_count: Some(token_infos_len),192        token_list: None,193        message: Some("Token files have been updated and reloaded".to_string()),194    }))195}196 197pub async fn handle_tokeninfo_page() -> impl IntoResponse {198    match AppConfig::get_page_content(ROUTE_TOKENINFO_PATH).unwrap_or_default() {199        PageContent::Default => Response::builder()200            .header(CONTENT_TYPE, CONTENT_TYPE_TEXT_HTML_WITH_UTF8)201            .body(include_str!("../../../static/tokeninfo.min.html").to_string())202            .unwrap(),203        PageContent::Text(content) => Response::builder()204            .header(CONTENT_TYPE, CONTENT_TYPE_TEXT_PLAIN_WITH_UTF8)205            .body(content.clone())206            .unwrap(),207        PageContent::Html(content) => Response::builder()208            .header(CONTENT_TYPE, CONTENT_TYPE_TEXT_HTML_WITH_UTF8)209            .body(content.clone())210            .unwrap(),211    }212}213 214#[derive(Deserialize)]215pub struct TokenRequest {216    pub token: Option<String>,217}218 219#[derive(Serialize)]220pub struct BasicCalibrationResponse {221    pub status: ApiStatus,222    pub message: Option<String>,223    #[serde(skip_serializing_if = "Option::is_none")]224    pub user_id: Option<String>,225    #[serde(skip_serializing_if = "Option::is_none")]226    pub create_at: Option<String>,227    #[serde(skip_serializing_if = "Option::is_none")]228    pub checksum_time: Option<u64>,229}230 231pub async fn handle_basic_calibration(232    Json(request): Json<TokenRequest>,233) -> Json<BasicCalibrationResponse> {234    // 从请求头中获取并验证 auth token235    let auth_token = match request.token {236        Some(token) => token,237        None => {238            return Json(BasicCalibrationResponse {239                status: ApiStatus::Error,240                message: Some("未提供授权令牌".to_string()),241                user_id: None,242                create_at: None,243                checksum_time: None,244            })245        }246    };247 248    // 校验 token 和 checksum249    let (token, checksum) = match validate_token_and_checksum(&auth_token) {250        Some(parts) => parts,251        None => {252            return Json(BasicCalibrationResponse {253                status: ApiStatus::Error,254                message: Some("无效令牌或无效校验和".to_string()),255                user_id: None,256                create_at: None,257                checksum_time: None,258            })259        }260    };261 262    // 提取用户ID和创建时间263    let user_id = extract_user_id(&token);264    let create_at = extract_time(&token).map(|dt| dt.to_string());265    let checksum_time = extract_time_ks(&checksum[..8]);266 267    // 返回校验结果268    Json(BasicCalibrationResponse {269        status: ApiStatus::Success,270        message: Some("校验成功".to_string()),271        user_id,272        create_at,273        checksum_time,274    })275}276