autter-core 2.0.2

use crate::{
    DataManager,
    model::{
        AuditLogEntry, Error, Result, Token, User, UserBadge, UserLinkedAccounts, UserPermission,
        UserSettings, organizations::Organization,
    },
};
use oiseau::cache::Cache;
use oiseau::{PostgresRow, execute, get, params, query_row};
use tetratto_core2::{auto_method, model::id::Id};
use tritools::encoding::{hash_salted, salt};

#[inline]
fn transform_uid(uid: &Id) -> usize {
    uid.as_usize() / 2
}

impl DataManager {
    /// Get a [`User`] from an SQL row.
    pub(crate) fn get_user_from_row(x: &PostgresRow) -> User {
        User {
            id: Id::Legacy(get!(x->0(i64)) as usize),
            created: get!(x->1(i64)) as u128,
            username: get!(x->2(String)),
            password: get!(x->3(String)),
            salt: get!(x->4(String)),
            settings: serde_json::from_str(&get!(x->5(String)).to_string()).unwrap(),
            tokens: serde_json::from_str(&get!(x->6(String)).to_string()).unwrap(),
            permissions: serde_json::from_str(&get!(x->7(String)).to_string()).unwrap(),
            is_verified: get!(x->8(i32)) as i8 == 1,
            notification_count: {
                let x = get!(x->9(i32)) as usize;
                // we're a little too close to the maximum count, clearly something's gone wrong
                if x > usize::MAX - 1000 { 0 } else { x }
            },
            totp: get!(x->10(String)),
            recovery_codes: serde_json::from_str(&get!(x->11(String)).to_string()).unwrap(),
            stripe_id: get!(x->12(String)),
            ban_reason: get!(x->13(String)),
            ban_expire: get!(x->14(i64)) as usize,
            is_deactivated: get!(x->15(i32)) as i8 == 1,
            checkouts: serde_json::from_str(&get!(x->16(String)).to_string()).unwrap(),
            last_policy_consent: get!(x->17(i64)) as u128,
            linked_accounts: serde_json::from_str(&get!(x->18(String)).to_string()).unwrap(),
            badges: serde_json::from_str(&get!(x->19(String)).to_string()).unwrap(),
            principal_org: Id::deserialize(
                &get!(x->20(String) default (Id::default().printable())),
            ),
            org_as_tenant: get!(x->21(i32)) as i8 == 1,
            org_creation_credits: get!(x->22(i32)),
            org_user_register_credits: get!(x->23(i32)),
            is_super_verified: get!(x->24(i32)) as i8 == 1,
        }
    }

    auto_method!(get_user_by_id(Id :: usize)@get_user_from_row -> "SELECT * FROM a_users WHERE id = $1::bigint" --name="user" --returns=User --cache-key-tmpl="srmp.user:{}");
    auto_method!(get_user_by_username(&str)@get_user_from_row -> "SELECT * FROM a_users WHERE username = $1" --name="user" --returns=User --cache-key-tmpl="srmp.user:{}");
    auto_method!(get_user_by_username_no_cache(&str)@get_user_from_row -> "SELECT * FROM a_users WHERE username = $1" --name="user" --returns=User);

    /// Get a user given just their ID. Returns the void user if the user doesn't exist.
    ///
    /// # Arguments
    /// * `id` - the ID of the user
    pub async fn get_user_by_id_with_void(&self, id: &Id) -> Result<User> {
        let conn = match self.0.connect().await {
            Ok(c) => c,
            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
        };

        let res = query_row!(
            &conn,
            "SELECT * FROM a_users WHERE id = $1::bigint",
            &[&id.printable()],
            |x| Ok(Self::get_user_from_row(x))
        );

        if res.is_err() {
            return Ok(User::deleted());
            // return Err(Error::UserNotFound);
        }

        Ok(res.unwrap())
    }

    /// Get a user given just their auth token.
    ///
    /// # Arguments
    /// * `token` - the token of the user
    pub async fn get_user_by_token(&self, token: &str) -> Result<User> {
        let conn = match self.0.connect().await {
            Ok(c) => c,
            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
        };

        let res = query_row!(
            &conn,
            "SELECT * FROM a_users WHERE tokens LIKE $1",
            &[&format!("%\"{token}\"%")],
            |x| Ok(Self::get_user_from_row(x))
        );

        if res.is_err() {
            return Err(Error::UserNotFound);
        }

        Ok(res.unwrap())
    }

