Oliver/discord-egress
1
0
Fork 0
generated from Oliver/discord-bot-template
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Compare commits

base: Oliver:66e72aeb54c1aa200e15b42ffab7b4c4719a0e1d
Branches Tags
Oliver:main
Oliver:surreal
..
compare: Oliver:c5cb4eb524300e87c1194b1ab33f01e5bf34134e
Branches Tags
Oliver:surreal
Oliver:main

No commits in common. "66e72aeb54c1aa200e15b42ffab7b4c4719a0e1d" and "c5cb4eb524300e87c1194b1ab33f01e5bf34134e" have entirely different histories.
66e72aeb54 ... c5cb4eb524

4 changed files with 80 additions and 2353 deletions
Show Stats Expand all files Collapse all files

2051
Cargo.lock generated
View File

File diff suppressed because it is too large Load Diff

1
Cargo.toml
View File

@ -18,4 +18,3 @@ 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
View File

@ -1,26 +1,27 @@
use std::{collections::HashMap, fmt::Display, fs, hint::black_box, sync::Arc};
use std::{collections::HashMap, fmt::Display, fs, sync::Arc};
use crate::Context;
use anyhow::Error;
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 poise::{serenity_prelude::{Cache, CacheHttp, ChannelId, ChannelType, GetMessages, GuildChannel, Http, Message}, CreateReply};
use serde::Serialize;
use tokio::time::Instant;
use tracing::{debug, error, info, trace};
#[derive(Serialize, Deserialize)]
#[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)]
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>,
}
@ -33,32 +34,14 @@ 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);
}
@ -84,6 +67,7 @@ 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 {
@ -96,12 +80,12 @@ impl Display for Server {
}
impl Server {
fn new(name: impl Into<String>) -> Self {
fn new() -> Self {
Self {
name: name.into(),
channels: Vec::new(),
orphanage: Vec::new(),
needs_clean: false,
needs_clean: false
}
}
@ -125,12 +109,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));
@ -139,9 +123,7 @@ 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 {
@ -178,20 +160,12 @@ 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);
@ -205,15 +179,12 @@ 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
@ -230,228 +201,17 @@ 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 {
@ -461,7 +221,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, guild.name).await;
let mut server = index(map).await;
match ctx.reply("Starting scrape...").await {
Ok(ok) => {
let start = Instant::now();
@ -469,48 +229,38 @@ pub async fn scrape_all(ctx: Context<'_>, pretty_print: bool) -> Result<(), Erro
let end = start.elapsed().as_millis();
let msg_count = server.message_count();
if let Err(e) = server.to_surreal().await {
error!("{e}");
let print = if pretty_print {
serde_json::to_string_pretty(&server)
} else {
serde_json::to_string(&server)
};
// 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}");
// },
// }
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>, name: impl Into<String>) -> Server {
let mut server = Server::new(name);
async fn index(map: HashMap<ChannelId, GuildChannel>) -> Server {
let mut server = Server::new();
// iterate thru all channels
map.into_iter().for_each(|(_id, current)| {
// println!("{} {} {:?}", current.name, current.id, current.parent_id);
@ -525,12 +275,7 @@ async fn index(map: HashMap<ChannelId, GuildChannel>, name: impl Into<String>) -
// 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);
@ -540,9 +285,9 @@ pub async fn index_cmd(ctx: Context<'_>) -> Result<(), Error> {
match guild.channels(ctx.http()).await {
Ok(ok) => {
let server = index(ok, guild.name).await;
let server = index(ok).await;
let _ = ctx.reply(server.to_string()).await;
}
},
Err(_) => todo!(),
}
Ok(())

2
src/main.rs
View File

@ -15,8 +15,6 @@ static ENV: Lazy<BotEnv> = Lazy::new(|| {
#[tokio::main]
async fn main() {
command::from_json().await;
return;
// Start the tracing subscriber
let filter = EnvFilter::builder()