mara/src/main.rs

#[macro_use]
extern crate tracing;

use axum::{
    extract::Extension,
    http::StatusCode,
    routing::{get, post},
    AddExtensionLayer, Json, Router,
};
use color_eyre::Result;
use futures::prelude::*;
use irc::{client::prelude::*, proto::message::Tag};
use rusqlite::{params, Connection};
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, env, net::SocketAddr, path::PathBuf, sync::Arc};
use tokio::sync::Mutex;
use tower_http::trace::TraceLayer;

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Config {
    pub database_path: PathBuf,
    pub addr: SocketAddr,
    pub announce_to: String,
    pub irc: irc::client::data::Config,
}

pub const APPLICATION_NAME: &str = concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION"),);

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt::init();
    color_eyre::install()?;

    let mut cfg: Config = serde_yaml::from_str(&std::fs::read_to_string("config.yaml")?)?;
    cfg.irc.version = Some(APPLICATION_NAME.to_string());

    let db = Arc::new(Mutex::new({
        let db = Connection::open(cfg.database_path.clone())?;
        db.execute_batch(include_str!("./schema.sql"))?;
        Ok::<Connection, color_eyre::eyre::ErrReport>(db)
    }?));

    let mut cli = Client::from_config(cfg.irc.clone()).await?;
    cli.send_cap_req(&[
        Capability::EchoMessage,
        Capability::AccountTag,
        Capability::ChgHost,
        Capability::ServerTime,
        Capability::Sasl,
    ])?;

    cli.identify()?;
    let cfg = Arc::new(cfg);

    let middleware = tower::ServiceBuilder::new()
        .layer(TraceLayer::new_for_http())
        .layer(AddExtensionLayer::new(cfg.clone()))
        .layer(AddExtensionLayer::new(cli.sender()))
        .layer(AddExtensionLayer::new(db.clone()));
    let app = Router::new()
        .route("/.within/health", get(health))
        .route("/.within/poke", post(poke_message))
        .layer(middleware);
    tokio::spawn(axum::Server::bind(&cfg.addr).serve(app.into_make_service()));

    let mut stream = cli.stream()?;
    while let Some(message) = stream.next().await.transpose()? {
        if let Err(err) = handle_message(db.clone(), cfg.clone(), message).await {
            error!("{}", err);
        }
    }

    Ok(())
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct DiscordWebhook {
    content: String,
    channel: Option<String>,
}

async fn poke_message(
    Extension(cfg): Extension<Arc<Config>>,
    Extension(sender): Extension<Sender>,
    Json(wh): Json<DiscordWebhook>,
) -> (StatusCode, String) {
    match sender.send_privmsg(wh.channel.unwrap_or(cfg.announce_to.clone()), wh.content) {
        Ok(_) => (StatusCode::NO_CONTENT, "".to_string()),
        Err(err) => {
            error!("error sending message: {}", err);
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                "internal server error sending to IRC".to_string(),
            )
        }
    }
}

async fn health(Extension(db): Extension<Arc<Mutex<Connection>>>) -> &'static str {
    let db = db.lock().await;
    match db.query_row("SELECT 1+1", [], |_| Ok(())) {
        Ok(_) => "OK",
        Err(why) => {
            error!("{}", why);
            "NOT OK"
        }
    }
}

async fn handle_message(db: Arc<Mutex<Connection>>, cfg: Arc<Config>, msg: Message) -> Result<()> {
    let db = db.lock().await;
    let msg2 = msg.clone();
    let nick = msg2.source_nickname().unwrap_or("");
    let tags = tags_map(msg2.tags.clone());
    match msg.command {
        Command::PRIVMSG(channel, message) => {
            db.execute(
                    "INSERT INTO messages(channel, nick, event, message, tags) VALUES (?1, ?2, 'message', ?3, ?4)",
                    params![
                        channel,
                        nick,
                        message,
                        serde_json::to_value(tags)?
                    ],
                )?;
        }
        Command::KICK(channel, user, reason) => {
            db.execute(
                    "INSERT INTO messages(channel, nick, event, message, tags) VALUES (?1, ?2, 'kick', ?3, ?4)",
                    params![
                        channel,
                        user,
                        format!("was kicked by {}: {}", nick, reason.unwrap_or("no reason given".to_string())),
                        serde_json::to_value(tags)?,
                    ],
                )?;
        }
        Command::PART(channel, reason) => {
            db.execute(
                    "INSERT INTO messages(channel, nick, event, message, tags) VALUES (?1, ?2, 'part', ?3, ?4)",
                    params![channel, nick, reason, serde_json::to_value(tags)?],
                )?;
        }
        Command::JOIN(channel, _, _) => {
            db.execute(
                "INSERT INTO messages(channel, nick, event, tags) VALUES (?1, ?2, 'join', ?3)",
                params![channel, nick, serde_json::to_value(tags)?],
            )?;
        }
        Command::QUIT(message) => {
            db.execute(
                "INSERT INTO messages(channel, nick, event, message tags) VALUES (?1, ?2, 'quit', ?3, ?4)",
                params![cfg.announce_to, nick, message, serde_json::to_value(tags)?],
            )?;
        }
        _ => {}
    }
    Ok(())
}

fn tags_map(tags: Option<Vec<Tag>>) -> HashMap<String, Option<String>> {
    let mut result = HashMap::<String, Option<String>>::new();

    if let None = tags {
        return result;
    }

    for tag in tags.unwrap() {
        result.insert(tag.0, tag.1);
    }

    result
}