    /// Create a new user in the database.
    ///
    /// # Arguments
    /// * `data` - a mock [`User`] object to insert
    pub async fn create_user(&self, mut data: User) -> Result<User> {
        if !self.0.0.security.registration_enabled {
            return Err(Error::RegistrationDisabled);
        }

        data.username = data.username.to_lowercase();

        // check values
        if data.username.len() < 2 {
            return Err(Error::DataTooShort("username".to_string()));
        } else if data.username.len() > 32 {
            return Err(Error::DataTooLong("username".to_string()));
        }

        if data.password.len() < 6 {
            return Err(Error::DataTooShort("password".to_string()));
        }

        if self.0.0.banned_usernames.contains(&data.username) {
            return Err(Error::MiscError("This username cannot be used".to_string()));
        }

        // make sure username isn't taken
        if self.get_user_by_username(&data.username).await.is_ok() {
            return Err(Error::UsernameInUse);
        }

        // ...
        let conn = match self.0.connect().await {
            Ok(c) => c,
            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
        };

        let res = execute!(
            &conn,
            "INSERT INTO a_users VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24)",
            params![
                &(transform_uid(&data.id) as i64),
                &(data.created as i64),
                &data.username.to_lowercase(),
                &data.password,
                &data.salt,
                &serde_json::to_string(&data.settings).unwrap(),
                &serde_json::to_string(&data.tokens).unwrap(),
                &serde_json::to_string(&data.permissions).unwrap(),
                &(data.is_verified as i32),
                &0_i32,
                &String::new(),
                "[]",
                &data.stripe_id,
                &data.ban_reason,
                &(data.ban_expire as i64),
                &(data.is_deactivated as i32),
                &serde_json::to_string(&data.checkouts).unwrap(),
                &(data.last_policy_consent as i64),
                &serde_json::to_string(&data.linked_accounts).unwrap(),
                &serde_json::to_string(&data.badges).unwrap(),
                &if data.principal_org != Id::default() {
                    Some(data.principal_org.printable())
                } else {
                    None
                },
                &data.org_creation_credits,
                &data.org_user_register_credits,
                &(data.is_super_verified as i32),
            ]
        );

        if let Err(e) = res {
            return Err(Error::DatabaseError(e.to_string()));
        }

        Ok(data)
    }

    /// Delete an existing user in the database.
    ///
    /// # Arguments
    /// * `id` - the ID of the user
    /// * `password` - the current password of the user
    /// * `force` - if we should delete even if the given password is incorrect
    pub async fn delete_user(&self, id: &Id, password: &str, force: bool) -> Result<User> {
        let user = self.get_user_by_id(&id).await?;

        if (hash_salted(password.to_string(), user.salt.clone()) != user.password) && !force {
            return Err(Error::IncorrectPassword);
        }

        let conn = match self.0.connect().await {
            Ok(c) => c,
            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
        };

        let res = execute!(
            &conn,
            "DELETE FROM a_users WHERE id = $1::bigint",
            &[&id.printable()]
        );

        if let Err(e) = res {
            return Err(Error::DatabaseError(e.to_string()));
        }

        self.cache_clear_user(&user).await;

        // delete uploads
        for upload in match self.1.get_uploads_by_owner_all(user.id.as_usize()).await {
            Ok(x) => x,
            Err(e) => return Err(Error::MiscError(e.to_string())),
        } {
            if let Err(e) = self.1.delete_upload(upload.id).await {
                return Err(Error::MiscError(e.to_string()));
            }
        }

        // ...
        Ok(user)
    }

    pub async fn update_user_verified_status(&self, id: &Id, x: bool, user: User) -> Result<()> {
        if !user.permissions.contains(&UserPermission::ManageVerified) {
            return Err(Error::NotAllowed);
        }

        let other_user = self.get_user_by_id(&id).await?;

        let conn = match self.0.connect().await {
            Ok(c) => c,
            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
        };

        let res = execute!(
            &conn,
            "UPDATE a_users SET is_verified = $1 WHERE id = $2::bigint",
            params![&{ if x { 1 } else { 0 } }, &id.printable()]
        );

        if let Err(e) = res {
            return Err(Error::DatabaseError(e.to_string()));
        }

        self.cache_clear_user(&other_user).await;

        // create audit log entry
        self.create_audit_log_entry(AuditLogEntry::new(
            user.id,
            format!(
                "invoked `update_user_verified_status` with x value `{}` and y value `{}`",
                other_user.id, x
            ),
        ))
        .await?;

        // ...
        Ok(())
    }

