From aae1f386456a605c12263e2ba5f09dc36bb180b3 Mon Sep 17 00:00:00 2001 From: Keuin Date: Sun, 27 Mar 2022 17:40:18 +0800 Subject: Finally, it works. Further cleanup work is needed. --- src/bot.rs | 63 ++++++++++++++++++++++++++++++++++++++++++++++++++ src/config.rs | 24 +++++++++++++++++++ src/database.rs | 72 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 55 +++++++++++++++++++++++++++++++++++++++++++ src/token.rs | 6 +++++ src/user.rs | 6 +++++ src/web.rs | 59 ++++++++++++++++++++++++++++++++++++++++++++++ 7 files changed, 285 insertions(+) create mode 100644 src/bot.rs create mode 100644 src/config.rs create mode 100644 src/database.rs create mode 100644 src/main.rs create mode 100644 src/token.rs create mode 100644 src/user.rs create mode 100644 src/web.rs (limited to 'src') diff --git a/src/bot.rs b/src/bot.rs new file mode 100644 index 0000000..b1a4ea2 --- /dev/null +++ b/src/bot.rs @@ -0,0 +1,63 @@ +use std::error::Error; +use std::ops::Deref; +use std::sync::Arc; + +use teloxide::{prelude2::*, types::MessageKind, utils::command::BotCommand}; + +use crate::{database, DbPool, token}; +use crate::user::User; + +#[derive(BotCommand, Clone)] +#[command(rename = "lowercase", description = "These commands are supported:")] +enum Command { + #[command(description = "display this text.")] + Help, + #[command(description = "get your personal token.")] + Start, +} + +async fn answer(bot: AutoSend, message: Message, command: Command, db: Arc) + -> Result<(), Box> { + match command { + Command::Help => { + bot.send_message(message.chat.id, Command::descriptions()).await?; + } + Command::Start => { + if let MessageKind::Common(msg) = message.kind { + if msg.from.is_none() { + return Ok(()); // ignore messages from channel + } + let from = msg.from.unwrap(); + let db = db.deref(); + let chat_id = message.chat.id; + match database::create_user(db, User { + id: from.id as u64, + name: from.username.unwrap_or_else(|| String::from("")), + token: token::generate(), + chat_id: chat_id as u64, + }).await { + Ok(_) => {} + Err(why) => println!("cannot create user: {:?}", why), + } + bot.send_message( + chat_id, + match database::get_user_by_chat_id(db, chat_id as u64).await? { + Some(user) => + format!("Your token is `{}`. Treat it as a secret!", user.token), + _ => + String::from("Error: cannot fetch token.") + }, + ).await?; + } + } + }; + Ok(()) +} + +pub async fn repl(bot: Bot, db: Arc) { + teloxide::repls2::commands_repl( + bot.auto_send(), + move |bot, msg, cmd| + answer(bot, msg, cmd, Arc::clone(&db)), Command::ty(), + ).await; +} diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..d32c525 --- /dev/null +++ b/src/config.rs @@ -0,0 +1,24 @@ +use std::fs::File; +use std::io::Read; + +use serde_derive::Deserialize; + +#[derive(Deserialize)] +pub struct Config { + pub bot_token: String, + pub log_file: String, + pub db_file: String, +} + +impl Config { + // Read config file. Panic if any error occurs. + pub fn from_file(file_path: &str) -> Config { + let mut file = File::open(file_path) + .unwrap_or_else(|_| panic!("cannot open file {}", file_path)); + let mut config = String::new(); + file.read_to_string(&mut config) + .unwrap_or_else(|_| panic!("cannot read config file {}", file_path)); + return serde_json::from_str(config.as_str()) + .expect("cannot decode config file in JSON"); + } +} \ No newline at end of file diff --git a/src/database.rs b/src/database.rs new file mode 100644 index 0000000..edb2fc7 --- /dev/null +++ b/src/database.rs @@ -0,0 +1,72 @@ +use std::str::FromStr; + +use futures::TryStreamExt; +use log::{debug, error, info, warn}; +use sqlx::Row; +use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions}; + +use crate::user::User; + +pub type DbPool = sqlx::sqlite::SqlitePool; + +const SQLITE_THREAD_POOL_SIZE: u32 = 16; + +pub async fn open(file_path: &str) -> Result { + let opt = + SqliteConnectOptions::from_str(format!("sqlite://{}", file_path).as_str())? + .create_if_missing(true); + let pool = SqlitePoolOptions::new() + .max_connections(SQLITE_THREAD_POOL_SIZE) + .connect_with(opt).await?; + + // create table + sqlx::query(r#"CREATE TABLE IF NOT EXISTS "user" ( + "id" INTEGER NOT NULL CHECK(id >= 0) UNIQUE, + "name" TEXT, + "token" TEXT NOT NULL UNIQUE, + "chat_id" INTEGER NOT NULL CHECK(chat_id >= 0) UNIQUE, + PRIMARY KEY("id","token","chat_id") + )"#).execute(&pool).await?; + + Ok(pool) +} + +pub async fn get_user(db: &DbPool, sql: &str, param: &str) -> Result, sqlx::Error> { + let mut rows = + sqlx::query(sql) + .bind(param).fetch(db); + let result = rows.try_next().await?; + match result { + None => Ok(None), + Some(row) => { + let id: i64 = row.try_get("id")?; + let chat_id: i64 = row.try_get("chat_id")?; + Ok(Some(User { + id: id as u64, + name: row.try_get("name")?, + token: row.try_get("token")?, + chat_id: chat_id as u64, + })) + } + } +} + +pub async fn get_user_by_token(db: &DbPool, token: &str) -> Result, sqlx::Error> { + get_user(db, "SELECT id, name, token, chat_id FROM user WHERE token=$1 LIMIT 1", + token).await +} + +pub async fn get_user_by_chat_id(db: &DbPool, chat_id: u64) -> Result, sqlx::Error> { + get_user(db, "SELECT id, name, token, chat_id FROM user WHERE chat_id=$1 LIMIT 1", + chat_id.to_string().as_str()).await +} + +pub async fn create_user(db: &DbPool, user: User) -> sqlx::Result<()> { + sqlx::query("INSERT INTO user (id, name, token, chat_id) VALUES ($1,$2,$3,$4)") + .bind(user.id.to_string()) + .bind(user.name) + .bind(user.token) + .bind(user.chat_id.to_string()) + .execute(db).await?; + Ok(()) +} \ No newline at end of file diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..fffdb65 --- /dev/null +++ b/src/main.rs @@ -0,0 +1,55 @@ +use std::convert::Infallible; +use std::sync::Arc; + +use log::{debug, error, info, warn}; +use teloxide::prelude2::*; +use warp::Filter; + +use config::Config; + +use crate::database::DbPool; + +mod database; +mod user; +mod web; +mod bot; +mod token; +mod config; + +const CONFIG_FILE_NAME: &str = "kimikuri.json"; +const MAX_BODY_LENGTH: u64 = 1024 * 16; + +fn with_db(db_pool: DbPool) -> impl Filter + Clone { + warp::any().map(move || db_pool.clone()) +} + +// TODO replace with generic +fn with_bot(bot: Bot) -> impl Filter + Clone { + warp::any().map(move || bot.clone()) +} + +#[tokio::main] +async fn main() { + pretty_env_logger::init(); + debug!("Loading bot config."); + let config = Config::from_file(CONFIG_FILE_NAME); + info!("Starting bot."); + let bot = Bot::new(config.bot_token); + + let db = config.db_file.as_str(); + println!("Database: {}", db); + let db: DbPool = database::open(db) + .await.expect(&*format!("cannot open database {}", db)); + + let send_message = warp::path("message") + .and(warp::post()) + .and(warp::body::content_length_limit(MAX_BODY_LENGTH)) + .and(warp::body::json()) + .and(with_db(db.clone())) + .and(with_bot(bot.clone())) + .and_then(web::handler); + + tokio::spawn(bot::repl(bot, Arc::new(db))); + + warp::serve(send_message).run(([127, 0, 0, 1], 3030)).await; +} diff --git a/src/token.rs b/src/token.rs new file mode 100644 index 0000000..3de71ce --- /dev/null +++ b/src/token.rs @@ -0,0 +1,6 @@ +pub fn generate() -> String { + let mut rng = urandom::csprng(); + let mut bytes = [0u8; 32]; + rng.fill_bytes(&mut bytes); + bs58::encode(bytes).into_string() +} \ No newline at end of file diff --git a/src/user.rs b/src/user.rs new file mode 100644 index 0000000..aef43fa --- /dev/null +++ b/src/user.rs @@ -0,0 +1,6 @@ +pub struct User { + pub id: u64, + pub name: String, + pub token: String, + pub chat_id: u64, +} \ No newline at end of file diff --git a/src/web.rs b/src/web.rs new file mode 100644 index 0000000..8b6505b --- /dev/null +++ b/src/web.rs @@ -0,0 +1,59 @@ +use serde_derive::{Deserialize, Serialize}; +use teloxide::{prelude2::*}; +use warp::{Rejection, Reply}; + +use crate::{Bot, database, DbPool}; + +#[derive(Deserialize, Serialize)] +pub struct SendMessage { + token: String, + message: String, +} + +#[derive(Deserialize, Serialize)] +pub struct SendMessageResponse { + success: bool, + message: String, +} + +pub async fn handler(req: SendMessage, db: DbPool, bot: Bot) -> std::result::Result { + println!("Token: {}, Message: {}", req.token, req.message); + let user = database::get_user_by_token(&db, req.token.as_str()).await; + Ok(warp::reply::json(&match user { + Ok(u) => match u { + Some(user) => { + log::info!("User: {} (id={}), message: {}", + user.name, user.id, req.message); + // TODO send message to Telegram + let bot = bot.auto_send(); + match bot.send_message(user.chat_id as i64, req.message).await { + Ok(_) => SendMessageResponse { + success: true, + message: String::new(), + }, + Err(why) => { + println!("Failed to send message to telegram: {:?}", why); + SendMessageResponse { + success: false, + message: String::from("Failed to send message to telegram."), + } + } + } + } + None => { + log::warn!("Invalid token {}, message: {}", req.token, req.message); + SendMessageResponse { + success: false, + message: String::from("Invalid token."), + } + } + }, + Err(_) => { + log::error!("Error when querying the database."); + SendMessageResponse { + success: false, + message: String::from("Invalid parameter."), + } + } + })) +} -- cgit v1.2.3