reminder-bot/src/main.rs

#[macro_use]
extern crate lazy_static;

mod commands;
mod consts;
mod framework;
mod models;
mod time_parser;

use serenity::{
    cache::Cache,
    client::{bridge::gateway::GatewayIntents, Client},
    framework::Framework,
    http::{client::Http, CacheHttp},
    model::{
        channel::Message,
        id::{GuildId, UserId},
    },
    prelude::TypeMapKey,
};

use sqlx::{
    mysql::{MySqlConnection, MySqlPool},
    Pool,
};

use dotenv::dotenv;

use std::{env, sync::Arc};

use crate::{
    commands::{info_cmds, moderation_cmds, reminder_cmds, todo_cmds},
    consts::{CNC_GUILD, DEFAULT_PREFIX, SUBSCRIPTION_ROLES, THEME_COLOR},
    framework::RegexFramework,
};

use serenity::futures::TryFutureExt;

use log::info;

struct SQLPool;

impl TypeMapKey for SQLPool {
    type Value = Pool<MySqlConnection>;
}

struct ReqwestClient;

impl TypeMapKey for ReqwestClient {
    type Value = Arc<reqwest::Client>;
}

struct FrameworkCtx;

impl TypeMapKey for FrameworkCtx {
    type Value = Arc<Box<dyn Framework + Send + Sync>>;
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    env_logger::init();

    dotenv()?;

    let token = env::var("DISCORD_TOKEN").expect("Missing DISCORD_TOKEN from environment");

    let http = Http::new_with_token(&token);

    let logged_in_id = http
        .get_current_user()
        .map_ok(|user| user.id.as_u64().to_owned())
        .await?;

    let framework = RegexFramework::new(logged_in_id)
        .default_prefix(DEFAULT_PREFIX.clone())
        .ignore_bots(env::var("IGNORE_BOTS").map_or(true, |var| var == "1"))
        // info commands
        .add_command("ping", &info_cmds::PING_COMMAND)
        .add_command("help", &info_cmds::HELP_COMMAND)
        .add_command("info", &info_cmds::INFO_COMMAND)
        .add_command("invite", &info_cmds::INFO_COMMAND)
        .add_command("donate", &info_cmds::DONATE_COMMAND)
        .add_command("dashboard", &info_cmds::DASHBOARD_COMMAND)
        .add_command("clock", &info_cmds::CLOCK_COMMAND)
        // reminder commands
        .add_command("timer", &reminder_cmds::TIMER_COMMAND)
        .add_command("remind", &reminder_cmds::REMIND_COMMAND)
        .add_command("r", &reminder_cmds::REMIND_COMMAND)
        .add_command("interval", &reminder_cmds::INTERVAL_COMMAND)
        .add_command("i", &reminder_cmds::INTERVAL_COMMAND)
        .add_command("natural", &reminder_cmds::NATURAL_COMMAND)
        .add_command("n", &reminder_cmds::NATURAL_COMMAND)
        .add_command("", &reminder_cmds::NATURAL_COMMAND)
        // management commands
        .add_command("look", &reminder_cmds::LOOK_COMMAND)
        .add_command("del", &reminder_cmds::DELETE_COMMAND)
        // to-do commands
        .add_command("todo", &todo_cmds::TODO_USER_COMMAND)
        .add_command("todo user", &todo_cmds::TODO_USER_COMMAND)
        .add_command("todoc", &todo_cmds::TODO_CHANNEL_COMMAND)
        .add_command("todo channel", &todo_cmds::TODO_CHANNEL_COMMAND)
        .add_command("todos", &todo_cmds::TODO_GUILD_COMMAND)
        .add_command("todo server", &todo_cmds::TODO_GUILD_COMMAND)
        .add_command("todo guild", &todo_cmds::TODO_GUILD_COMMAND)
        // moderation commands
        .add_command("blacklist", &moderation_cmds::BLACKLIST_COMMAND)
        .add_command("restrict", &moderation_cmds::RESTRICT_COMMAND)
        .add_command("timezone", &moderation_cmds::TIMEZONE_COMMAND)
        .add_command("prefix", &moderation_cmds::PREFIX_COMMAND)
        .add_command("lang", &moderation_cmds::LANGUAGE_COMMAND)
        .add_command("pause", &reminder_cmds::PAUSE_COMMAND)
        .add_command("offset", &reminder_cmds::OFFSET_COMMAND)
        .add_command("nudge", &reminder_cmds::NUDGE_COMMAND)
        .add_command("alias", &moderation_cmds::ALIAS_COMMAND)
        .add_command("a", &moderation_cmds::ALIAS_COMMAND)
        .build();

    let framework_arc = Arc::new(Box::new(framework) as Box<dyn Framework + Send + Sync>);

    let mut client = Client::new(&token)
        .intents(
            GatewayIntents::GUILD_MESSAGES
                | GatewayIntents::GUILDS
                | GatewayIntents::DIRECT_MESSAGES,
        )
        .framework_arc(framework_arc.clone())
        .await
        .expect("Error occurred creating client");

    {
        let pool = MySqlPool::new(
            &env::var("DATABASE_URL").expect("Missing DATABASE_URL from environment"),
        )
        .await
        .unwrap();

        let mut data = client.data.write().await;

        data.insert::<SQLPool>(pool);
        data.insert::<ReqwestClient>(Arc::new(reqwest::Client::new()));
        data.insert::<FrameworkCtx>(framework_arc);
    }

    if let Ok((Some(lower), Some(upper))) = env::var("SHARD_RANGE").map(|sr| {
        let mut split = sr
            .split(',')
            .map(|val| val.parse::<u64>().expect("SHARD_RANGE not an integer"));

        (split.next(), split.next())
    }) {
        let total_shards = env::var("SHARD_COUNT")
            .map(|shard_count| shard_count.parse::<u64>().ok())
            .ok()
            .flatten()
            .expect("No SHARD_COUNT provided, but SHARD_RANGE was provided");

        assert!(
            lower < upper,
            "SHARD_RANGE lower limit is not less than the upper limit"
        );

        info!(
            "Starting client fragment with shards {}-{}/{}",
            lower, upper, total_shards
        );

        client
            .start_shard_range([lower, upper], total_shards)
            .await?;
    } else if let Ok(total_shards) = env::var("SHARD_COUNT").map(|shard_count| {
        shard_count
            .parse::<u64>()
            .expect("SHARD_COUNT not an integer")
    }) {
        info!("Starting client with {} shards", total_shards);

        client.start_shards(total_shards).await?;
    } else {
        info!("Starting client as autosharded");

        client.start_autosharded().await?;
    }

    Ok(())
}

pub async fn check_subscription(cache_http: impl CacheHttp, user_id: impl Into<UserId>) -> bool {
    if let Some(subscription_guild) = *CNC_GUILD {
        let guild_member = GuildId(subscription_guild)
            .member(cache_http, user_id)
            .await;

        if let Ok(member) = guild_member {
            for role in member.roles {
                if SUBSCRIPTION_ROLES.contains(role.as_u64()) {
                    return true;
                }
            }
        }

        false
    } else {
        true
    }
}

pub async fn check_subscription_on_message(
    cache_http: impl CacheHttp + AsRef<Cache>,
    msg: &Message,
) -> bool {
    check_subscription(&cache_http, &msg.author).await
        || if let Some(guild) = msg.guild(&cache_http).await {
            check_subscription(&cache_http, guild.owner_id).await
        } else {
            false
        }
}