    /// Update a user's super verified status. This is intended to be called after they
    /// successfully pay for verification.
    pub async fn update_user_super_verified_status(&self, id: &Id, x: bool) -> Result<()> {
        let other_user = self.get_user_by_id(&id).await?;

        let conn = match self.0.connect().await {
            Ok(c) => c,
            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
        };

        let res = execute!(
            &conn,
            "UPDATE a_users SET is_super_verified = $1 WHERE id = $2::bigint",
            params![&{ if x { 1 } else { 0 } }, &id.printable()]
        );

        if let Err(e) = res {
            return Err(Error::DatabaseError(e.to_string()));
        }

        self.cache_clear_user(&other_user).await;

        // create audit log entry
        self.create_audit_log_entry(AuditLogEntry::new(
            other_user.id.clone(),
            format!(
                "invoked `update_user_super_verified_status` with x value `{}` and y value `{}`",
                other_user.id.printable(),
                x
            ),
        ))
        .await?;

        // ...
        Ok(())
    }

    pub async fn update_user_is_deactivated(&self, id: &Id, x: bool, user: User) -> Result<()> {
        if id != &user.id && !user.permissions.contains(&UserPermission::ManageUsers) {
            return Err(Error::NotAllowed);
        }

        let other_user = self.get_user_by_id(&id).await?;

        let conn = match self.0.connect().await {
            Ok(c) => c,
            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
        };

        let res = execute!(
            &conn,
            "UPDATE a_users SET is_deactivated = $1 WHERE id = $2::bigint",
            params![&{ if x { 1 } else { 0 } }, &id.printable()]
        );

        if let Err(e) = res {
            return Err(Error::DatabaseError(e.to_string()));
        }

        self.cache_clear_user(&other_user).await;

        // create audit log entry (if we aren't the user that is being updated)
        if user.id != other_user.id {
            self.create_audit_log_entry(AuditLogEntry::new(
                user.id,
                format!(
                    "invoked `update_user_is_deactivated` with x value `{}` and y value `{}`",
                    other_user.id, x
                ),
            ))
            .await?;
        }

        // ...
        Ok(())
    }

    pub async fn update_user_password(
        &self,
        id: &Id,
        from: String,
        to: String,
        user: User,
        force: bool,
    ) -> Result<()> {
        // verify password
        if !user.check_password(from.clone()) && !force {
            return Err(Error::MiscError("Password does not match".to_string()));
        }

        // ...
        let conn = match self.0.connect().await {
            Ok(c) => c,
            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
        };

        let new_salt = salt();
        let new_password = hash_salted(to, new_salt.clone());
        let res = execute!(
            &conn,
            "UPDATE a_users SET password = $1, salt = $2 WHERE id = $3",
            params![&new_password.as_str(), &new_salt.as_str(), &id.printable()]
        );

        if let Err(e) = res {
            return Err(Error::DatabaseError(e.to_string()));
        }

        self.cache_clear_user(&user).await;
        Ok(())
    }

    pub async fn update_user_username(&self, id: &Id, to: String, user: User) -> Result<()> {
        // check value
        if to.len() < 2 {
            return Err(Error::DataTooShort("username".to_string()));
        } else if to.len() > 32 {
            return Err(Error::DataTooLong("username".to_string()));
        }

        if self.0.0.banned_usernames.contains(&to) {
            return Err(Error::MiscError("This username cannot be used".to_string()));
        }

        let regex = regex::RegexBuilder::new(r"[^\w_\-\.!]+")
            .multi_line(true)
            .build()
            .unwrap();

        if regex.captures(&to).is_some() {
            return Err(Error::MiscError(
                "This username contains invalid characters".to_string(),
            ));
        }

        // ...
        let conn = match self.0.connect().await {
            Ok(c) => c,
            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
        };

        let res = execute!(
            &conn,
            "UPDATE a_users SET username = $1 WHERE id = $2::bigint",
            params![&to.to_lowercase(), &id.printable()]
        );

        if let Err(e) = res {
            return Err(Error::DatabaseError(e.to_string()));
        }

        self.cache_clear_user(&user).await;
        Ok(())
    }

