SomeBot/main.cpp

#include "bot.hpp"
#include "util.hpp"

#include <string>
#include <string_view>
#include <sstream>
#include <fstream>
#include <vector>
#include <unordered_map>
#include <algorithm>
#include <thread>
#include <mutex>
#include <ctime>

#include <exception>
#include <dpp/dpp.h>
#include <dpp/nlohmann/json.hpp>
#include <sys/mman.h>



std::vector<std::function<void (Bot*)>> Bot::on_init,
                                        Bot::on_deinit;

Bot::Bot(const std::string& token, const std::string& database) : cluster(token), db(database) {
    // Set up bot
    cluster.on_slashcommand([this](const dpp::slashcommand_t& event) {
        try {
            for (const auto& command : commands.chat_commands) {
                if (command == event.command) {
                    command(event);
                }
            }
            for (const auto& [command, guild_id] : commands.guild_chat_commands) {
                if (command == event.command) {
                    command(event);
                }
            }
        } catch (std::exception& e) {
            event.reply(":sos: Unhandled exception at `"+Util::int_as_hex(&e)+"`: "+e.what());
        }
    });
    cluster.on_user_context_menu([this](const dpp::user_context_menu_t& event) {
        try {
            for (const auto& command : commands.user_commands) {
                if (command == event.command) {
                    command(event);
                }
            }
            for (const auto& [command, guild_id] : commands.guild_user_commands) {
                if (command == event.command) {
                    command(event);
                }
            }
        } catch (std::exception& e) {
            event.reply(":sos: Unhandled exception at `"+Util::int_as_hex(&e)+"`: "+e.what());
        }
    });
    cluster.on_message_context_menu([this](const dpp::message_context_menu_t& event) {
        try {
            for (const auto& command : commands.message_commands) {
                if (command == event.command) {
                    command(event);
                }
            }
            for (const auto& [command, guild_id] : commands.guild_message_commands) {
                if (command == event.command) {
                    command(event);
                }
            }
        } catch (std::exception& e) {
            event.reply(":sos: Unhandled exception at `"+Util::int_as_hex(&e)+"`: "+e.what());
        }
    });
    cluster.on_ready([this](const dpp::ready_t&) {  //TODO: Clear this macro mess without compromising performance... actually, does that even matter here?
        // Register bot commands once
        if (dpp::run_once<struct register_bot_commands>()) {
            std::thread([this]() {
                // Clear all commands first if requested
                if (config.reregister_commands) {
                    for (const auto& command : cluster.global_commands_get_sync()) {
                        try {
                            cluster.global_command_delete_sync(command.first);
                        } catch (...) {}
                    }
                    for (const auto& guild : cluster.current_user_get_guilds_sync()) {
                        for (const auto& command : cluster.guild_commands_get_sync(guild.first)) {
                            try {
                                cluster.guild_command_delete_sync(command.first, guild.first);
                            } catch (...) {}
                        }
                    }
                }
                // Add global commands
                #define ADD_SC_RESPECTING_CMD_ALIAS_MODE \
                    if (command.names.empty()) continue; \
                    switch (config.command_alias_mode) { \
                    case Config::CommandAliasMode::all: { \
                        for (const auto& name : command.names) { \
                            sc.set_name(name); \
                            slashcommands.push_back(sc); \
                        } \
                    } break; \
                    case Config::CommandAliasMode::first: { \
                        sc.set_name(command.names.front()); \
                        slashcommands.push_back(sc); \
                    } break; \
                    case Config::CommandAliasMode::last: { \
                        sc.set_name(command.names.back()); \
                        slashcommands.push_back(sc); \
                    } break; \
                    }
                {
                    std::vector<dpp::slashcommand> slashcommands;
                    for (const auto& command : commands.chat_commands) {
                        auto sc = command.misc;
                        sc.set_description(command.description);
                        sc.set_type(dpp::slashcommand_contextmenu_type::ctxm_chat_input);
                        ADD_SC_RESPECTING_CMD_ALIAS_MODE
                    }
                    for (const auto& command : commands.user_commands) {
                        auto sc = command.misc;
                        sc.set_description(command.description);
                        sc.set_type(dpp::slashcommand_contextmenu_type::ctxm_user);
                        ADD_SC_RESPECTING_CMD_ALIAS_MODE
                    }
                    for (const auto& command : commands.message_commands) {
                        auto sc = command.misc;
                        sc.set_description(command.description);
                        sc.set_type(dpp::slashcommand_contextmenu_type::ctxm_message);
                        ADD_SC_RESPECTING_CMD_ALIAS_MODE
                    }
                    cluster.global_bulk_command_create(slashcommands);
                }
                // Add guild commands
                #undef ADD_SC_RESPECTING_CMD_ALIAS_MODE
                #define ADD_SC_RESPECTING_CMD_ALIAS_MODE \
                    if (command.names.empty()) continue; \
                    switch (config.command_alias_mode) { \
                    case Config::CommandAliasMode::all: { \
                        for (const auto& name : command.names) { \
                            sc.set_name(name); \
                            cluster.guild_command_create(sc, guild_id); \
                        } \
                    } break; \
                    case Config::CommandAliasMode::first: { \
                        sc.set_name(command.names.front()); \
                        cluster.guild_command_create(sc, guild_id); \
                    } break; \
                    case Config::CommandAliasMode::last: { \
                        sc.set_name(command.names.back()); \
                        cluster.guild_command_create(sc, guild_id); \
                    } break; \
                    }
                {
                    for (const auto& [command, guild_id] : commands.guild_chat_commands) {
                        auto sc = command.misc;
                        sc.set_description(command.description);
                        sc.set_type(dpp::slashcommand_contextmenu_type::ctxm_chat_input);
                        sc.set_application_id(cluster.me.id);
                        ADD_SC_RESPECTING_CMD_ALIAS_MODE
                    }
                    for (const auto& [command, guild_id] : commands.guild_user_commands) {
                        auto sc = command.misc;
                        sc.set_description(command.description);
                        sc.set_type(dpp::slashcommand_contextmenu_type::ctxm_user);
                        sc.set_application_id(cluster.me.id);
                        ADD_SC_RESPECTING_CMD_ALIAS_MODE
                    }
                    for (const auto& [command, guild_id] : commands.guild_message_commands) {
                        auto sc = command.misc;
                        sc.set_description(command.description);
                        sc.set_type(dpp::slashcommand_contextmenu_type::ctxm_message);
                        sc.set_application_id(cluster.me.id);
                        ADD_SC_RESPECTING_CMD_ALIAS_MODE
                    }
                }
                #undef ADD_SC_RESPECTING_CMD_ALIAS_MODE
                // We're done
                config.reregister_commands = false;
            }).detach();
        }
        // Log user details
        cluster.log(dpp::loglevel::ll_info, "Identified as "+cluster.me.format_username()+" ("+std::to_string(cluster.me.id)+").");
        cluster.log(dpp::loglevel::ll_info, "Invite me using this link: "+dpp::utility::bot_invite_url(cluster.me.id, dpp::permissions::p_administrator));
    });
}
Bot::~Bot() {
    if (inited) {
        // Run on_deinit callbacks
        for (const auto& cb : on_deinit) {
            cb(this);
        }
    }
}

