2022-01-26 12:34:54 -06:00
|
|
|
#![allow(clippy::match_wildcard_for_single_variants)]
|
|
|
|
|
2022-01-02 14:58:54 -06:00
|
|
|
use std::fmt::Write;
|
2021-12-27 18:18:39 -06:00
|
|
|
use std::fs::File;
|
|
|
|
use std::io::Read;
|
2022-01-13 10:52:39 -06:00
|
|
|
use std::net::SocketAddr;
|
2022-01-26 12:17:07 -06:00
|
|
|
use std::sync::Arc;
|
2022-01-02 14:58:54 -06:00
|
|
|
use std::thread;
|
2021-12-27 18:18:39 -06:00
|
|
|
use std::{collections::HashMap, env};
|
2022-01-26 05:58:49 -06:00
|
|
|
|
|
|
|
use arrayvec::ArrayString;
|
|
|
|
use futures_util::stream::StreamExt;
|
|
|
|
use irc::client::prelude::Config;
|
2022-01-26 12:17:07 -06:00
|
|
|
use irc::client::{Client, ClientStream};
|
2022-01-26 05:58:49 -06:00
|
|
|
use irc::proto::{ChannelExt, Command, Prefix};
|
|
|
|
use rspotify::Credentials;
|
|
|
|
use serde::Deserialize;
|
2022-01-26 12:17:07 -06:00
|
|
|
use tokio::sync::broadcast;
|
2022-01-27 17:44:50 -06:00
|
|
|
use tokio::sync::mpsc::unbounded_channel;
|
2021-12-27 06:39:01 -06:00
|
|
|
use tracing_subscriber::EnvFilter;
|
|
|
|
|
2022-01-26 05:58:49 -06:00
|
|
|
use crate::bots::{leek, misc, sed, title};
|
2022-01-26 17:58:00 -06:00
|
|
|
use crate::database::{DbExecutor, ExecutorConnection, Quote};
|
2022-01-26 05:58:49 -06:00
|
|
|
|
|
|
|
mod bots;
|
|
|
|
mod database;
|
|
|
|
mod web_service;
|
|
|
|
|
2021-12-31 17:51:54 -06:00
|
|
|
// this will be displayed when the help command is used
|
|
|
|
const HELP: &[&str] = &[
|
|
|
|
concat!("=- \x1d\x02Ü\x02berbot\x0f ", env!("CARGO_PKG_VERSION"), " -="),
|
|
|
|
" * waifu <category>",
|
|
|
|
" * owo/mock/leet [user]",
|
|
|
|
" * ev <math expression>",
|
2022-01-01 06:38:43 -06:00
|
|
|
" - This bot also provides titles of URLs and details for Spotify URIs/links. It can also resolve sed expressions."
|
2021-12-31 17:51:54 -06:00
|
|
|
];
|
2021-12-27 12:08:46 -06:00
|
|
|
|
2021-12-27 06:39:01 -06:00
|
|
|
#[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();
|
2021-12-31 17:51:54 -06:00
|
|
|
tracing::debug!("Installed ctrl+c handler");
|
2021-12-27 06:39:01 -06:00
|
|
|
select! {
|
2021-12-27 18:18:39 -06:00
|
|
|
_ = sigterm.recv() => return,
|
|
|
|
_ = sigint.recv() => return
|
2021-12-27 06:39:01 -06:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
#[cfg(windows)]
|
|
|
|
async fn terminate_signal() {
|
|
|
|
use tokio::signal::windows::ctrl_c;
|
|
|
|
let mut ctrlc = ctrl_c().unwrap();
|
2021-12-31 17:51:54 -06:00
|
|
|
tracing::debug!("Installed ctrl+c handler");
|
2021-12-27 06:39:01 -06:00
|
|
|
let _ = ctrlc.recv().await;
|
|
|
|
}
|
|
|
|
|
2022-01-02 14:58:54 -06:00
|
|
|
pub struct AppState {
|
2021-12-27 14:37:50 -06:00
|
|
|
prefix: String,
|
2022-01-26 12:17:07 -06:00
|
|
|
client: Arc<Client>,
|
|
|
|
stream: ClientStream,
|
2021-12-27 14:37:50 -06:00
|
|
|
last_msgs: HashMap<String, String>,
|
2021-12-31 17:51:54 -06:00
|
|
|
last_eval: HashMap<String, f64>,
|
2022-01-26 05:58:49 -06:00
|
|
|
titlebot: title::Titlebot,
|
2022-01-02 14:58:54 -06:00
|
|
|
db: ExecutorConnection,
|
2021-12-27 18:18:39 -06:00
|
|
|
}
|
|
|
|
|
|
|
|
#[derive(Deserialize)]
|
2021-12-28 06:38:50 -06:00
|
|
|
struct ClientConf {
|
|
|
|
channels: Vec<String>,
|
|
|
|
host: String,
|
2022-01-26 05:58:49 -06:00
|
|
|
tls: bool,
|
2021-12-28 06:38:50 -06:00
|
|
|
mode: Option<String>,
|
|
|
|
nickname: Option<String>,
|
|
|
|
port: u16,
|
|
|
|
username: String,
|
2021-12-27 18:18:39 -06:00
|
|
|
spotify_client_id: String,
|
|
|
|
spotify_client_secret: String,
|
2021-12-28 06:43:34 -06:00
|
|
|
prefix: String,
|
2022-01-02 14:58:54 -06:00
|
|
|
db_path: Option<String>,
|
2022-01-15 10:09:46 -06:00
|
|
|
http_listen: Option<SocketAddr>,
|
2022-01-13 10:52:39 -06:00
|
|
|
git_channel: String,
|
2021-12-27 14:37:50 -06:00
|
|
|
}
|
|
|
|
|
2022-01-26 12:17:07 -06:00
|
|
|
#[tokio::main]
|
2021-12-27 06:39:01 -06:00
|
|
|
async fn main() -> anyhow::Result<()> {
|
2021-12-27 12:08:46 -06:00
|
|
|
tracing_subscriber::fmt()
|
|
|
|
.with_env_filter(EnvFilter::from_env("UBERBOT_LOG"))
|
|
|
|
.init();
|
2021-12-27 06:39:01 -06:00
|
|
|
|
2022-01-02 14:58:54 -06:00
|
|
|
let mut file =
|
|
|
|
File::open(env::var("UBERBOT_CONFIG").unwrap_or_else(|_| "uberbot.toml".to_string()))?;
|
2021-12-28 06:38:50 -06:00
|
|
|
let mut client_conf = String::new();
|
2021-12-28 16:41:51 -06:00
|
|
|
file.read_to_string(&mut client_conf)?;
|
2021-12-27 18:18:39 -06:00
|
|
|
|
2021-12-28 16:41:51 -06:00
|
|
|
let client_config: ClientConf = toml::from_str(&client_conf)?;
|
2021-12-27 18:18:39 -06:00
|
|
|
|
2022-01-02 14:58:54 -06:00
|
|
|
let (db_exec, db_conn) =
|
|
|
|
DbExecutor::create(client_config.db_path.as_deref().unwrap_or("uberbot.db3"))?;
|
2022-01-01 16:03:15 -06:00
|
|
|
let exec_thread = thread::spawn(move || {
|
2022-01-01 15:35:27 -06:00
|
|
|
db_exec.run();
|
|
|
|
tracing::info!("Database executor has been shut down");
|
|
|
|
});
|
|
|
|
|
2021-12-28 06:43:34 -06:00
|
|
|
let spotify_creds = Credentials::new(
|
|
|
|
&client_config.spotify_client_id,
|
|
|
|
&client_config.spotify_client_secret,
|
|
|
|
);
|
2021-12-27 18:18:39 -06:00
|
|
|
|
2022-01-13 10:52:39 -06:00
|
|
|
let http_listen = client_config
|
2022-01-15 10:09:46 -06:00
|
|
|
.http_listen
|
2022-01-13 10:52:39 -06:00
|
|
|
.unwrap_or_else(|| SocketAddr::from(([127, 0, 0, 1], 5000)));
|
|
|
|
|
2022-01-26 05:58:49 -06:00
|
|
|
let uber_ver = concat!("Überbot ", env!("CARGO_PKG_VERSION"));
|
|
|
|
let irc_config = Config {
|
2022-01-26 12:17:07 -06:00
|
|
|
nickname: Some(
|
|
|
|
client_config
|
|
|
|
.nickname
|
|
|
|
.unwrap_or_else(|| client_config.username.clone()),
|
|
|
|
),
|
2022-01-26 05:58:49 -06:00
|
|
|
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()
|
|
|
|
};
|
2022-01-26 12:17:07 -06:00
|
|
|
let mut client = Client::from_config(irc_config).await?;
|
|
|
|
let stream = client.stream()?;
|
2022-01-26 05:58:49 -06:00
|
|
|
client.identify()?;
|
2022-01-26 12:17:07 -06:00
|
|
|
let client = Arc::new(client);
|
|
|
|
|
|
|
|
let (ctx, _) = broadcast::channel(1);
|
|
|
|
let (etx, mut erx) = unbounded_channel();
|
|
|
|
|
|
|
|
let web_task = tokio::spawn(web_service::run(
|
|
|
|
db_conn.clone(),
|
|
|
|
client.clone(),
|
|
|
|
client_config.git_channel,
|
|
|
|
http_listen,
|
2022-01-27 17:44:50 -06:00
|
|
|
ctx.subscribe(),
|
2022-01-26 12:17:07 -06:00
|
|
|
));
|
2021-12-27 06:39:01 -06:00
|
|
|
|
2021-12-27 14:37:50 -06:00
|
|
|
let state = AppState {
|
2021-12-28 06:43:34 -06:00
|
|
|
prefix: client_config.prefix,
|
2022-01-26 12:17:07 -06:00
|
|
|
client: client.clone(),
|
|
|
|
stream,
|
2021-12-27 14:37:50 -06:00
|
|
|
last_msgs: HashMap::new(),
|
2021-12-31 17:51:54 -06:00
|
|
|
last_eval: HashMap::new(),
|
2022-01-26 05:58:49 -06:00
|
|
|
titlebot: title::Titlebot::create(spotify_creds).await?,
|
2022-01-02 14:58:54 -06:00
|
|
|
db: db_conn,
|
2021-12-27 14:37:50 -06:00
|
|
|
};
|
2022-01-26 12:17:07 -06:00
|
|
|
let message_loop_task = tokio::spawn(async move {
|
|
|
|
if let Err(e) = message_loop(state).await {
|
2022-01-26 12:34:54 -06:00
|
|
|
let _err = etx.send(e);
|
2022-01-15 11:18:26 -06:00
|
|
|
}
|
2022-01-26 12:17:07 -06:00
|
|
|
});
|
|
|
|
|
|
|
|
tokio::select! {
|
2021-12-31 17:51:54 -06:00
|
|
|
_ = terminate_signal() => {
|
2022-01-26 12:17:07 -06:00
|
|
|
tracing::info!("Received shutdown signal, sending QUIT message");
|
|
|
|
client.send_quit("überbot shutting down")?;
|
|
|
|
}
|
|
|
|
e = erx.recv() => {
|
|
|
|
if let Some(e) = e {
|
|
|
|
tracing::error!("An error has occurred, shutting down: {}", e);
|
|
|
|
} else {
|
|
|
|
tracing::error!("Error channel has been dropped due to an unknown error, shutting down");
|
|
|
|
}
|
2021-12-27 06:39:01 -06:00
|
|
|
}
|
|
|
|
}
|
2022-01-26 12:17:07 -06:00
|
|
|
|
|
|
|
tracing::info!("Closing services...");
|
|
|
|
let _ = ctx.send(());
|
|
|
|
web_task
|
|
|
|
.await
|
|
|
|
.unwrap_or_else(|e| tracing::warn!("Couldn't join the web service: {:?}", e));
|
|
|
|
message_loop_task
|
|
|
|
.await
|
|
|
|
.unwrap_or_else(|e| tracing::warn!("Couldn't join the web service: {:?}", e));
|
|
|
|
exec_thread
|
|
|
|
.join()
|
|
|
|
.unwrap_or_else(|e| tracing::warn!("Couldn't join the database: {:?}", e));
|
|
|
|
tracing::info!("Shutdown complete!");
|
|
|
|
|
2021-12-27 12:08:46 -06:00
|
|
|
Ok(())
|
|
|
|
}
|
2021-12-27 06:39:01 -06:00
|
|
|
|
2022-01-26 12:17:07 -06:00
|
|
|
async fn message_loop(mut state: AppState) -> anyhow::Result<()> {
|
|
|
|
while let Some(message) = state.stream.next().await.transpose()? {
|
2022-01-26 05:58:49 -06:00
|
|
|
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,
|
|
|
|
}) {
|
2022-01-26 12:17:07 -06:00
|
|
|
if let Err(e) = handle_privmsg(&mut state, author, origin, content).await {
|
2022-01-26 05:58:49 -06:00
|
|
|
state
|
|
|
|
.client
|
|
|
|
.send_privmsg(origin, &format!("Error: {}", e))?;
|
|
|
|
}
|
|
|
|
} else {
|
|
|
|
tracing::warn!("Couldn't get the author for a message");
|
|
|
|
}
|
2022-01-01 06:38:43 -06:00
|
|
|
}
|
2021-12-27 12:08:46 -06:00
|
|
|
}
|
|
|
|
}
|
2022-01-26 12:17:07 -06:00
|
|
|
tracing::info!("Message loop finished");
|
2021-12-27 06:39:01 -06:00
|
|
|
Ok(())
|
|
|
|
}
|
|
|
|
|
2022-01-01 15:35:27 -06:00
|
|
|
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)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-01-26 12:34:54 -06:00
|
|
|
#[allow(clippy::too_many_lines)]
|
2021-12-27 12:08:46 -06:00
|
|
|
async fn handle_privmsg(
|
2021-12-27 14:37:50 -06:00
|
|
|
state: &mut AppState,
|
2022-01-26 05:58:49 -06:00
|
|
|
author: &str,
|
|
|
|
origin: &str,
|
|
|
|
content: String,
|
2021-12-27 12:08:46 -06:00
|
|
|
) -> anyhow::Result<()> {
|
2022-01-26 05:58:49 -06:00
|
|
|
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)?;
|
2021-12-27 14:37:50 -06:00
|
|
|
}
|
2022-01-01 00:52:43 -06:00
|
|
|
|
2022-01-26 05:58:49 -06:00
|
|
|
if let Some(prev_msg) = state.last_msgs.get(author) {
|
|
|
|
if let Some(formatted) = sed::resolve(prev_msg, &content)? {
|
2022-01-01 01:08:28 -06:00
|
|
|
let mut result = ArrayString::<512>::new();
|
2022-01-26 05:58:49 -06:00
|
|
|
write!(result, "<{}> {}", author, formatted)?;
|
|
|
|
state.client.send_privmsg(origin, &result)?;
|
|
|
|
state.last_msgs.insert(author.into(), formatted.to_string());
|
2022-01-01 01:08:28 -06:00
|
|
|
return Ok(());
|
2022-01-01 00:52:43 -06:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-01-26 05:58:49 -06:00
|
|
|
state.last_msgs.insert(author.into(), content);
|
2021-12-27 12:08:46 -06:00
|
|
|
return Ok(());
|
|
|
|
}
|
2022-01-26 05:58:49 -06:00
|
|
|
let (command, remainder) = separate_to_space(&content, state.prefix.len());
|
2022-01-01 06:38:43 -06:00
|
|
|
tracing::debug!("Command received ({:?}; {:?})", command, remainder);
|
2021-12-27 12:08:46 -06:00
|
|
|
|
|
|
|
match command {
|
2021-12-27 14:37:50 -06:00
|
|
|
"help" => {
|
2021-12-31 17:51:54 -06:00
|
|
|
for help_line in HELP {
|
2022-01-26 05:58:49 -06:00
|
|
|
state.client.send_privmsg(origin, help_line)?;
|
2021-12-31 17:51:54 -06:00
|
|
|
}
|
2021-12-27 14:37:50 -06:00
|
|
|
}
|
2021-12-27 12:08:46 -06:00
|
|
|
"waifu" => {
|
|
|
|
let category = remainder.unwrap_or("waifu");
|
2021-12-31 17:51:54 -06:00
|
|
|
let url = misc::get_waifu_pic(category).await?;
|
2022-01-27 17:44:50 -06:00
|
|
|
let response = url
|
|
|
|
.as_deref()
|
2021-12-27 12:08:46 -06:00
|
|
|
.unwrap_or("Invalid category. Valid categories: https://waifu.pics/docs");
|
2022-01-26 05:58:49 -06:00
|
|
|
state.client.send_privmsg(origin, response)?;
|
2021-12-27 12:08:46 -06:00
|
|
|
}
|
2021-12-29 14:25:22 -06:00
|
|
|
"mock" => {
|
2022-01-26 12:34:54 -06:00
|
|
|
leek::execute(
|
2022-01-02 14:58:54 -06:00
|
|
|
state,
|
2022-01-26 12:34:54 -06:00
|
|
|
leek::Command::Mock,
|
2022-01-26 05:58:49 -06:00
|
|
|
origin,
|
|
|
|
remainder.unwrap_or(author),
|
|
|
|
)?;
|
2021-12-29 14:25:22 -06:00
|
|
|
}
|
2021-12-29 14:08:49 -06:00
|
|
|
"leet" => {
|
2022-01-26 12:34:54 -06:00
|
|
|
leek::execute(
|
2022-01-02 14:58:54 -06:00
|
|
|
state,
|
2022-01-26 12:34:54 -06:00
|
|
|
leek::Command::Leet,
|
2022-01-26 05:58:49 -06:00
|
|
|
origin,
|
|
|
|
remainder.unwrap_or(author),
|
|
|
|
)?;
|
2021-12-30 17:02:12 -06:00
|
|
|
}
|
|
|
|
"owo" => {
|
2022-01-26 12:34:54 -06:00
|
|
|
leek::execute(
|
2022-01-26 05:58:49 -06:00
|
|
|
state,
|
2022-01-26 12:34:54 -06:00
|
|
|
leek::Command::Owo,
|
2022-01-26 05:58:49 -06:00
|
|
|
origin,
|
|
|
|
remainder.unwrap_or(author),
|
|
|
|
)?;
|
2021-12-29 14:08:49 -06:00
|
|
|
}
|
2021-12-31 17:51:54 -06:00
|
|
|
"ev" => {
|
2022-01-26 05:58:49 -06:00
|
|
|
let result = misc::mathbot(author.into(), remainder, &mut state.last_eval)?;
|
|
|
|
state.client.send_privmsg(origin, &result)?;
|
2021-12-31 17:51:54 -06:00
|
|
|
}
|
2022-01-01 15:35:27 -06:00
|
|
|
"grab" => {
|
|
|
|
if let Some(target) = remainder {
|
2022-01-26 05:58:49 -06:00
|
|
|
if target == author {
|
2022-01-02 14:58:54 -06:00
|
|
|
state
|
|
|
|
.client
|
2022-01-30 11:26:09 -06:00
|
|
|
.send_privmsg(origin, "You can't grab yourself")?;
|
2022-01-02 14:58:54 -06:00
|
|
|
return Ok(());
|
2022-01-01 15:35:27 -06:00
|
|
|
}
|
|
|
|
if let Some(prev_msg) = state.last_msgs.get(target) {
|
2022-01-27 17:44:50 -06:00
|
|
|
if state
|
|
|
|
.db
|
|
|
|
.add_quote(Quote {
|
|
|
|
quote: prev_msg.clone(),
|
|
|
|
author: target.into(),
|
|
|
|
})
|
|
|
|
.await
|
|
|
|
{
|
2022-01-30 11:26:09 -06:00
|
|
|
state.client.send_privmsg(origin, "Quote added")?;
|
2022-01-01 15:35:27 -06:00
|
|
|
} else {
|
2022-01-02 14:58:54 -06:00
|
|
|
state
|
|
|
|
.client
|
2022-01-30 11:26:09 -06:00
|
|
|
.send_privmsg(origin, "A database error has occurred")?;
|
2022-01-01 15:35:27 -06:00
|
|
|
}
|
|
|
|
} else {
|
2022-01-02 14:58:54 -06:00
|
|
|
state
|
|
|
|
.client
|
2022-01-30 11:26:09 -06:00
|
|
|
.send_privmsg(origin, "No previous messages to grab")?;
|
2022-01-01 15:35:27 -06:00
|
|
|
}
|
|
|
|
} else {
|
2022-01-26 05:58:49 -06:00
|
|
|
state.client.send_privmsg(origin, "No nickname to grab")?;
|
2022-01-01 15:35:27 -06:00
|
|
|
}
|
|
|
|
}
|
|
|
|
"quot" => {
|
2022-01-26 12:34:54 -06:00
|
|
|
if let Some(quote) = state.db.get_quote(remainder.map(ToString::to_string)).await {
|
2022-01-01 15:35:27 -06:00
|
|
|
let mut resp = ArrayString::<512>::new();
|
2022-01-26 17:58:00 -06:00
|
|
|
write!(resp, "\"{}\" ~{}", quote.quote, quote.author)?;
|
2022-01-26 05:58:49 -06:00
|
|
|
state.client.send_privmsg(origin, &resp)?;
|
2022-01-01 15:35:27 -06:00
|
|
|
} else {
|
2022-01-26 05:58:49 -06:00
|
|
|
state.client.send_privmsg(origin, "No quotes found")?;
|
2022-01-01 15:35:27 -06:00
|
|
|
}
|
|
|
|
}
|
2021-12-27 12:08:46 -06:00
|
|
|
_ => {
|
2022-01-26 05:58:49 -06:00
|
|
|
state.client.send_privmsg(origin, "Unknown command")?;
|
2021-12-27 12:08:46 -06:00
|
|
|
}
|
2021-12-27 06:39:01 -06:00
|
|
|
}
|
|
|
|
Ok(())
|
|
|
|
}
|