use axum::{ extract::{Query, State}, http::StatusCode, routing::get, Json, Router, }; use chrono::{prelude::*, Duration, ParseError}; use serde::{Deserialize, Serialize}; use std::{env, ops::Deref, sync::Arc}; mod db; mod generate; #[tokio::main] async fn main() -> Result<(), Box> { // 根据操作系统选择数据库路径 let db_path: String = env::current_dir() .expect("无法获取当前目录") .join("authorize_data/database.db") .to_string_lossy() .into_owned(); let db = Arc::new(db::Db::new(db_path).await?); let generator = Arc::new( generate::TokenGenerator::new() .with_uppercase(true) .with_lowercase(true) .with_numbers(true), ); let state = AppState { db, generator }; let app = Router::new() .route("/create", get(create_token)) .route("/verify", get(verify_token)) .route("/info", get(get_token_info)) .route("/reset", get(update_token_status)) .route("/renewal", get(renewal_token)) .route("/project", get(get_projects)) .with_state(state); let listener = tokio::net::TcpListener::bind("0.0.0.0:3009").await?; axum::serve(listener, app).await?; Ok(()) } async fn create_token( State(state): State, Query(args): Query, ) -> (StatusCode, Json) { let exists = state .db .exists_project(&args.project, &args.device_id) .await .unwrap(); match exists { Some(info) => { return ( StatusCode::OK, Json(CreateTokenInfo { code: 404, project: info.project, device_id: info.device_id, token: info.token, msg: "token已存在,请勿重复创建".to_owned(), }), ) } None => (), } let str_time = get_current_datetime(); let exp_time = add_day(&str_time, 7).unwrap(); let token: String = state.generator.generate(16); let _token_id = state .db .insert_authorize(InsertArgs { project: args.project.clone(), token: token.clone(), device_id: args.device_id.clone(), disable: 1, expire: exp_time, insert_time: str_time, }) .await; ( StatusCode::OK, Json(CreateTokenInfo { code: 200, project: args.project, device_id: args.device_id, token, msg: "token创建成功".to_owned(), }), ) } async fn verify_token( State(state): State, Query(args): Query, ) -> (StatusCode, Json) { if state.db.verify_token(&args.token).await.unwrap() { ( StatusCode::OK, Json(VerifyResult { code: 200, msg: "正常".to_owned(), }), ) } else { ( StatusCode::OK, Json(VerifyResult { code: 404, msg: "token已过期".to_owned(), }), ) } } async fn get_token_info( State(state): State, Query(args): Query, ) -> (StatusCode, Json) { let auth_info = state.db.get_token_info(&args.token).await.unwrap(); match auth_info { Some(auth) => { let v = auth.into_inner(); ( StatusCode::OK, Json(TokenResponse::Success(TokenInfo { code: 200, project: v.project, token: v.token, device_id: v.device_id, disable: v.disable, expire: v.expire, insert_time: v.insert_time, })), ) } None => ( StatusCode::OK, Json(TokenResponse::Error(QueryError { code: 404, msg: "未查询token相关信息".to_owned(), })), ), } } async fn update_token_status( State(state): State, Query(args): Query, ) -> (StatusCode, Json) { let result = if args.enable { state.db.update_token_state(&args.token, 1).await.unwrap() } else { state.db.update_token_state(&args.token, 0).await.unwrap() }; if result { ( StatusCode::OK, Json(VerifyResult { code: 200, msg: "操作成功".to_owned(), }), ) } else { ( StatusCode::OK, Json(VerifyResult { code: 404, msg: "操作失败".to_owned(), }), ) } } async fn renewal_token( State(state): State, Query(args): Query, ) -> (StatusCode, Json) { let auth_info = state .db .update_token_expiry(&args.token, args.days) .await .unwrap(); match auth_info { Some(auth) => { let v = auth.into_inner(); ( StatusCode::OK, Json(TokenResponse::Success(TokenInfo { code: 200, project: v.project, token: v.token, device_id: v.device_id, disable: v.disable, expire: v.expire, insert_time: v.insert_time, })), ) } None => ( StatusCode::OK, Json(TokenResponse::Error(QueryError { code: 404, msg: "操作失败".to_owned(), })), ), } } /// 获取所有的项目名 async fn get_projects( State(state): State, ) -> (StatusCode, Json) { let projects = state.db.get_all_project().await.unwrap(); ( StatusCode::OK, Json(ProjectResponse { code: 200, projects: projects, }), ) } fn get_current_datetime() -> String { Local::now().format("%Y-%m-%d %H:%M:%S").to_string() } fn add_day(t: &str, days: i64) -> Result { let date_time = NaiveDateTime::parse_from_str(t, "%Y-%m-%d %H:%M:%S")?; let new_time = date_time + Duration::days(days); Ok(new_time.format("%Y-%m-%d %H:%M:%S").to_string()) } pub struct InsertArgs { pub project: String, pub token: String, pub device_id: String, pub disable: i8, pub expire: String, pub insert_time: String, } #[derive(Clone)] struct AppState { db: Arc, generator: Arc, } #[derive(Deserialize)] struct CreateToken { project: String, device_id: String, } #[derive(Deserialize, Debug, Clone)] struct Token(String); // 让 Token 可以像 &str 一样使用 impl Deref for Token { type Target = str; fn deref(&self) -> &Self::Target { &self.0 } } #[derive(Deserialize)] struct VerifyToken { token: Token, } #[derive(Deserialize)] struct RenewalToken { token: Token, days: i64, } #[derive(Serialize)] struct VerifyResult { code: i16, msg: String, } #[derive(Serialize)] struct TokenInfo { code: i16, project: String, token: String, device_id: String, disable: i8, expire: String, insert_time: String, } #[derive(Serialize)] struct QueryError { code: i16, msg: String, } #[derive(Serialize)] #[serde(untagged)] enum TokenResponse { Success(TokenInfo), Error(QueryError), } #[derive(Serialize)] struct CreateTokenInfo { code: i16, project: String, device_id: String, token: String, msg: String, } #[derive(Deserialize)] struct UpdateTokenStatus { token: String, enable: bool, } #[derive(Serialize)] struct ProjectResponse { code: i32, projects: Vec, }