void Bot::start() {
    if (!inited) {
        // Run on_init callbacks
        for (const auto& cb : on_init) {
            cb(this);
        }
        inited = true;
    }
    // Start bot
    cluster.start(dpp::start_type::st_wait);
}

bool Bot::is_module_enabled(std::string_view name) {
    auto fres = std::find(config.modules.begin(), config.modules.end(), name) != config.modules.end();
    if (config.modules_mode == Config::ModulesMode::blacklist) {
        fres = !fres;
    }
    return fres;
}

void Bot::add_chatcommand(ChatCommand command, const ChatCommand::Callback& cb, dpp::snowflake guild_id) {
    if (guild_id == dpp::snowflake(123)) return;
    auto sc = std::move(command.misc);
    command.callback = cb;
    command.misc = std::move(sc);
    if (!guild_id) {
        commands.chat_commands.push_back(std::move(command));
    } else {
        commands.guild_chat_commands.push_back({std::move(command), guild_id});
    }
}
void Bot::add_usercommand(UserCommand command, const UserCommand::Callback& cb, dpp::snowflake guild_id) {
    if (guild_id == dpp::snowflake(123)) return;
    auto sc = std::move(command.misc);
    command.callback = cb;
    command.misc = std::move(sc);
    commands.user_commands.push_back(std::move(command));
    if (!guild_id) {
        commands.user_commands.push_back(std::move(command));
    } else {
        commands.guild_user_commands.push_back({std::move(command), guild_id});
    }
}
void Bot::add_messagecommand(MessageCommand command, const MessageCommand::Callback& cb, dpp::snowflake guild_id) {
    if (guild_id == dpp::snowflake(123)) return;
    auto sc = std::move(command.misc);
    command.callback = cb;
    command.misc = std::move(sc);
    if (!guild_id) {
        commands.message_commands.push_back(std::move(command));
    } else {
        commands.guild_message_commands.push_back({std::move(command), guild_id});
    }
}


