use std::fmt::Write; use std::fs::File; use std::io::Read; use std::net::SocketAddr; use std::thread; use std::{collections::HashMap, env}; use arrayvec::ArrayString; use futures_util::stream::StreamExt; use irc::client::prelude::Config; use irc::client::Client; use irc::proto::{ChannelExt, Command, Prefix}; use rspotify::Credentials; use serde::Deserialize; use tokio::select; use tokio::sync::mpsc::{channel, Receiver, Sender}; use tracing_log::LogTracer; use tracing_subscriber::EnvFilter; use crate::bots::{leek, misc, sed, title}; use crate::database::{DbExecutor, ExecutorConnection}; mod bots; mod database; mod web_service; // this will be displayed when the help command is used const HELP: &[&str] = &[ concat!("=- \x1d\x02Ü\x02berbot\x0f ", env!("CARGO_PKG_VERSION"), " -="), " * waifu ", " * owo/mock/leet [user]", " * ev ", " - This bot also provides titles of URLs and details for Spotify URIs/links. It can also resolve sed expressions." ]; #[cfg(unix)] async fn terminate_signal() { use tokio::signal::unix::{signal, SignalKind}; let mut sigterm = signal(SignalKind::terminate()).unwrap(); let mut sigint = signal(SignalKind::interrupt()).unwrap(); tracing::debug!("Installed ctrl+c handler"); select! { _ = sigterm.recv() => return, _ = sigint.recv() => return } } #[cfg(windows)] async fn terminate_signal() { use tokio::signal::windows::ctrl_c; let mut ctrlc = ctrl_c().unwrap(); tracing::debug!("Installed ctrl+c handler"); let _ = ctrlc.recv().await; } pub struct AppState { prefix: String, client: Client, last_msgs: HashMap, last_eval: HashMap, titlebot: title::Titlebot, db: ExecutorConnection, git_channel: String, } #[derive(Deserialize)] struct ClientConf { channels: Vec, host: String, tls: bool, mode: Option, nickname: Option, port: u16, username: String, spotify_client_id: String, spotify_client_secret: String, prefix: String, db_path: Option, http_listen: Option, git_channel: String, } #[tokio::main(flavor = "current_thread")] async fn main() -> anyhow::Result<()> { LogTracer::init()?; tracing_subscriber::fmt() .with_env_filter(EnvFilter::from_env("UBERBOT_LOG")) .init(); let mut file = File::open(env::var("UBERBOT_CONFIG").unwrap_or_else(|_| "uberbot.toml".to_string()))?; let mut client_conf = String::new(); file.read_to_string(&mut client_conf)?; let client_config: ClientConf = toml::from_str(&client_conf)?; let (db_exec, db_conn) = DbExecutor::create(client_config.db_path.as_deref().unwrap_or("uberbot.db3"))?; let exec_thread = thread::spawn(move || { db_exec.run(); tracing::info!("Database executor has been shut down"); }); let spotify_creds = Credentials::new( &client_config.spotify_client_id, &client_config.spotify_client_secret, ); let http_listen = client_config .http_listen .unwrap_or_else(|| SocketAddr::from(([127, 0, 0, 1], 5000))); let uber_ver = concat!("Überbot ", env!("CARGO_PKG_VERSION")); let irc_config = Config { nickname: client_config.nickname, username: Some(client_config.username.clone()), realname: Some(client_config.username), server: Some(client_config.host), port: Some(client_config.port), use_tls: Some(client_config.tls), channels: client_config.channels, umodes: client_config.mode, user_info: Some(uber_ver.into()), version: Some(uber_ver.into()), ..Config::default() }; let client = Client::from_config(irc_config).await?; client.identify()?; let state = AppState { prefix: client_config.prefix, client, last_msgs: HashMap::new(), last_eval: HashMap::new(), titlebot: title::Titlebot::create(spotify_creds).await?, db: db_conn, git_channel: client_config.git_channel, }; let (git_tx, git_recv) = channel(512); if let Err(e) = executor(state, git_tx, git_recv, http_listen).await { tracing::error!("Error in message loop: {}", e); } if let Err(e) = exec_thread.join() { tracing::error!("Error while shutting down the database: {:?}", e); } tracing::info!("Shutting down"); Ok(()) } async fn executor( mut state: AppState, git_tx: Sender, mut git_recv: Receiver, http_listen: SocketAddr, ) -> anyhow::Result<()> { let web_db = state.db.clone(); select! { r = web_service::run(web_db, git_tx, http_listen) => r?, r = message_loop(&mut state) => r?, r = git_recv.recv() => { if let Some(message) = r { state.client.send_privmsg(&state.git_channel, &message)?; } } _ = terminate_signal() => { tracing::info!("Sending QUIT message"); state.client.send_quit("überbot shutting down")?; } } Ok(()) } async fn message_loop(state: &mut AppState) -> anyhow::Result<()> { let mut stream = state.client.stream()?; while let Some(message) = stream.next().await.transpose()? { if let Command::PRIVMSG(ref origin, content) = message.command { if origin.is_channel_name() { if let Some(author) = message.prefix.as_ref().and_then(|p| match p { Prefix::Nickname(name, _, _) => Some(&name[..]), _ => None, }) { if let Err(e) = handle_privmsg(state, author, origin, content).await { state .client .send_privmsg(origin, &format!("Error: {}", e))?; } } else { tracing::warn!("Couldn't get the author for a message"); } } } } Ok(()) } fn separate_to_space(str: &str, prefix_len: usize) -> (&str, Option<&str>) { if let Some(o) = str.find(' ') { (&str[prefix_len..o], Some(&str[o + 1..])) } else { (&str[prefix_len..], None) } } async fn handle_privmsg( state: &mut AppState, author: &str, origin: &str, content: String, ) -> anyhow::Result<()> { if !content.starts_with(state.prefix.as_str()) { if let Some(titlebot_msg) = state.titlebot.resolve(&content).await? { state.client.send_privmsg(origin, &titlebot_msg)?; } if let Some(prev_msg) = state.last_msgs.get(author) { if let Some(formatted) = sed::resolve(prev_msg, &content)? { let mut result = ArrayString::<512>::new(); write!(result, "<{}> {}", author, formatted)?; state.client.send_privmsg(origin, &result)?; state.last_msgs.insert(author.into(), formatted.to_string()); return Ok(()); } } state.last_msgs.insert(author.into(), content); return Ok(()); } let (command, remainder) = separate_to_space(&content, state.prefix.len()); tracing::debug!("Command received ({:?}; {:?})", command, remainder); match command { "help" => { for help_line in HELP { state.client.send_privmsg(origin, help_line)?; } } "waifu" => { let category = remainder.unwrap_or("waifu"); let url = misc::get_waifu_pic(category).await?; let response = url .as_ref() .map(|v| v.as_str()) .unwrap_or("Invalid category. Valid categories: https://waifu.pics/docs"); state.client.send_privmsg(origin, response)?; } "mock" => { leek::execute_leek( state, leek::LeekCommand::Mock, origin, remainder.unwrap_or(author), )?; } "leet" => { leek::execute_leek( state, leek::LeekCommand::Leet, origin, remainder.unwrap_or(author), )?; } "owo" => { leek::execute_leek( state, leek::LeekCommand::Owo, origin, remainder.unwrap_or(author), )?; } "ev" => { let result = misc::mathbot(author.into(), remainder, &mut state.last_eval)?; state.client.send_privmsg(origin, &result)?; } "grab" => { if let Some(target) = remainder { if target == author { state .client .send_privmsg(target, "You can't grab yourself")?; return Ok(()); } if let Some(prev_msg) = state.last_msgs.get(target) { if state.db.add_quote(prev_msg.clone(), target.into()).await { state.client.send_privmsg(target, "Quote added")?; } else { state .client .send_privmsg(target, "A database error has occurred")?; } } else { state .client .send_privmsg(target, "No previous messages to grab")?; } } else { state.client.send_privmsg(origin, "No nickname to grab")?; } } "quot" => { if let Some(quote) = state.db.get_quote(remainder.map(|v| v.to_string())).await { let mut resp = ArrayString::<512>::new(); write!(resp, "\"{}\" ~{}", quote.0, quote.1)?; state.client.send_privmsg(origin, &resp)?; } else { state.client.send_privmsg(origin, "No quotes found")?; } } _ => { state.client.send_privmsg(origin, "Unknown command")?; } } Ok(()) }