    /// Validate a given TOTP code for the given profile.
    pub fn check_totp(&self, ua: &User, code: &str) -> bool {
        let totp = ua.totp(Some(
            self.0
                .0
                .host
                .replace("http://", "")
                .replace("https://", "")
                .replace(":", "_"),
        ));

        if let Some(totp) = totp {
            return !code.is_empty()
                && (totp.check_current(code).unwrap()
                    | ua.recovery_codes.contains(&code.to_string()));
        }

        true
    }

    /// Generate 8 random recovery codes for TOTP.
    pub fn generate_totp_recovery_codes() -> Vec<String> {
        let mut out: Vec<String> = Vec::new();

        for _ in 0..9 {
            out.push(salt())
        }

        out
    }

    /// Update the profile's TOTP secret.
    ///
    /// # Arguments
    /// * `id` - the ID of the user
    /// * `secret` - the TOTP secret
    /// * `recovery` - the TOTP recovery codes
    pub async fn update_user_totp(
        &self,
        id: &Id,
        secret: &str,
        recovery: &Vec<String>,
    ) -> Result<()> {
        let user = self.get_user_by_id(&id).await?;

        // update
        let conn = match self.0.connect().await {
            Ok(c) => c,
            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
        };

        let res = execute!(
            &conn,
            "UPDATE a_users SET totp = $1, recovery_codes = $2 WHERE id = $3",
            params![
                &secret,
                &serde_json::to_string(recovery).unwrap(),
                &id.printable()
            ]
        );

        if let Err(e) = res {
            return Err(Error::DatabaseError(e.to_string()));
        }

        self.cache_clear_user(&user).await;
        Ok(())
    }

    /// Enable TOTP for a profile.
    ///
    /// # Arguments
    /// * `id` - the ID of the user to enable TOTP for
    /// * `user` - the user doing this
    ///
    /// # Returns
    /// `Result<(secret, qr base64)>`
    pub async fn enable_totp(&self, id: &Id, user: User) -> Result<(String, String, Vec<String>)> {
        let other_user = self.get_user_by_id(&id).await?;

        if other_user.id != user.id {
            if other_user
                .permissions
                .contains(&UserPermission::ManageUsers)
            {
                // create audit log entry
                self.create_audit_log_entry(AuditLogEntry::new(
                    user.id,
                    format!("invoked `enable_totp` with x value `{}`", other_user.id,),
                ))
                .await?;
            } else {
                return Err(Error::NotAllowed);
            }
        }

        let secret = totp_rs::Secret::default().to_string();
        let recovery = Self::generate_totp_recovery_codes();
        self.update_user_totp(id, &secret, &recovery).await?;

        // fetch profile again (with totp information)
        let other_user = self.get_user_by_id(&id).await?;

        // get totp
        let totp = other_user.totp(Some(
            self.0
                .0
                .host
                .replace("http://", "")
                .replace("https://", "")
                .replace(":", "_"),
        ));

        if totp.is_none() {
            return Err(Error::MiscError("Failed to get TOTP code".to_string()));
        }

        let totp = totp.unwrap();

        // generate qr
        let qr = match totp.get_qr_base64() {
            Ok(q) => q,
            Err(e) => return Err(Error::MiscError(e.to_string())),
        };

        // return
        Ok((totp.get_secret_base32(), qr, recovery))
    }

    /// Get the given user's principal organization.
    pub async fn get_principal_org(&self, user: &User) -> Option<Organization> {
        if user.principal_org == Id::default() {
            return None;
        }

        if let Ok(x) = self.get_organization_by_id(&user.principal_org).await {
            Some(x)
        } else {
            self.update_user_principal_org(&user.id, None)
                .await
                .expect("failed to clear user principal org");

            None
        }
    }

    pub async fn cache_clear_user(&self, user: &User) {
        self.0.1.remove(format!("srmp.user:{}", user.id)).await;
        self.0
            .1
            .remove(format!("srmp.user:{}", user.username))
            .await;
    }

