1
0
Fork 0
mirror of https://gitlab.com/niansa/dpplogger.git synced 2025-03-06 20:48:29 +01:00
dpplogger/main.cpp

314 lines
13 KiB
C++

#include "Random.hpp"
#include "sqlite_modern_cpp/sqlite_modern_cpp.h"
#include <iostream>
#include <unordered_map>
#include <optional>
#include <chrono>
#include <memory>
#include <stdexcept>
#include <boost/asio/io_service.hpp>
#include <dcboost/discord.hpp>
#include <dcboost/snowflake.hpp>
using namespace ChatFuse;
using namespace QLog;
class Cache {
struct Miss : public std::runtime_error {
Miss() : std::runtime_error("Object cache miss") {}
};
Logger<Cache> logger;
std::unordered_map<Discord::Snowflake, Json::Value> cache;
public:
Cache() {
logger.inst_ptr = this;
}
const Json::Value& store(const Json::Value& object) {
const auto& id_str = object["id"].asString();
logger.log(Loglevel::verbose, "Stored "+id_str+" in cache");
return cache.insert_or_assign(std::move(id_str), object).first->second;
}
const Json::Value& store(Json::Value&& object) {
const auto& id_str = object["id"].asString();
logger.log(Loglevel::verbose, "Moved "+id_str+" into cache");
return cache.insert_or_assign(std::move(id_str), std::move(object)).first->second;
}
const Json::Value& fetch(Discord::Snowflake id) const {
const auto& entry = cache.find(id);
if (entry == cache.end()) {
logger.log(Loglevel::warn, "Missing "+id.str());
throw Miss();
}
return entry->second;
}
bool has(Discord::Snowflake id) const {
return cache.find(id) != cache.end();
}
};
static inline std::optional<std::string> GetJSONAsOptionalString(const Json::Value& data) {
return data.isString()?data.asString():std::optional<std::string>();
}
class MyClient final : public Discord::Client {
sqlite::database db;
Cache cache;
Json::Value *me;
public:
MyClient(boost::asio::io_service& io, const ChatFuse::Discord::Settings& settings = {})
: ChatFuse::Discord::Client(io, settings),
db("log.sqlite3")
{
// Improve database performance
db << "pragma journal_mode = WAL;";
db << "pragma synchronous = normal;";
db << "pragma temp_store = memory;";
// Create tables
{
db << "CREATE TABLE IF NOT EXISTS messages ("
" id TEXT PRIMARY KEY NOT NULL,"
" type INTEGER NOT NULL,"
" channel_id TEXT NOT NULL,"
" author_id TEXT NOT NULL,"
" replied_to_id TEXT,"
" is_deleted INTEGER DEFAULT 0 NOT NULL,"
" is_edited INTEGER DEFAULT 0 NOT NULL,"
" creation_timestamp TEXT NOT NULL,"
" update_timestamp TEXT DEFAULT '0' NOT NULL,"
" UNIQUE(id)"
");";
db << "CREATE TABLE IF NOT EXISTS message_contents ("
" message_id TEXT NOT NULL,"
" is_initial INTEGER DEFAULT 1 NOT NULL,"
" timestamp TEXT NOT NULL,"
" content TEXT NOT NULL,"
" embed TEXT" // Discord JSON embed object
");";
db << "CREATE TABLE IF NOT EXISTS users ("
" id TEXT NOT NULL,"
" is_initial INTEGER DEFAULT 1 NOT NULL,"
" timestamp TEXT NOT NULL,"
" full_username TEXT NOT NULL," // Username#Discriminator
" avatar TEXT,"
" bio TEXT,"
" has_nitro INTEGER NOT NULL,"
" is_bot INTEGER NOT NULL"
");";
db << "CREATE TABLE IF NOT EXISTS user_statuses ("
" user_id TEXT NOT NULL,"
" is_initial INTEGER DEFAULT 1 NOT NULL,"
" timestamp TEXT NOT NULL,"
" status TEXT," // JSON array or null if not updated: [int:online/dnd/afk/offline enum, ?:emoji (null/empty if none), string:text (empty if none)]
" presence TEXT" // Discord JSON presence object
");";
db << "CREATE TABLE IF NOT EXISTS members ("
" user_id TEXT NOT NULL,"
" guild_id TEXT NOT NULL,"
" is_initial INTEGER DEFAULT 1 NOT NULL,"
" timestamp TEXT NOT NULL,"
" nickname TEXT,"
" in_guild INTEGER NOT NULL,"
" is_untracked INTEGER NOT NULL"
");";
db << "CREATE TABLE IF NOT EXISTS member_voice_connections ("
" channel_id TEXT NOT NULL,"
" user_id TEXT NOT NULL,"
" is_initial INTEGER DEFAULT 1 NOT NULL,"
" timestamp TEXT NOT NULL,"
" voice_channel_id TEXT,"
" self_muted INTEGER NOT NULL,"
" self_deafened INTEGER NOT NULL,"
" server_muted INTEGER NOT NULL,"
" server_deafened INTEGER NOT NULL"
");";
db << "CREATE TABLE IF NOT EXISTS channels ("
" id TEXT NOT NULL,"
" guild_id TEXT,"
" parent_id TEXT,"
" timestamp TEXT NOT NULL,"
" is_initial INTEGER DEFAULT 1 NOT NULL,"
" type INTEGER NOT NULL,"
" name TEXT,"
" topic TEXT"
");";
db << "CREATE TABLE IF NOT EXISTS guilds ("
" id TEXT NOT NULL,"
" timestamp TEXT NOT NULL,"
" is_initial INTEGER DEFAULT 1 NOT NULL,"
" name TEXT NOT NULL,"
" owner_user_id TEXT NOT NULL,"
" in_guild INTEGER DEFAULT 1 NOT NULL"
");";
}
}
protected:
/*
* Guilds
*/
void processGuild(const Json::Value& data) {
const auto& id = data["id"];
// Insert into database
insertGuildUpdate(data, !cache.has(id));
// Store in cache
cache.store(data);
// Insert and store all channels
for (auto channel_data : data["channels"]) {
// Add guild_id to channel (it'll be missing)
channel_data["guild_id"] = id;
// Insert into database
insertChannelUpdate(channel_data, !cache.has(channel_data["id"]));
// Store in cache
cache.store(std::move(channel_data));
}
}
boost::asio::awaitable<void> fetchAllGuilds(const Json::Value& unavailable_guild_array) {
RandomGenerator rng;
rng.seed();
for (const auto& incomplete_data : unavailable_guild_array) {
Discord::Snowflake guild_id = incomplete_data["id"];
// Make sure guild isn't cached yet
if (cache.has(guild_id)) continue;
// Make sure data is actually incomplete
if (incomplete_data["name"].isString() && incomplete_data["channels"].isArray()) {
// We'll just use it as-is
processGuild(incomplete_data);
continue;
}
// Fetch guild from API and process it
processGuild(co_await api.call(boost::beast::http::verb::get, "/guilds/"+guild_id.str()));
// Delay
co_await asyncSleep(rng.getUInt(6000, 15000));
}
}
void insertGuildUpdate(const Json::Value& data, bool is_initial) {
db << "INSERT OR IGNORE INTO guilds (id, timestamp, is_initial, name, owner_user_id)"
" VALUES (?, ?, ?, ?, ? );"
<< data["id"].asString() << std::to_string(time(nullptr)) << is_initial << data["name"].asString() << data["owner_id"].asString();
cache.store(data);
}
void insertGuildDelete(const Json::Value& data) {
auto cached_data = cache.fetch(data["id"]);
db << "INSERT OR IGNORE INTO guilds (id, timestamp, is_initial, name, owner_user_id, in_guild)"
" VALUES (?, ?, 0, ?, ?, 0 );"
<< data["id"].asString() << std::to_string(time(nullptr)) << cached_data["name"].asString() << cached_data["owner_id"].asString();
}
/*
* Messages
*/
void insertMessageContent(const Json::Value& data, bool is_initial) {
const auto& embed = data["embed"];
db << "INSERT OR IGNORE INTO message_contents (message_id, is_initial, timestamp, content, embed)"
" VALUES (?, ?, ?, ?, ? );"
<< data["id"].asString() << is_initial << std::to_string(time(nullptr)) << data["content"].asString() << (embed.isObject()?embed.toStyledString():std::optional<std::string>());
}
void insertMessageUpdate(const Json::Value& data, bool is_initial) {
if (is_initial) {
const auto& author = data["author"];
int type = data["type"].asInt();
db << "INSERT OR IGNORE INTO messages (id, type, channel_id, author_id, replied_to_id, creation_timestamp)"
" VALUES (?, ?, ?, ?, ?, ? );"
<< data["id"].asString() << type << data["channel_id"].asString() << author["id"].asString() << (type==19?data["referenced_message"]["id"].asString():std::optional<std::string>()) << std::to_string(time(nullptr));
if (!data["content"].asString().empty()) insertMessageContent(data, true);
cache.store(author);
} else {
if (!data["content"].asString().empty()) {
db << "UPDATE messages "
"SET is_edited = 1, update_timestamp = ? "
"WHERE id = ?;"
<< std::to_string(time(nullptr)) << data["id"].asString();
insertMessageContent(data, true);
}
}
}
void insertMessageDelete(const Json::Value& data) {
db << "UPDATE messages "
"SET is_deleted = 1 "
"WHERE id = ?;"
<< data["id"].asString();
}
/*
* Channels
*/
void insertChannelUpdate(const Json::Value& data, bool is_initial) {
db << "INSERT OR IGNORE INTO channels (id, guild_id, parent_id, timestamp, is_initial, type, name, topic)"
" VALUES (?, ?, ?, ?, ?, ?, ?, ? );"
<< data["id"].asString() << GetJSONAsOptionalString(data["guild_id"])<< GetJSONAsOptionalString(data["parent_id"]) << std::to_string(time(nullptr)) << is_initial << data["type"].asInt() << GetJSONAsOptionalString(data["name"]) << GetJSONAsOptionalString(data["topic"]);
cache.store(data);
}
/*
* Intent handler
*/
virtual boost::asio::awaitable<void> intentHandler(const std::string& intent, const Json::Value& data) override {
if (intent == "READY") [[unlikely]] {
const auto& user = cache.store(data["user"]);
logger.log(Loglevel::info, "Connected to Discord as: "+user["username"].asString()+'#'+user["discriminator"].asString());
settings.is_bot = user["bot"].asBool();
co_await fetchAllGuilds(data["guilds"]);
}
else if (intent == "GUILD_CREATE") [[unlikely]] {
insertGuildUpdate(data, true);
} else if (intent == "GUILD_UPDATE") [[unlikely]] {
insertGuildUpdate(data, false);
} else if (intent == "GUILD_DELETE") [[unlikely]] {
insertGuildDelete(data);
}
else if (intent == "MESSAGE_CREATE") [[likely]] {
insertMessageUpdate(data, true);
} else if (intent == "MESSAGE_UPDATE") [[likely]] {
insertMessageUpdate(data, false);
} else if (intent == "MESSAGE_DELETE") {
insertMessageDelete(data);
}
co_return;
}
};
int main(int argc, char **argv) {
// Check args
if (argc == 1) {
std::cout << "Usage: " << argv[0] << " <token>" << std::endl;
return EXIT_FAILURE;
}
// Create io service
boost::asio::io_service io;
// Set up client
auto client = std::make_shared<MyClient>(io, ChatFuse::Discord::Settings{
.bot_token = argv[1],
.intents = ChatFuse::Discord::intents::all
});
client->detach();
// Erase bot token from argv (so it's no longer visible in /proc/{pid}/cmdline)
memset(argv[1], 0, strlen(argv[1]));
// Run!!!
io.run();
return EXIT_SUCCESS;
}