Authentication System

#[derive(Serialize, Deserialize, Debug)]
pub struct User {
    id: u64,
    role: Roles,
}

#[derive(Debug, Serialize, Deserialize)]
pub enum Roles {
    Unprivileged,  // No access
    User,          // Basic authenticated user
    Moderator,     // Moderation privileges
    Admin,         // Full administrative access
}

pub fn create_token(&self, role: Roles) -> String {
    let secret_key = std::env::var("SECRET_KEY")
        .expect("SECRET_KEY must be set");
    let key: Hmac<Sha256> = Hmac::new_from_slice(secret_key.as_bytes())
        .expect("HMAC can take key of any size");
    
    let mut claims = BTreeMap::new();
    claims.insert("id", self.id);
    claims.insert("exp", chrono::Utc::now().timestamp() as u64 + 86400); // 24h
    claims.insert("iat", chrono::Utc::now().timestamp() as u64);
    claims.insert("role", role as u64);

    claims.sign_with_key(&key).unwrap()
}

pub fn verify_token(token: &str) -> bool {
    let secret_key = std::env::var("SECRET_KEY")
        .expect("SECRET_KEY must be set");
    let key: Hmac<Sha256> = Hmac::new_from_slice(secret_key.as_bytes())
        .expect("HMAC can take key of any size");

    token.verify_with_key(&key)
        .ok()
        .and_then(|claims: BTreeMap<String, u64>| {
            let exp = claims.get("exp")?;
            let current_time = chrono::Utc::now().timestamp() as u64;
            
            if current_time > *exp || !claims.contains_key("id") {
                None
            } else {
                Some(true)
            }
        })
        .unwrap_or(false)
}

pub fn user_has_role(token: &str, role: Roles) -> bool {
    let secret_key = std::env::var("SECRET_KEY")
        .expect("SECRET_KEY must be set");
    let key: Hmac<Sha256> = Hmac::new_from_slice(secret_key.as_bytes())
        .expect("HMAC can take key of any size");

    token.verify_with_key(&key)
        .ok()
        .and_then(|claims: BTreeMap<String, u64>| {
            let exp = claims.get("exp")?;
            let role_claim = claims.get("role")?;
            let current_time = chrono::Utc::now().timestamp() as u64;

            if current_time > *exp {
                None
            } else if role_claim >= &(role as u64) {
                Some(true)
            } else {
                None
            }
        })
        .unwrap_or(false)
}

#[rocket::async_trait]
impl<'r> FromRequest<'r> for AuthenticatedUser<'r> {
    type Error = ApiKeyError;

    async fn from_request(req: &'r Request<'_>) -> Outcome<Self, Self::Error> {
        fn is_valid(key: &str) -> bool {
            let is_valid = User::verify_token(key);
            let has_role = User::user_has_role(key, Roles::User);
            is_valid && has_role
        }

        let key = req.cookies().get("token").map(|cookie| cookie.value());
        match key {
            None => Outcome::Error((Status::Unauthorized, ApiKeyError::Missing)),
            Some(key) if is_valid(key) => Outcome::Success(AuthenticatedUser(key)),
            Some(_) => Outcome::Error((Status::Unauthorized, ApiKeyError::Invalid)),
        }
    }
}

#[rocket::async_trait]
impl<'r> FromRequest<'r> for AdminUser<'r> {
    type Error = ApiKeyError;

    async fn from_request(req: &'r Request<'_>) -> Outcome<Self, Self::Error> {
        fn is_valid(key: &str) -> bool {
            let is_valid = User::verify_token(key);
            let has_role = User::user_has_role(key, Roles::Admin);
            is_valid && has_role
        }

        let key = req.cookies().get("token").map(|cookie| cookie.value());
        match key {
            None => Outcome::Error((Status::Unauthorized, ApiKeyError::Missing)),
            Some(key) if is_valid(key) => Outcome::Success(AdminUser(key)),
            Some(_) => Outcome::Error((Status::Unauthorized, ApiKeyError::Invalid)),
        }
    }
}

// Requires authenticated user
#[get("/verify")]
pub fn verify(_user: AuthenticatedUser<'_>) -> Template {
    Template::render("verify", &{})
}

// Requires admin
#[get("/admin")]
pub fn admin(_user: AdminUser<'_>) -> Template {
    Template::render("admin", &{})
}

#[get("/logout")]
pub fn logout(jar: &CookieJar<'_>) -> Template {
    jar.remove("token");
    Template::render("logout", &{})
}

jar.add(Cookie::build(("token", token))
    .path("/")
    .secure(true)      // HTTPS only
    .http_only(true)); // No JavaScript access

let example_user = User::new(1234567890).create_token(web::auth::Roles::Admin);
println!("Example admin token: {}", example_user);