    auto_method!(update_user_permissions(Id :: usize, Vec<UserPermission>)@get_user_by_id -> "UPDATE a_users SET permissions = $1 WHERE id = $2::bigint" --serde --cache-key-tmpl=cache_clear_user);
    auto_method!(update_user_tokens(Id :: usize, Vec<Token>)@get_user_by_id -> "UPDATE a_users SET tokens = $1 WHERE id = $2::bigint" --serde --cache-key-tmpl=cache_clear_user);
    auto_method!(update_user_settings(Id :: usize, UserSettings)@get_user_by_id -> "UPDATE a_users SET settings = $1 WHERE id = $2::bigint" --serde --cache-key-tmpl=cache_clear_user);
    auto_method!(update_user_ban_reason(Id :: usize, &str)@get_user_by_id -> "UPDATE a_users SET ban_reason = $1 WHERE id = $2::bigint" --cache-key-tmpl=cache_clear_user);
    auto_method!(update_user_ban_expire(Id :: usize, i64)@get_user_by_id -> "UPDATE a_users SET ban_expire = $1 WHERE id = $2::bigint" --cache-key-tmpl=cache_clear_user);
    auto_method!(update_user_checkouts(Id :: usize, Vec<String>)@get_user_by_id -> "UPDATE a_users SET checkouts = $1 WHERE id = $2::bigint" --serde --cache-key-tmpl=cache_clear_user);
    auto_method!(update_user_last_policy_consent(Id :: usize, i64)@get_user_by_id -> "UPDATE a_users SET last_policy_consent = $1 WHERE id = $2::bigint" --cache-key-tmpl=cache_clear_user);
    auto_method!(update_user_linked_accounts(Id :: usize, UserLinkedAccounts)@get_user_by_id -> "UPDATE a_users SET linked_accounts = $1 WHERE id = $2::bigint" --serde --cache-key-tmpl=cache_clear_user);
    auto_method!(update_user_badges(Id :: usize, Vec<UserBadge>)@get_user_by_id -> "UPDATE a_users SET badges = $1 WHERE id = $2::bigint" --serde --cache-key-tmpl=cache_clear_user);
    auto_method!(update_user_principal_org(Id :: usize, Option<String>)@get_user_by_id -> "UPDATE a_users SET principal_org = $1 WHERE id = $2::bigint" --cache-key-tmpl=cache_clear_user);
    auto_method!(update_user_org_as_tenant(Id :: usize, i32)@get_user_by_id -> "UPDATE a_users SET org_as_tenant = $1 WHERE id = $2::bigint" --cache-key-tmpl=cache_clear_user);

    auto_method!(get_user_by_stripe_id(&str)@get_user_from_row -> "SELECT * FROM a_users WHERE stripe_id = $1::bigint" --name="user" --returns=User);
    auto_method!(update_user_stripe_id(&str)@get_user_by_id -> "UPDATE a_users SET stripe_id = $1 WHERE id = $2::bigint" --cache-key-tmpl=cache_clear_user);

    auto_method!(update_user_notification_count(Id :: usize, i32)@get_user_by_id -> "UPDATE a_users SET notification_count = $1 WHERE id = $2::bigint" --cache-key-tmpl=cache_clear_user);
    auto_method!(incr_user_notifications(Id :: usize)@get_user_by_id -> "UPDATE a_users SET notification_count = notification_count + 1 WHERE id = $1::bigint" --cache-key-tmpl=cache_clear_user --incr);
    auto_method!(decr_user_notifications(Id :: usize)@get_user_by_id -> "UPDATE a_users SET notification_count = notification_count - 1 WHERE id = $1::bigint" --cache-key-tmpl=cache_clear_user --decr=notification_count);

    auto_method!(incr_user_org_creation_credits(Id :: usize)@get_user_by_id -> "UPDATE a_users SET org_creation_credits = org_creation_credits + 1 WHERE id = $1::bigint" --cache-key-tmpl=cache_clear_user --incr);
    auto_method!(decr_user_org_creation_credits(Id :: usize)@get_user_by_id -> "UPDATE a_users SET org_creation_credits = org_creation_credits - 1 WHERE id = $1::bigint" --cache-key-tmpl=cache_clear_user --decr=org_creation_credits);

    auto_method!(incr_user_org_user_register_credits(Id :: usize)@get_user_by_id -> "UPDATE a_users SET org_user_register_credits = org_user_register_credits + 1 WHERE id = $1::bigint" --cache-key-tmpl=cache_clear_user --incr);
    auto_method!(decr_user_org_user_register_credits(Id :: usize)@get_user_by_id -> "UPDATE a_users SET org_user_register_credits = org_user_register_credits - 1 WHERE id = $1::bigint" --cache-key-tmpl=cache_clear_user --decr=org_user_register_credits);
}