From b2bb016573aaeb6cc9481a5a54b8715435bbf83c Mon Sep 17 00:00:00 2001 From: AB Date: Sat, 5 Dec 2020 15:57:11 +0300 Subject: [PATCH] Many improvements --- .gitignore | 7 ++- Cargo.toml | 8 +++- assets/scheme.sql | 6 +++ src/db.rs | 64 ++++++++++++++++++++------- src/errors.rs | 20 ++++++++- src/main.rs | 110 ++++++++++++++++++++++++++++++++++++++++++---- src/utils.rs | 80 +++++++++++++++++++++++++++++++-- 7 files changed, 264 insertions(+), 31 deletions(-) diff --git a/.gitignore b/.gitignore index 6f247c8..5688ef8 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,8 @@ /target memory.sqlite3 - +/doc +/photo +/sticker +/video +/voice +/.idea diff --git a/Cargo.toml b/Cargo.toml index bb51a5c..777aa87 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,5 +21,11 @@ hyper-tls = { version = "0.4", optional = true } futures = "0.3" hyper-rustls = { version = "0.19", optional = true } -rusqlite = "0.24.1" +rusqlite = { version = "0.24.1", features = ["bundled"]} html-escape = "0.2" +reqwest = "0.10.9" +uuid = { version = "0.8", features = ["v4"] } +sha1 = "*" +env_logger = "0.7" +log = { version = "^0.4.5", features = ["std"] } + diff --git a/assets/scheme.sql b/assets/scheme.sql index c7d446c..f15c3cc 100644 --- a/assets/scheme.sql +++ b/assets/scheme.sql @@ -17,6 +17,12 @@ CREATE TABLE `conf` ( `date` INTEGER NOT NULL, PRIMARY KEY(`id`) ); +CREATE TABLE `file` ( + `path` TEXT NOT NULL UNIQUE, + `user_id` TEXT NOT NULL, + `conf_id` TEXT NOT NULL, + PRIMARY KEY(`path`) +); CREATE TABLE IF NOT EXISTS "relations" ( `id` INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT, `word_id` INTEGER NOT NULL, diff --git a/src/db.rs b/src/db.rs index 8ef2c19..f1b353c 100644 --- a/src/db.rs +++ b/src/db.rs @@ -1,6 +1,7 @@ use crate::errors; use crate::utils; use rusqlite::{named_params, params, Connection, Error, Result}; +use sha1::{Digest, Sha1}; use std::time::SystemTime; use telegram_bot::*; @@ -63,7 +64,8 @@ pub(crate) fn get_conf(id: telegram_bot::ChatId) -> Result Err(errors::Error::ConfNotFound) } else { Ok(confs[0].clone()) - }} + } +} pub(crate) fn get_confs() -> Result> { let conn = open()?; @@ -113,7 +115,7 @@ pub(crate) fn get_members(id: telegram_bot::ChatId) -> Result Result<(), Error> { +pub(crate) async fn add_conf(message: Message) -> Result<(), Error> { let conn = open()?; let title = utils::get_title(&message); @@ -122,7 +124,7 @@ pub(crate) async fn add_conf(api: Api, message: Message) -> Result<(), Error> { let update = Conf { id: message.chat.id(), title, - date: 0 + date: 0, }; let mut stmt = conn.prepare( "UPDATE conf @@ -131,17 +133,14 @@ pub(crate) async fn add_conf(api: Api, message: Message) -> Result<(), Error> { WHERE id = :id", )?; - stmt.execute_named(&[ - (":id", &update.id.to_string()), - (":title", &update.title), - ])?; + stmt.execute_named(&[(":id", &update.id.to_string()), (":title", &update.title)])?; //println!("Conf {:?} updated: {:?}", update.title, get_conf(update.id)); } Err(e) => { let update = Conf { id: message.chat.id(), title, - date: 0 + date: 0, }; let unix_time = SystemTime::now() .duration_since(SystemTime::UNIX_EPOCH) @@ -168,8 +167,7 @@ pub(crate) async fn add_conf(api: Api, message: Message) -> Result<(), Error> { Ok(()) } - -pub(crate) async fn add_user(api: Api, message: Message) -> Result<(), Error> { +pub(crate) async fn add_user(message: Message) -> Result<(), Error> { let conn = open()?; match get_user(message.from.id) { Ok(user) => { @@ -179,7 +177,7 @@ pub(crate) async fn add_user(api: Api, message: Message) -> Result<(), Error> { last_name: message.from.last_name, username: message.from.username, is_bot: false, - language_code: None + language_code: None, }; let mut stmt = conn.prepare( "UPDATE user @@ -209,7 +207,7 @@ pub(crate) async fn add_user(api: Api, message: Message) -> Result<(), Error> { last_name: message.from.last_name, username: message.from.username, is_bot: false, - language_code: None + language_code: None, }; let mut stmt = conn.prepare( "INSERT OR IGNORE INTO @@ -224,10 +222,44 @@ pub(crate) async fn add_user(api: Api, message: Message) -> Result<(), Error> { (":date", &unix_time), ])?; //println!("User added: {:?}", user); - - }, - _ => {}, + } + _ => {} } Ok(()) - +} + +pub(crate) async fn add_file( + message: &Message, + path: String, + file_id: String, +) -> Result<(), Error> { + let conn = open()?; + let mut stmt = conn.prepare( + "INSERT OR IGNORE INTO + file('path', 'user_id', 'conf_id', 'file_id') + VALUES (:path, :user_id, :conf_id, :file_id)", + )?; + stmt.execute_named(&[ + (":path", &path), + (":user_id", &message.from.id.to_string()), + (":conf_id", &message.chat.id().to_string()), + (":file_id", &file_id), + ])?; + Ok(()) +} + +pub(crate) async fn get_file(file_id: String) -> Result<(), errors::Error> { + let conn = open()?; + let mut stmt = conn.prepare("SELECT path FROM file WHERE file_id = :file_id")?; + let mut rows = stmt.query_named(&[(":file_id", &file_id)])?; + let mut files = Vec::new(); + + while let Some(row) = rows.next()? { + files.push("should be rewritten"); + } + if files.len() > 0 { + Ok(()) + } else { + Err(errors::Error::FileNotFound) + } } diff --git a/src/errors.rs b/src/errors.rs index 181848b..2651c2e 100644 --- a/src/errors.rs +++ b/src/errors.rs @@ -1,14 +1,18 @@ +use reqwest::Error as reqwest_error; use rusqlite::Error as sqlite_error; use rusqlite::{named_params, params, Connection, Result}; +use std::{fmt, io, io::Error as io_error}; use telegram_bot::Error as tg_error; -use std::{fmt, io}; #[derive(Debug)] pub(crate) enum Error { UserNotFound, SQLITE3Error(sqlite_error), TelegramError(tg_error), + ReqwestError(reqwest_error), ConfNotFound, + IOError(io_error), + FileNotFound, } impl fmt::Display for Error { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { @@ -26,4 +30,16 @@ impl From for Error { fn from(e: tg_error) -> Error { return Error::TelegramError(e); } -} \ No newline at end of file +} + +impl From for Error { + fn from(e: reqwest_error) -> Error { + return Error::ReqwestError(e); + } +} + +impl From for Error { + fn from(e: io_error) -> Error { + return Error::IOError(e); + } +} diff --git a/src/main.rs b/src/main.rs index ee70aef..2234a5c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,15 +1,19 @@ -use std::env; +use std::{env,process}; use futures::StreamExt; +use reqwest; use telegram_bot::types::chat::MessageChat; use telegram_bot::*; +#[macro_use] +extern crate log; +use env_logger::Env; mod commands; mod db; mod errors; mod utils; -async fn handler(api: Api, message: Message) -> Result<(), Error> { +async fn handler(api: Api, message: Message, token: String) -> Result<(), errors::Error> { match message.kind { MessageKind::Text { ref data, .. } => { let title = utils::get_title(&message); @@ -27,16 +31,106 @@ async fn handler(api: Api, message: Message) -> Result<(), Error> { _ => (), } } + MessageKind::Photo { + ref caption, + ref data, + .. + } => { + let title = utils::get_title(&message); + println!( + "<{}({})>[{}({})]: *PHOTO* {}", + &message.chat.id(), + title, + &message.from.id, + &message.from.first_name, + caption.clone().unwrap_or("NO_TITLE".to_string()) + ); + utils::get_files(api, message, token).await?; + } + + MessageKind::Document { ref caption, .. } => { + let title = utils::get_title(&message); + println!( + "<{}({})>[{}({})]: *DOCUMENT* {}", + &message.chat.id(), + title, + &message.from.id, + &message.from.first_name, + caption.clone().unwrap_or("NO_TITLE".to_string()) + ); + /*match utils::get_files(api, message, token).await { + Ok(count) => println!("Got {} files successfully.", count), + Err(e) => println!("Couldn't get files: {:?}", e) + } + + */ + utils::get_files(api, message, token).await?; + + } + + MessageKind::Sticker { ref data, .. } => { + let title = utils::get_title(&message); + println!( + "<{}({})>[{}({})]: *STICKER*", + &message.chat.id(), + title, + &message.from.id, + &message.from.first_name, + ); + utils::get_files(api, message, token).await?; + } + + MessageKind::Voice { .. } => { + let title = utils::get_title(&message); + println!( + "<{}({})>[{}({})]: *VOICE*", + &message.chat.id(), + title, + &message.from.id, + &message.from.first_name, + ); + utils::get_files(api, message, token).await?; + } + + MessageKind::Video { .. } => { + let title = utils::get_title(&message); + println!( + "<{}({})>[{}({})]: *VIDEO*", + &message.chat.id(), + title, + &message.from.id, + &message.from.first_name, + ); + utils::get_files(api, message, token).await?; + } + + MessageKind::VideoNote { .. } => { + let title = utils::get_title(&message); + println!( + "<{}({})>[{}({})]: *VIDEO_NOTE*", + &message.chat.id(), + title, + &message.from.id, + &message.from.first_name, + ); + utils::get_files(api, message, token).await?; + } _ => (), }; - Ok(()) } #[tokio::main] async fn main() -> Result<(), errors::Error> { - let token = env::var("TELEGRAM_BOT_TOKEN").expect("TELEGRAM_BOT_TOKEN not set"); - let api = Api::new(token); + env_logger::from_env(Env::default().default_filter_or("debug")).init(); + let token = match env::var("TELEGRAM_BOT_TOKEN") { + Ok(token) => token, + Err(_) => { + error!("TELEGRAM_BOT_TOKEN not set"); + process::exit(0x0001); + }, + }; + let api = Api::new(token.clone()); // Fetch new updates via long poll method let mut stream = api.stream(); @@ -44,10 +138,10 @@ async fn main() -> Result<(), errors::Error> { // If the received update contains a new message... let update = update?; if let UpdateKind::Message(message) = update.kind { - db::add_user(api.clone(), message.clone()).await?; - db::add_conf(api.clone(), message.clone()).await?; + db::add_user(message.clone()).await?; + db::add_conf(message.clone()).await?; - handler(api.clone(), message).await?; + handler(api.clone(), message, token.clone()).await?; } } Ok(()) diff --git a/src/utils.rs b/src/utils.rs index d5af2b0..fab6ce6 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,10 +1,84 @@ +use reqwest::Client; +use sha1::Sha1; +use std::fs::File; +use std::io::prelude::*; +use std::path::Path; +use std::{env, io}; use telegram_bot::*; +use uuid::Uuid; -pub(crate) fn get_title(message: &telegram_bot::Message) -> String { +use crate::db; +use crate::errors; +extern crate reqwest; + +pub(crate) fn get_title(message: &Message) -> String { match &message.chat { MessageChat::Supergroup(chat) => chat.title.clone(), MessageChat::Group(chat) => chat.title.clone(), MessageChat::Private(chat) => "PRIVATE".to_string(), - _ => "PRIVATE".to_string() + _ => "PRIVATE".to_string(), } -} \ No newline at end of file +} + +pub(crate) async fn get_files( + api: Api, + message: Message, + token: String, +) -> Result { + let mut file_count = 0; + let file_type = match message.kind { + MessageKind::Photo { .. } => "photo".to_string(), + MessageKind::Document { .. } => "doc".to_string(), + MessageKind::Voice { .. } => "voice".to_string(), + MessageKind::Video { .. } => "video".to_string(), + MessageKind::VideoNote { .. } => "video".to_string(), + MessageKind::Sticker { .. } => "sticker".to_string(), + _ => "docs".to_string(), + }; + if let Some(files) = message.get_files() { + let group_title = get_title(&message); + let author = message.from.id; + for file in files { + file_count += 1; + let uuid = Uuid::new_v4(); + match api.send(&file).await { + Ok(api_response) => { + let url = format!( + "https://api.telegram.org/file/bot{}/{}", + token, + api_response.file_path.unwrap() + ); + let mut file_response = reqwest::get(&url).await?; + let ext = { + file_response + .url() + .path_segments() + .and_then(|segments| segments.last()) + .and_then(|name| if name.is_empty() { None } else { Some(name) }) + .unwrap_or("tmp.bin") + .split('.') + .last() + .unwrap() + }; + let path = format!("{}/{}_{}_{}.{}", file_type, group_title, author, uuid, ext); + let mut hasher = Sha1::new(); + let content = file_response.bytes().await?; + hasher.update(&content); + let file_hash = hasher.digest().to_string(); + match db::get_file(file_hash.clone()).await { + Ok(_) => { + println!("File exist"); + } + Err(_) => { + let mut dest = File::create(path.clone())?; + dest.write(&content); + } + }; + db::add_file(&message, path, file_hash).await?; + } + Err(e) => println!("Couldn't get file: {}", e) + } + } + }; + Ok(file_count) +}