1  use std::{
  2      sync::{Arc, Mutex},
  3      time::{SystemTime, UNIX_EPOCH},
  4  };
  5  
  6  use axum::{
  7      Router,
  8      extract::{FromRequest, Request, State},
  9      http::StatusCode,
 10      middleware::{self, Next},
 11      response::{IntoResponse, Redirect, Response},
 12  };
 13  use bcrypt;
 14  
 15  use crate::config::Config;
 16  use axum::extract::Form;
 17  use axum::http::header;
 18  use jsonwebtoken::{DecodingKey, EncodingKey, Header, decode, encode};
 19  use serde::{Deserialize, Serialize};
 20  use sysinfo::System;
 21  
 22  #[derive(Serialize, Deserialize)]
 23  struct Claims {
 24      exp: usize,
 25      iat: usize,
 26  }
 27  
 28  #[derive(Deserialize)]
 29  struct LoginForm {
 30      password: String,
 31  }
 32  
 33  pub async fn auth_handler(
 34      State((_, config)): State<(Arc<Mutex<System>>, Config)>,
 35      request: Request,
 36  ) -> impl IntoResponse {
 37      // Extract form data
 38      let pass = match Form::<LoginForm>::from_request(request, &()).await {
 39          Ok(form) => form.password.clone(),
 40          Err(_) => "".to_string(),
 41      };
 42  
 43      // Check if password matches
 44      if bcrypt::verify(pass, &config.password_hash.unwrap_or_default()).unwrap_or(false) {
 45          // Create JWT token
 46          let now = SystemTime::now()
 47              .duration_since(UNIX_EPOCH)
 48              .unwrap()
 49              .as_secs();
 50  
 51          let claims = Claims {
 52              exp: now as usize + 60 * 86400, // 60 days
 53              iat: now as usize,
 54          };
 55  
 56          let token = encode(
 57              &Header::default(),
 58              &claims,
 59              &EncodingKey::from_secret(config.jwt_secret.as_bytes()),
 60          )
 61          .unwrap_or_default();
 62  
 63          return Response::builder()
 64              .status(StatusCode::OK)
 65              .header(
 66                  header::SET_COOKIE,
 67                  format!(
 68                      "prheri_auth_token={}; Path=/; HttpOnly; SameSite=Strict; Max-Age=5184000",
 69                      token
 70                  ),
 71              )
 72              .body("logged in".to_string())
 73              .unwrap();
 74      }
 75  
 76      return Response::builder()
 77          .status(StatusCode::UNAUTHORIZED)
 78          .body("Unauthorized".to_string())
 79          .unwrap();
 80  }
 81  
 82  // Middleware function to check authentication
 83  async fn auth_middleware(
 84      State(config): State<Config>,
 85      request: Request,
 86      next: Next,
 87  ) -> Result<Response, StatusCode> {
 88      // Skip authentication for root path to allow login page access
 89      if request.uri().path() == "/auth" {
 90          return Ok(next.run(request).await);
 91      }
 92  
 93      // Extract JWT token from cookie
 94      let token = match request.headers().get("cookie") {
 95          Some(cookie) => {
 96              let cookie_str = cookie.to_str().unwrap_or_default();
 97              let token = cookie_str
 98                  .split(';')
 99                  .find(|c| c.contains("prheri_auth_token"))
100                  .unwrap_or_default()
101                  .split('=')
102                  .nth(1)
103                  .unwrap_or_default();
104  
105              token.to_string()
106          }
107          None => return Ok(Redirect::temporary("/auth").into_response()),
108      };
109  
110      // Verify JWT token
111      let token_data = match decode::<Claims>(
112          &token,
113          &DecodingKey::from_secret(config.jwt_secret.as_bytes()),
114          &jsonwebtoken::Validation::default(),
115      ) {
116          Ok(data) => data.claims,
117          Err(_) => return Ok(Redirect::temporary("/auth").into_response()),
118      };
119  
120      // Check if token is expired
121      let now = SystemTime::now()
122          .duration_since(UNIX_EPOCH)
123          .unwrap()
124          .as_secs() as usize;
125      if token_data.exp < now {
126          return Ok(Redirect::temporary("/auth").into_response());
127      }
128  
129      return Ok(next.run(request).await);
130  }
131  
132  pub fn apply_auth_middleware(app: Router, config: Config) -> Router {
133      app.layer(middleware::from_fn_with_state(config, auth_middleware))
134  }