use std::sync::Arc; use argon2::{password_hash::SaltString, PasswordHash, PasswordHasher, PasswordVerifier}; use axum_extra::extract::cookie::Cookie; use chrono::{DateTime, Utc}; use rand::rngs::OsRng; use serde::{Deserialize, Serialize}; use sqlx::{query, query_as, Pool}; use uuid::Uuid; pub async fn setup_db(pool: &Pool) { query("create table if not exists parties (id serial primary key, name text)") .execute(pool) .await .unwrap(); query("create table if not exists tickets (id serial primary key, party_id int references parties(id), student text, inside bool)") .execute(pool) .await .unwrap(); query("create table if not exists users (id serial primary key, username text, password text)") .execute(pool) .await .unwrap(); query("create table if not exists user_tokens (id serial primary key, user_id int references users(id), token uuid, expires_at timestamp)") .execute(pool) .await .unwrap(); } #[derive(sqlx::FromRow, Debug, Deserialize, Serialize)] pub struct Party { pub id: i32, pub name: String, } #[derive(sqlx::FromRow, Debug, Deserialize, Serialize)] pub struct Ticket { pub id: i32, party_id: i32, pub student: String, pub inside: bool, } #[derive(sqlx::FromRow, Debug, Deserialize, Serialize)] pub struct User { id: i32, pub username: String, password: String, } impl User { pub fn verify_password(&self, password: &str) -> bool { let hashed_password = PasswordHash::new(&self.password).unwrap(); argon2::Argon2::default() .verify_password(password.as_bytes(), &hashed_password) .is_ok() } pub fn update_password(&mut self, password: &str) { let salt = SaltString::generate(&mut OsRng); let argon2 = argon2::Argon2::default(); let password = argon2 .hash_password(password.as_bytes(), &salt) .unwrap() .to_string(); self.password = password; } } #[derive(sqlx::FromRow, Debug, Deserialize, Serialize)] pub struct UserToken { id: i32, user_id: i32, token: Uuid, expires_at: chrono::NaiveDateTime, } impl UserToken { pub fn new(user_id: i32, expires_at: chrono::DateTime) -> Self { Self { id: 0, user_id, token: Uuid::new_v4(), expires_at: expires_at.naive_utc(), } } pub fn to_cookie(self) -> Cookie<'static> { let mut c = Cookie::new("token", self.token.to_string()); c.make_permanent(); c } } pub type Database = Arc; pub struct Db(Pool); impl Db { pub fn new(pool: Pool) -> Self { Self(pool) } pub async fn get_all_parties(&self) -> Vec { query_as("select id, name from parties") .fetch_all(&self.0) .await .unwrap() } pub async fn get_party(&self, id: i32) -> Option { query_as("select id, name from parties where id = $1") .bind(id) .fetch_optional(&self.0) .await .unwrap() } pub async fn get_party_by_name(&self, name: &str) -> Option { query_as("select id, name from parties where name = $1") .bind(name) .fetch_optional(&self.0) .await .unwrap() } pub async fn add_party(&self, name: &str) -> Party { query_as("insert into parties (name) values ($1) returning *") .bind(name) .fetch_one(&self.0) .await .unwrap() } pub async fn get_all_tickets_for_party(&self, party_id: i32) -> Vec { query_as("select * from tickets where party_id = $1") .bind(party_id) .fetch_all(&self.0) .await .unwrap() } pub async fn is_student_in_party(&self, party_id: i32, student: &str) -> Option { query_as("select * from tickets where party_id = $1 and student = $2") .bind(party_id) .bind(student) .fetch_optional(&self.0) .await .unwrap() } pub async fn update_ticket(&self, ticket: &Ticket) { query("update tickets set inside = $1 where id = $2") .bind(ticket.inside) .bind(ticket.id) .execute(&self.0) .await .unwrap(); } pub async fn add_ticket(&self, party_id: i32, student: &str) -> Ticket { query_as("insert into tickets (party_id, student, inside) values ($1, $2, $3) returning *") .bind(party_id) .bind(student) .bind(false) .fetch_one(&self.0) .await .unwrap() } pub async fn get_user_by_username(&self, username: &str) -> Option { query_as("select id, username, password from users where username = $1") .bind(username) .fetch_optional(&self.0) .await .unwrap() } pub async fn add_new_user(&self, username: &str, password: &str) -> User { let salt = SaltString::generate(&mut OsRng); let argon2 = argon2::Argon2::default(); let password = argon2 .hash_password(password.as_bytes(), &salt) .unwrap() .to_string(); query_as("insert into users (username, password) values ($1, $2) returning *") .bind(username) .bind(password) .fetch_one(&self.0) .await .unwrap() } pub async fn add_new_user_token(&self, user: &User, expires_at: DateTime) -> UserToken { let user_token = UserToken::new(user.id, expires_at); query_as( "insert into user_tokens (user_id, token, expires_at) values ($1, $2, $3) returning *", ) .bind(user.id) .bind(&user_token.token) .bind(expires_at) .fetch_one(&self.0) .await .unwrap() } pub async fn get_user_token_by_token(&self, token: Uuid) -> Option { query_as("select * from user_tokens where token = $1") .bind(token) .fetch_optional(&self.0) .await .unwrap() } pub async fn get_user_by_token(&self, token: UserToken) -> Option { query_as("select * from users where id = $1") .bind(token.user_id) .fetch_optional(&self.0) .await .unwrap() } pub async fn save_user(&self, user: &User) { query("update users set username = $1, password = $2 where id = $3") .bind(&user.username) .bind(&user.password) .bind(user.id) .execute(&self.0) .await .unwrap(); } pub async fn remove_ticket(&self, id: i32) { query("delete from tickets where id = $1") .bind(id) .execute(&self.0) .await .unwrap(); } }