void print_usage(char **argv) {
    std::cout << "Usage: " << argv[0] << " [\"--reregister-commands\"]" << std::endl;
}

int main(int argc, char **argv) {
    if (argc > 1 && strcmp(argv[1], "--reregister-commands") != 0) {
        print_usage(argv);
        return -2;
    }

    // Configure ASAN
    setenv("ASAN_OPTIONS", "new_delete_type_mismatch=0", true);

    // Parse config file
    nlohmann::json config_json;
    std::ifstream config_file("config.json");
    if (!config_file) {
        std::cerr << "Failed to load config file \"config.json\"!" << std::endl;
        return -3;
    }
    config_file >> config_json;
    config_file.close();

    // Create, configure and start each bot
    std::vector<std::thread> threads;
    std::vector<Bot*> instances;
    std::mutex instances_mutex;
    for (const auto& config_entry : config_json) {
        std::thread thread([&]() {
            // setup from JSON
            if (!config_entry.contains("token") || !config_entry.contains("id")) {
                std::cerr << "Each config entry needs at least a token and an id." << std::endl;
                exit(-4);
            }
            Bot bot(config_entry["token"], std::string(config_entry["id"])+"-db.sqlite3");
            bot.config.id = config_entry["id"];
            if (config_entry.contains("owner_id")) {
                bot.config.owner_id = std::string(config_entry["owner_id"]);
            }
            if (config_entry.contains("management_guild_id")) {
                bot.config.management_guild_id = std::string(config_entry["management_guild_id"]);
            }
            if (config_entry.contains("private")) {
                bot.config._private = bool(config_entry["private"]);
            }
            if (config_entry.contains("modules")) {
                for (const auto& module_name : config_entry["modules"]) {
                    bot.add_module(std::string(module_name));
                }
                bot.config.modules_mode = Bot::Config::ModulesMode::whitelist;
            }
            if (config_entry.contains("modules_mode")) {
                const std::string& mode = config_entry["modules_mode"];
                if (mode == "blacklist") {
                    bot.config.modules_mode = Bot::Config::ModulesMode::blacklist;
                } else if (mode == "whitelist") {
                    bot.config.modules_mode = Bot::Config::ModulesMode::whitelist;
                } else {
                    std::cerr << "Invalid modules mode \"" << mode << "\"! Accepted values: \"whitelist\", \"blacklist\"" << std::endl;
                    exit(-5);
                }
            }
            if (config_entry.contains("command_alias_mode")) {
                std::string mode = config_entry["command_alias_mode"];
                if (mode == "all") {
                    bot.config.command_alias_mode = Bot::Config::CommandAliasMode::all;
                } else if (mode == "first") {
                    bot.config.command_alias_mode = Bot::Config::CommandAliasMode::first;
                } else if (mode == "last") {
                    bot.config.command_alias_mode = Bot::Config::CommandAliasMode::last;
                } else {
                    std::cerr << "Invalid command alias mode \"" << mode << "\"! Accepted values: \"all\", \"first\", \"last\"" << std::endl;
                    exit(-5);
                }
            }
            // Misc setup
            std::ofstream logfile(bot.config.id+"-log.txt", std::ios_base::app);
            bot.cluster.on_log([&](const dpp::log_t& event) {
                auto t = std::time(nullptr);
                logfile << "(" << std::put_time(std::localtime(&t), "%Y-%m-%d %H:%M:%S") << ") [" << dpp::utility::loglevel(event.severity) << "]: " << event.message << std::endl;
                std::cout << "(" << std::put_time(std::localtime(&t), "%Y-%m-%d %H:%M:%S") << ") [" << bot.config.id << "/" << dpp::utility::loglevel(event.severity) << "]: " << event.message << std::endl;
            });
            bot.config.reregister_commands = argc == 2;
            bot.add_property("main_all_instances", &instances, instances_mutex);
            // Add bot to bot list
            {
                std::scoped_lock L(instances_mutex);
                instances.push_back(&bot);
            }
            // Run
            start:
            try {
                logfile << " -------- Cut here -------- " << std::endl;
                bot.start();
            } catch (std::exception& e) {
                bot.cluster.log(dpp::loglevel::ll_critical, "BOT RESTARTING IN RESPONSE TO LEAKED EXCEPTION: "+std::string(e.what()));
                goto start;
            }
            // Remove bot from bot list
            std::scoped_lock L(instances_mutex);
            std::remove(instances.begin(), instances.end(), &bot);
        });
        threads.push_back(std::move(thread));
    }

    // Wait for all threads to complete
    for (auto& thread : threads) {
        if (thread.joinable()) thread.join();
    }
}