is-this-a-repost-bot/src/main.rs

mod async_error;
mod db;
mod media_hash;
mod sql_types;

use async_error::AsyncError;
use clap::Parser as _;
use db::{get_db, Db};
use media_hash::MediaHash;
use teloxide::{
    dispatching::UpdateFilterExt,
    net::Download as _,
    prelude::*,
    types::{Chat, ChatMemberStatus, MediaKind, MessageCommon, MessageKind},
};

#[derive(clap::Parser)]
struct CommandLineArgs {
    /// Path to the database file
    #[clap(short, long)]
    db_path: String,
}

#[tokio::main]
async fn main() -> Result<(), AsyncError> {
    let args = CommandLineArgs::parse();

    pretty_env_logger::init();
    log::info!("Starting bot");

    // open up the database
    let db = get_db(&args.db_path).await?;
    db.print_info().await?;

    let bot = Bot::from_env();

    let handler = dptree::entry()
        .branch(Update::filter_my_chat_member().endpoint(handle_bot_status_change))
        .branch(Update::filter_channel_post().endpoint(handle_channel_message))
        .branch(Update::filter_message().endpoint(handle_direct_message));

    log::info!("Starting dispatcher");
    Dispatcher::builder(bot, handler)
        .enable_ctrlc_handler()
        .dependencies(dptree::deps![db])
        .build()
        .dispatch()
        .await;
    Ok(())
}

fn chat_name(chat: &Chat) -> &str {
    chat.username().or(chat.title()).unwrap_or("<no name>")
}

async fn handle_bot_status_change(updated: ChatMemberUpdated, db: Db) -> Result<(), AsyncError> {
    let status = updated.new_chat_member.status();
    let by_user = updated.from.id;
    let chat_id = updated.chat.id;
    log::info!(
        "Bot status in Chat({}) changed by User({}) to {:?}",
        chat_id,
        by_user,
        status
    );
    match status {
        ChatMemberStatus::Administrator => {
            log::info!(
                "Bot is now an administrator of Chat({}, {})",
                chat_id,
                chat_name(&updated.chat)
            );
            db.insert_channel_ownership(by_user, chat_id).await?;
        }
        ChatMemberStatus::Left => {
            log::info!(
                "Bot is no longer in Chat({}, {})",
                chat_id,
                chat_name(&updated.chat)
            );
            db.delete_channel_ownership(by_user, chat_id).await?;
        }
        _ => {}
    }
    Ok(())
}

async fn handle_direct_message(bot: Bot, message: Message, db: Db) -> Result<(), AsyncError> {
    let from_user = message.from.clone().ok_or("No sender in direct message")?;
    log::info!(
        "Got direct message from User({}, {}): {}",
        from_user.id,
        from_user.username.as_deref().unwrap_or("<no username>"),
        message.text().unwrap_or("<no text>")
    );

    let media_hash = download_message_photo(&bot, &message).await?;
    if let Some(media_hash) = media_hash {
        let results = db
            .search_chat_messages(media_hash, from_user.id, 1, 1.0)
            .await?;

        log::info!("Found {} similar medias", results.len());

        bot.send_message(
            message.chat.id,
            format!("Found {} similar medias", results.len()),
        )
        .await?;

        for result in results {
            log::info!("  {:?}", result);

            bot.forward_message(message.chat.id, result.chat_id, result.message_id)
                .await?;
        }
    }

    Ok(())
}

async fn handle_channel_message(bot: Bot, message: Message, db: Db) -> Result<(), AsyncError> {
    let chat_id = message.chat.id;
    let media_hash = match download_message_photo(&bot, &message).await? {
        Some(media_hash) => media_hash,
        None => return Ok(()),
    };

    let media_id = db
        .insert_chat_message(chat_id, message.id, media_hash)
        .await?;

    log::info!("Inserted message: {:?}", media_id);

    Ok(())
}

async fn download_message_photo(
    bot: &Bot,
    message: &Message,
) -> Result<Option<MediaHash>, AsyncError> {
    let photo = match &message.kind {
        MessageKind::Common(MessageCommon {
            media_kind: MediaKind::Photo(photo),
            ..
        }) => {
            let photo = match photo.photo.iter().max_by_key(|p| p.file.size) {
                Some(photo) => photo,
                None => {
                    log::info!("No photo found in message");
                    return Ok(None);
                }
            };

            photo
        }
        _ => return Ok(None),
    };

    log::info!("Downloading photo {}...", photo.file.id,);
    let file_path = bot.get_file(&photo.file.id).await?.path;
    let mut dst = Vec::new();
    bot.download_file(&file_path, &mut dst).await?;
    log::info!(
        "Downloaded {}.",
        humansize::format_size(dst.len(), humansize::BINARY)
    );
    Ok(Some(MediaHash::from_bytes(&dst)?))
}