generated from Oliver/discord-bot-template
Compare commits
4 Commits
Author | SHA1 | Date | |
---|---|---|---|
66e72aeb54 | |||
a701716958 | |||
8e53f2a234 | |||
9492f378d3 |
2051
Cargo.lock
generated
2051
Cargo.lock
generated
File diff suppressed because it is too large
Load Diff
@ -18,3 +18,4 @@ tracing = "0.1.40"
|
||||
tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
|
||||
serde_json = "1.0.120"
|
||||
serde = { version = "1.0.204", features = ["derive"] }
|
||||
surrealdb = "1.5.4"
|
||||
|
355
src/command.rs
355
src/command.rs
@ -1,27 +1,26 @@
|
||||
use std::{collections::HashMap, fmt::Display, fs, sync::Arc};
|
||||
use std::{collections::HashMap, fmt::Display, fs, hint::black_box, sync::Arc};
|
||||
|
||||
use crate::Context;
|
||||
use anyhow::Error;
|
||||
use poise::{serenity_prelude::{Cache, CacheHttp, ChannelId, ChannelType, GetMessages, GuildChannel, Http, Message}, CreateReply};
|
||||
use serde::Serialize;
|
||||
use poise::{
|
||||
serenity_prelude::{
|
||||
Cache, CacheHttp, ChannelId, ChannelType, GetMessages, GuildChannel, Http, Message,
|
||||
},
|
||||
CreateReply,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use surrealdb::{engine::remote::ws::Ws, opt::auth::Root, sql::Thing, Surreal};
|
||||
use tokio::time::Instant;
|
||||
use tracing::{debug, error, info, trace};
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct Server {
|
||||
channels: Vec<Channel>,
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
orphanage: Vec<GuildChannel>,
|
||||
#[serde(skip_serializing)]
|
||||
needs_clean: bool,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct Channel {
|
||||
this: GuildChannel,
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
#[serde(default)]
|
||||
children: Vec<Channel>,
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
#[serde(default)]
|
||||
messages: Vec<Message>,
|
||||
}
|
||||
|
||||
@ -34,14 +33,32 @@ impl Channel {
|
||||
messages: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_category(&self) -> bool {
|
||||
!self.children.is_empty() && self.messages.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct Server {
|
||||
#[serde(default)]
|
||||
name: String,
|
||||
channels: Vec<Channel>,
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
#[serde(default)]
|
||||
orphanage: Vec<GuildChannel>,
|
||||
#[serde(skip_serializing)]
|
||||
#[serde(default)]
|
||||
needs_clean: bool,
|
||||
}
|
||||
|
||||
impl Display for Server {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
|
||||
fn print(f: &mut std::fmt::Formatter<'_>, tab: usize, channel: &Vec<Channel>) {
|
||||
for i in channel {
|
||||
for _ in 0..tab { let _ = write!(f, "\t"); }
|
||||
for _ in 0..tab {
|
||||
let _ = write!(f, "\t");
|
||||
}
|
||||
let _ = writeln!(f, "{} {}", prefix(i.this.kind), i.this.name);
|
||||
print(f, tab + 1, &i.children);
|
||||
}
|
||||
@ -67,7 +84,6 @@ impl Display for Server {
|
||||
|
||||
print(f, 0, &self.channels);
|
||||
|
||||
|
||||
if self.needs_clean {
|
||||
let _ = writeln!(f, "Orphans: (please clean() before displaying...)");
|
||||
for i in &self.orphanage {
|
||||
@ -80,12 +96,12 @@ impl Display for Server {
|
||||
}
|
||||
|
||||
impl Server {
|
||||
|
||||
fn new() -> Self {
|
||||
fn new(name: impl Into<String>) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
channels: Vec::new(),
|
||||
orphanage: Vec::new(),
|
||||
needs_clean: false
|
||||
needs_clean: false,
|
||||
}
|
||||
}
|
||||
|
||||
@ -109,12 +125,12 @@ impl Server {
|
||||
match Self::search_by_id(&mut self.channels, parent_id) {
|
||||
Some(parent_node) => {
|
||||
parent_node.children.push(Channel::new(insert));
|
||||
},
|
||||
}
|
||||
None => {
|
||||
// couldn't find parent, store somewhere else until it's parent is added...
|
||||
self.orphanage.push(insert);
|
||||
self.needs_clean = true;
|
||||
},
|
||||
}
|
||||
}
|
||||
} else {
|
||||
self.channels.push(Channel::new(insert));
|
||||
@ -123,7 +139,9 @@ impl Server {
|
||||
|
||||
/// Cleans out the orphan channels, finding them parents. You'll want to use this before displaying anything.
|
||||
fn clean(&mut self) {
|
||||
if !self.needs_clean {return;}
|
||||
if !self.needs_clean {
|
||||
return;
|
||||
}
|
||||
|
||||
// Look thru the orphanage and try to find parents
|
||||
for orphan in &self.orphanage {
|
||||
@ -160,12 +178,20 @@ impl Server {
|
||||
}
|
||||
|
||||
/// Get all messages for 1 channel and children
|
||||
async fn get_messages(channel: &mut Channel, cache: impl CacheHttp + Clone, settings: GetMessages) {
|
||||
async fn get_messages(
|
||||
channel: &mut Channel,
|
||||
cache: impl CacheHttp + Clone,
|
||||
settings: GetMessages,
|
||||
) {
|
||||
// Loop thru all the messages in the channel in batches.
|
||||
// Adding each batch to the current channel's messages the whole time.
|
||||
let mut last_id = channel.this.last_message_id;
|
||||
while let Some(last) = last_id {
|
||||
match channel.this.messages(cache.clone(), settings.before(last)).await {
|
||||
match channel
|
||||
.this
|
||||
.messages(cache.clone(), settings.before(last))
|
||||
.await
|
||||
{
|
||||
Ok(mut ok) => {
|
||||
if ok.is_empty() {
|
||||
debug!("Reached the beginning of {}", channel.this.name);
|
||||
@ -179,12 +205,15 @@ impl Server {
|
||||
last_id = Some(l.id);
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
Err(e) => {
|
||||
error!("While reading messages in \"{}\" before `{}` - {e}", channel.this.name, last);
|
||||
error!(
|
||||
"While reading messages in \"{}\" before `{}` - {e}",
|
||||
channel.this.name, last
|
||||
);
|
||||
// Stop reading this channel on an error.
|
||||
last_id = None;
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
// Then recurse into children channels
|
||||
@ -201,17 +230,228 @@ impl Server {
|
||||
for channel in this {
|
||||
total += channel.messages.len();
|
||||
total += walk(&channel.children);
|
||||
};
|
||||
}
|
||||
total
|
||||
}
|
||||
walk(&self.channels)
|
||||
}
|
||||
|
||||
async fn to_surreal(&self) -> surrealdb::Result<()> {
|
||||
trace!("Connecting to surrealdb...");
|
||||
// Connect to the server
|
||||
let db = Surreal::new::<Ws>("127.0.0.1:8000").await?;
|
||||
db.signin(Root {
|
||||
username: "root",
|
||||
password: "root",
|
||||
})
|
||||
.await?;
|
||||
|
||||
db.use_ns("egress").use_db(self.name.clone()).await?;
|
||||
|
||||
// =========================================================
|
||||
// Ingress data
|
||||
|
||||
// Data will only be in three layers
|
||||
// Layer 1: Categories (no parent)
|
||||
// Layer 2: Channels (might have parent)
|
||||
// Layer 3: Messages (has parent)
|
||||
|
||||
trace!("Starting ingress...");
|
||||
for cat in self.channels.iter() {
|
||||
|
||||
match cat.this.kind {
|
||||
ChannelType::Text => {
|
||||
// This is a text channel
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ChannelWrapper {
|
||||
name: String,
|
||||
nsfw: bool,
|
||||
discord_id: u64,
|
||||
discord_parent_id: Option<u64>,
|
||||
topic: String,
|
||||
}
|
||||
let chan = &cat.this;
|
||||
let dpi = if let Some(val) = chan.parent_id {Some(val.get())} else { None };
|
||||
let new_channel: Vec<Thing> = db
|
||||
.create("channel")
|
||||
.content(ChannelWrapper {
|
||||
name: chan.name.to_owned(),
|
||||
nsfw: chan.nsfw,
|
||||
discord_id: chan.id.get(),
|
||||
discord_parent_id: dpi,
|
||||
topic: chan.topic.to_owned().unwrap_or(String::new()),
|
||||
})
|
||||
.await?;
|
||||
|
||||
for msg in cat.messages.iter() {
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct Author {
|
||||
nickname: String,
|
||||
username: String,
|
||||
/// B64 encoded string of image (for now)
|
||||
avatar: String,
|
||||
id: u64,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct Attachment {
|
||||
content_type: String,
|
||||
filename: String,
|
||||
url: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct Reaction {
|
||||
count: u64,
|
||||
emoji: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct MessageWrapper {
|
||||
// FIXME learn how to do references
|
||||
parent: Thing,
|
||||
author: Author,
|
||||
content: String,
|
||||
utc_timestamp: String,
|
||||
mentions: Vec<u64>,
|
||||
attachments: Vec<Attachment>,
|
||||
reactions: Vec<Reaction>,
|
||||
pinned: bool,
|
||||
}
|
||||
|
||||
let _: Vec<Thing> = db
|
||||
.create("message")
|
||||
.content(MessageWrapper {
|
||||
parent: new_channel[0].clone(),
|
||||
author: Author {
|
||||
id: msg.author.id.get(),
|
||||
nickname: msg.author.name.to_owned(),
|
||||
username: msg.author.global_name.clone().unwrap_or(String::new()),
|
||||
avatar: {
|
||||
match msg.author.avatar {
|
||||
Some(hash) => {
|
||||
format!(
|
||||
"https://cdn.discordapp.com/avatars/{}/{}.webp",
|
||||
msg.author.id,
|
||||
hash,
|
||||
)
|
||||
},
|
||||
None => String::new(),
|
||||
}
|
||||
}
|
||||
},
|
||||
content: msg.content.clone(),
|
||||
utc_timestamp: msg.timestamp.to_utc().to_string(),
|
||||
mentions: msg.mentions.iter().map(|f| f.id.get()).collect(),
|
||||
attachments: msg.attachments.iter().map(|f| Attachment {
|
||||
content_type: f.content_type.clone().unwrap_or(String::new()),
|
||||
filename: f.filename.to_owned(),
|
||||
url: f.url.to_owned(),
|
||||
}).collect(),
|
||||
reactions: msg.reactions.iter().map(|f| Reaction {
|
||||
count: f.count,
|
||||
emoji: f.reaction_type.as_data(),
|
||||
}).collect(),
|
||||
pinned: msg.pinned,
|
||||
})
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
ChannelType::Private => todo!(),
|
||||
ChannelType::Voice => todo!(),
|
||||
ChannelType::GroupDm => todo!(),
|
||||
ChannelType::Category => todo!(),
|
||||
ChannelType::News => todo!(),
|
||||
ChannelType::NewsThread => todo!(),
|
||||
ChannelType::PublicThread => todo!(),
|
||||
ChannelType::PrivateThread => todo!(),
|
||||
ChannelType::Stage => todo!(),
|
||||
ChannelType::Directory => todo!(),
|
||||
ChannelType::Forum => todo!(),
|
||||
ChannelType::Unknown(_) => todo!(),
|
||||
_ => todo!(),
|
||||
}
|
||||
|
||||
// TODO learn why this is a vec
|
||||
// Do the first iteration of channels a bit different, so as to name it "category".
|
||||
let new_category: Vec<Thing> = db.create("category").content(&cat.this).await?;
|
||||
import_messages(&cat.messages, &new_category[0], &db).await?;
|
||||
|
||||
// Ok, now automatically recurse the rest of the structure and auto import as channels
|
||||
// and messages.
|
||||
import_channel(&cat.children, &new_category[0], &db).await?;
|
||||
}
|
||||
|
||||
async fn import_channel(
|
||||
channels: &Vec<Channel>,
|
||||
parent: &Thing,
|
||||
db: &Surreal<surrealdb::engine::remote::ws::Client>,
|
||||
) -> surrealdb::Result<()> {
|
||||
for channel in channels.iter() {
|
||||
trace!("Importing channel \"{}\"", channel.this.name);
|
||||
#[derive(Serialize)]
|
||||
struct ChannelWrapper<'a, 'b> {
|
||||
channel: &'a GuildChannel,
|
||||
surreal_parent: &'b Thing,
|
||||
}
|
||||
let new_channel: Vec<Thing> = db
|
||||
.create("channel")
|
||||
.content(ChannelWrapper {
|
||||
channel: &channel.this,
|
||||
surreal_parent: &parent,
|
||||
})
|
||||
.await?;
|
||||
|
||||
import_messages(&channel.messages, &new_channel[0], &db).await?;
|
||||
// async recursion - thus box
|
||||
Box::pin(import_channel(&channel.children, &new_channel[0], &db)).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn import_messages(
|
||||
msgs: &Vec<Message>,
|
||||
parent: &Thing,
|
||||
db: &Surreal<surrealdb::engine::remote::ws::Client>,
|
||||
) -> surrealdb::Result<()> {
|
||||
trace!("Importing {} messages...", msgs.len());
|
||||
for msg in msgs.iter() {
|
||||
#[derive(Serialize)]
|
||||
struct MessageWrapper<'a, 'b> {
|
||||
message: &'a Message,
|
||||
surreal_parent: &'b Thing,
|
||||
}
|
||||
|
||||
let created: Vec<Thing> = db
|
||||
.create("message")
|
||||
.content(MessageWrapper {
|
||||
message: &msg,
|
||||
surreal_parent: &parent,
|
||||
})
|
||||
.await?;
|
||||
|
||||
trace!("Imported message {:?}", created);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// Data is all in
|
||||
// =========================================================
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[poise::command(slash_command, rename = "scrape_all", guild_only, owners_only)]
|
||||
pub async fn scrape_all(ctx: Context<'_>, pretty_print: bool) -> Result<(), Error> {
|
||||
let guild = ctx.guild_id().unwrap().to_partial_guild(ctx.serenity_context()).await.unwrap();
|
||||
let guild = ctx
|
||||
.guild_id()
|
||||
.unwrap()
|
||||
.to_partial_guild(ctx.serenity_context())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let invoker = ctx.author().name.clone();
|
||||
if let Some(nickname) = ctx.author().nick_in(ctx.http(), guild.id).await {
|
||||
@ -221,7 +461,7 @@ pub async fn scrape_all(ctx: Context<'_>, pretty_print: bool) -> Result<(), Erro
|
||||
}
|
||||
|
||||
if let Ok(map) = guild.channels(ctx.http()).await {
|
||||
let mut server = index(map).await;
|
||||
let mut server = index(map, guild.name).await;
|
||||
match ctx.reply("Starting scrape...").await {
|
||||
Ok(ok) => {
|
||||
let start = Instant::now();
|
||||
@ -229,38 +469,48 @@ pub async fn scrape_all(ctx: Context<'_>, pretty_print: bool) -> Result<(), Erro
|
||||
let end = start.elapsed().as_millis();
|
||||
let msg_count = server.message_count();
|
||||
|
||||
let print = if pretty_print {
|
||||
serde_json::to_string_pretty(&server)
|
||||
} else {
|
||||
serde_json::to_string(&server)
|
||||
if let Err(e) = server.to_surreal().await {
|
||||
error!("{e}");
|
||||
};
|
||||
|
||||
match print {
|
||||
Ok(ok) => {
|
||||
if let Err(e) = fs::write("server.json", ok) {
|
||||
error!("Problem writing server to disk: {e}");
|
||||
}
|
||||
},
|
||||
Err(err) => {
|
||||
error!("Trying to serialize server: {err}");
|
||||
},
|
||||
}
|
||||
// let print = if pretty_print {
|
||||
// serde_json::to_string_pretty(&server)
|
||||
// } else {
|
||||
// serde_json::to_string(&server)
|
||||
// };
|
||||
|
||||
// match print {
|
||||
// Ok(ok) => {
|
||||
// if let Err(e) = fs::write("server.json", ok) {
|
||||
// error!("Problem writing server to disk: {e}");
|
||||
// }
|
||||
// },
|
||||
// Err(err) => {
|
||||
// error!("Trying to serialize server: {err}");
|
||||
// },
|
||||
// }
|
||||
|
||||
// Done. Print stats.
|
||||
let _ = ok.edit(ctx, CreateReply::default().content(
|
||||
&format!("Done. Stats: \n```toml\nMessages saved: {msg_count}\nElapsed time: {end}ms\n```")
|
||||
)).await;
|
||||
debug!("Scraped server in {}ms", end);
|
||||
},
|
||||
}
|
||||
Err(e) => error!("{e} - While trying to reply to scrape command"),
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn from_json() {
|
||||
let data = fs::read_to_string("server.json").unwrap();
|
||||
let server: Server = serde_json::from_str(&data).unwrap();
|
||||
server.to_surreal().await.unwrap();
|
||||
}
|
||||
|
||||
/// Get server's topology (and runs clean)
|
||||
async fn index(map: HashMap<ChannelId, GuildChannel>) -> Server {
|
||||
let mut server = Server::new();
|
||||
async fn index(map: HashMap<ChannelId, GuildChannel>, name: impl Into<String>) -> Server {
|
||||
let mut server = Server::new(name);
|
||||
// iterate thru all channels
|
||||
map.into_iter().for_each(|(_id, current)| {
|
||||
// println!("{} {} {:?}", current.name, current.id, current.parent_id);
|
||||
@ -275,7 +525,12 @@ async fn index(map: HashMap<ChannelId, GuildChannel>) -> Server {
|
||||
// NOTE!!! Make sure these names in quotes are lowercase!
|
||||
#[poise::command(slash_command, rename = "index", guild_only)]
|
||||
pub async fn index_cmd(ctx: Context<'_>) -> Result<(), Error> {
|
||||
let guild = ctx.guild_id().unwrap().to_partial_guild(ctx.serenity_context()).await.unwrap();
|
||||
let guild = ctx
|
||||
.guild_id()
|
||||
.unwrap()
|
||||
.to_partial_guild(ctx.serenity_context())
|
||||
.await
|
||||
.unwrap();
|
||||
let invoker = ctx.author().name.clone();
|
||||
if let Some(nickname) = ctx.author().nick_in(ctx.http(), guild.id).await {
|
||||
info!("{invoker} ({nickname}) is indexing {}", guild.name);
|
||||
@ -285,9 +540,9 @@ pub async fn index_cmd(ctx: Context<'_>) -> Result<(), Error> {
|
||||
|
||||
match guild.channels(ctx.http()).await {
|
||||
Ok(ok) => {
|
||||
let server = index(ok).await;
|
||||
let server = index(ok, guild.name).await;
|
||||
let _ = ctx.reply(server.to_string()).await;
|
||||
},
|
||||
}
|
||||
Err(_) => todo!(),
|
||||
}
|
||||
Ok(())
|
||||
|
@ -15,6 +15,8 @@ static ENV: Lazy<BotEnv> = Lazy::new(|| {
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
command::from_json().await;
|
||||
return;
|
||||
|
||||
// Start the tracing subscriber
|
||||
let filter = EnvFilter::builder()
|
||||
|
Loading…
Reference in New Issue
Block a user