chat_log_downloader.rs

#![feature(dedup_by)]

extern crate csv;
extern crate serenity;

use serenity::Client;
use serenity::client::Context;
use serenity::model::{GuildId, Message, MessageId, UserId};
use serenity::utils::builder::{SortingMode, SortingOrder};
use std::{env, thread};
use std::collections::HashSet;
use std::path::Path;
use std::time::Duration;

const SEARCH_DELAY: u64 = 250;

fn main() {
    let token = env::var("DISCORD_USER_TOKEN").expect("Missing env var DISCORD_USER_TOKEN");
    let args: Vec<String> = env::args().collect();
    if args.len() != 4 {
        panic!("Invalid number of arguments");
    }
    let author_id = args[1]
        .parse::<u64>()
        .expect("Failed to parse author ID arg")
        .into();
    let guild_id = args[2]
        .parse::<u64>()
        .expect("Failed to parse guild ID arg")
        .into();
    let file_path = args[3].clone();

    let mut client = Client::login_user(&token);

    client.on_ready(move |context, ready| {
        println!(
            "Started as {}#{}, serving {} guilds",
            ready.user.name,
            ready.user.discriminator,
            ready.guilds.len(),
        );

        let mut messages = perform_search(&context, &guild_id, &author_id)
            .expect("Failed to perform search");
        println!("Message count: {}", messages.len());

        write_messages_to_file(file_path.clone(), messages.clone());

        messages.sort_by_key(|ref m| m.timestamp.clone());
        messages.dedup_by_key(|ref mut m| m.id);
        println!("Unique messages: {}", messages.len());

        // TODO: make this idiomatic.
        panic!("Quitting");
    });

    if let Err(err) = client.start() {
        println!("Client error: {:?}", err);
    }
}

fn write_messages_to_file<P>(path: P, messages: Vec<Message>)
    where P: AsRef<Path>,
{
    let mut writer = csv::Writer::from_file(path).expect("Failed to build writer");
    for record in messages.into_iter().map(|m| (m.timestamp, m.id.0, m.content)) {
        let result = writer.encode(record);
        assert!(result.is_ok());
    }
}

fn perform_search(context: &Context, guild_id: &GuildId, author_id: &UserId) -> Option<Vec<Message>> {
    let offset_search = |context, guild_id, author_id, offset, max_id| {
        Context::search_guild(context, guild_id, Vec::new(), |s| {
                let s = if let Some(max_id) = max_id { s.max_id(max_id) } else { s };
                s.author_id(author_id)
                    .context_size(0)
                    .offset(offset)
                    .sort_by(SortingMode::Timestamp)
                    .sort_order(SortingOrder::Descending)
            })
            .expect("Search failed")
    };

    let num_results = offset_search(context, *guild_id, *author_id, 0, None).total;

    println!("Found {} results", num_results);

    let mut messages: Vec<Message> = Vec::new();
    let mut results_left = num_results;
    while results_left > 0 {
        let max_id: Option<MessageId> = if messages.is_empty() {
            None
        } else {
            Some(messages[messages.len() - 1].id)
        };

        // We want to go through every offset from 0 to 5000.
        let mut offset = 0;
        while offset <= 5000 {
            // Sleep for one second.
            thread::sleep(Duration::from_millis(SEARCH_DELAY));
            let result = offset_search(context, *guild_id, *author_id, offset, max_id);
            println!("Performed search with offset: {} -- Messages left: {}", offset, results_left);
            for message_set in result.results {
                for message in message_set {
                    println!("TS: {} -- {}", message.timestamp, message.id.0);
                    //messages.push((message.timestamp.clone(), message.id.0, message));
                    messages.push(message);
                    results_left -= 1;

                    if results_left == 0 {
                        break;
                    }
                }
            }

            if results_left == 0 {
                break;
            }

            offset += 25;
        }
    }

    Some(messages)
}