Many improvements

.gitignore

@@ -1,3 +1,8 @@
 /target
 memory.sqlite3
-
+/doc
+/photo
+/sticker
+/video
+/voice
+/.idea

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"] }
+

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,

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<Conf, errors::Error>
         Err(errors::Error::ConfNotFound)
     } else {
         Ok(confs[0].clone())
-    }}
+    }
+}
 
 pub(crate) fn get_confs() -> Result<Vec<Conf>> {
     let conn = open()?;
@@ -113,7 +115,7 @@ pub(crate) fn get_members(id: telegram_bot::ChatId) -> Result<Vec<telegram_bot::
     Ok(users)
 }
 
-pub(crate) async fn add_conf(api: Api, message: Message) -> 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(&[
+            stmt.execute_named(&[(":id", &update.id.to_string()), (":title", &update.title)])?;
-                (":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(message: Message) -> Result<(), Error> {
-pub(crate) async fn add_user(api: Api, 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)
+    }
 }

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<tg_error> for Error {
     fn from(e: tg_error) -> Error {
         return Error::TelegramError(e);
     }
-}
+}
+
+impl From<reqwest_error> for Error {
+    fn from(e: reqwest_error) -> Error {
+        return Error::ReqwestError(e);
+    }
+}
+
+impl From<io_error> for Error {
+    fn from(e: io_error) -> Error {
+        return Error::IOError(e);
+    }
+}

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");
+    env_logger::from_env(Env::default().default_filter_or("debug")).init();
-    let api = Api::new(token);
+    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_user(message.clone()).await?;
-            db::add_conf(api.clone(), message.clone()).await?;
+            db::add_conf(message.clone()).await?;
 
-            handler(api.clone(), message).await?;
+            handler(api.clone(), message, token.clone()).await?;
         }
     }
     Ok(())

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(),
     }
-}
+}
+
+pub(crate) async fn get_files(
+    api: Api,
+    message: Message,
+    token: String,
+) -> Result<i32, errors::Error> {
+    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)
+}