Add /sql command.

src/commands.rs

@@ -1,16 +1,22 @@
 use crate::db;
 use crate::errors::Error;
+use crate::errors::Error::SQLInvalidCommand;
 use async_trait::async_trait;
 use html_escape::encode_text;
 use markov::Chain;
 use mystem::Case::Nominative;
 use mystem::Gender::Feminine;
-use mystem::Tense::{Inpresent, Past};
-use mystem::Person::First;
 use mystem::MyStem;
+use mystem::Person::First;
+use mystem::Tense::{Inpresent, Past};
 use rand::seq::SliceRandom;
 use rand::Rng;
 use regex::Regex;
+use rusqlite::types::FromSql;
+use rusqlite::{CachedStatement, Rows, ToSql};
+use sqlparser::ast::Statement;
+use sqlparser::dialect::GenericDialect;
+use sqlparser::parser::Parser;
 use telegram_bot::prelude::*;
 use telegram_bot::{Api, Message, ParseMode};
 
@@ -29,14 +35,9 @@ pub struct Markov {
 pub struct Omedeto {
     pub data: String,
 }
-
-// pub enum Command {
-//     Here(Here),
-//     Top { data: String },
-//     MarkovAll { data: String },
-//     Markov { data: String },
-//     Omedeto { data: String },
-// }
+pub struct Sql {
+    pub data: String,
+}
 
 #[async_trait]
 pub trait Execute {
@@ -49,6 +50,189 @@ pub trait Execute {
     ) -> Result<(), Error>;
 }
 
+#[async_trait]
+impl Execute for Sql {
+    async fn run(&self, api: Api, message: Message) -> Result<(), Error> {
+        let mut sql = self.data.to_uppercase();
+        let is_head = if sql.starts_with('-') {
+            sql = sql.replacen("-", "", 1);
+            false
+        } else {
+            true
+        };
+        let dialect = GenericDialect {};
+        let ast: Result<Vec<Statement>, Error> = match Parser::parse_sql(&dialect, &sql) {
+            Ok(ast) => Ok(ast),
+            Err(_) => {
+                warn!("Invalid SQL - {}", sql);
+                Err(SQLInvalidCommand)
+            }
+        };
+        let ast = match ast {
+            Err(e) => {
+                let _ = api
+                    .send(
+                        message
+                            .text_reply(format!("❌ Invalid SQL. Syntax error ❌"))
+                            .parse_mode(ParseMode::Html),
+                    )
+                    .await;
+                return Err(SQLInvalidCommand);
+            }
+            Ok(ast) => ast,
+        };
+        let msg: Result<String, Error> = match ast.len() {
+            l if l > 1 => {
+                //Max 1 request per message allowed only.
+                Err(Error::SQLBannedCommand)
+            }
+            _ => match ast[0] {
+                sqlparser::ast::Statement::Query { .. } => {
+                    let conn = db::open()?;
+                    let mut x = match conn.prepare_cached(&sql) {
+                        Ok(mut stmt) => {
+                            let query = match stmt.query(rusqlite::NO_PARAMS) {
+                                Err(e) => Err(SQLInvalidCommand),
+                                Ok(mut rows) => {
+                                    let mut res: Vec<Vec<String>> = match rows.column_names() {
+                                        Some(n) => vec![n
+                                            .into_iter()
+                                            .map(|s| {
+                                                let t = String::from(s);
+                                                if t.len() > 10 {
+                                                    "EMSGSIZE".to_string()
+                                                } else {
+                                                    t
+                                                }
+                                            })
+                                            .collect()],
+                                        None => return Err(SQLInvalidCommand),
+                                    };
+                                    let index_count = match rows.column_count() {
+                                        Some(c) => c,
+                                        None => return Err(SQLInvalidCommand),
+                                    };
+                                    while let Some(row) = rows.next().unwrap() {
+                                        let mut tmp: Vec<String> = Vec::new();
+                                        for i in 0..index_count {
+                                            match row.get(i).unwrap_or(None) {
+                                                Some(rusqlite::types::Value::Text(t)) => {
+                                                    tmp.push(t)
+                                                }
+                                                Some(rusqlite::types::Value::Integer(t)) => {
+                                                    tmp.push(t.to_string())
+                                                }
+                                                Some(rusqlite::types::Value::Blob(t)) => {
+                                                    tmp.push("Binary".to_string())
+                                                }
+                                                Some(rusqlite::types::Value::Real(t)) => {
+                                                    tmp.push(t.to_string())
+                                                }
+                                                Some(rusqlite::types::Value::Null) => {
+                                                    tmp.push("Null".to_string())
+                                                }
+                                                None => tmp.push("Null".to_string()),
+                                            };
+                                        }
+                                        res.push(tmp);
+                                    }
+
+                                    // add Header
+                                    let mut msg = if is_head {
+                                        let mut x = String::from("<b>");
+                                        for head in res[0].iter() {
+                                            x = format!("{} {}", x, head);
+                                        }
+                                        format!("{}{}", x, "</b>\n")
+                                    } else {
+                                        String::new()
+                                    };
+
+                                    // remove header
+                                    res.remove(0);
+
+                                    msg = format!("{}{}", msg, "<pre>");
+                                    for line in res.iter() {
+                                        for field in line.iter() {
+                                            msg = format!("{}{}", msg, format!("{} ", field));
+                                        }
+                                        msg = format!("{}{}", msg, "\n");
+                                    }
+                                    msg = format!("{}{}", msg, "</pre>");
+                                    msg = if msg.len() > 4096 {
+                                        "🚫 Result is too big. Use LIMIT 🚫".into()
+                                    } else {
+                                        msg
+                                    };
+                                    Ok(msg)
+                                }
+                            };
+                            query
+                        }
+                        Err(e) => Err(Error::SQLITE3Error(e)),
+                    };
+                    x
+                }
+                _ => {
+                    warn!("SELECT requests allowed only.");
+                    Err(Error::SQLBannedCommand)
+                }
+            },
+        };
+        match msg {
+            Ok(msg) => {
+                match api
+                    .send(message.text_reply(msg).parse_mode(ParseMode::Html))
+                    .await
+                {
+                    Ok(_) => debug!("/sql command sent to {}", message.chat.id()),
+                    Err(_) => warn!("/sql command sent failed to {}", message.chat.id()),
+                }
+            }
+            Err(e) => match e {
+                Error::SQLITE3Error(e) => {
+                    let _ = api
+                        .send(
+                            message
+                                .text_reply(format!("❌ An error ocurred {}❌", e))
+                                .parse_mode(ParseMode::Html),
+                        )
+                        .await;
+                }
+                Error::SQLBannedCommand => {
+                    let _ = api
+                        .send(
+                            message
+                                .text_reply(format!("🚫 SELECT requests allowed only 🚫"))
+                                .parse_mode(ParseMode::Html),
+                        )
+                        .await;
+                }
+                Error::SQLInvalidCommand => {
+                    let _ = api
+                        .send(
+                            message
+                                .text_reply(format!("🚫 Invalid SQL. Check DB scheme. 🚫"))
+                                .parse_mode(ParseMode::Html),
+                        )
+                        .await;
+                }
+                _ => {}
+            },
+        }
+        Ok(())
+    }
+
+    async fn run_mystem(
+        &self,
+        api: Api,
+        message: Message,
+        mystem: &mut MyStem,
+    ) -> Result<(), Error> {
+        unimplemented!()
+    }
+}
+
 #[async_trait]
 impl Execute for Here {
     async fn run(&self, api: Api, message: Message) -> Result<(), Error> {
@@ -76,8 +260,6 @@ impl Execute for Here {
             Ok(_) => debug!("/here command sent to {}", message.chat.id()),
             Err(_) => warn!("/here command sent failed to {}", message.chat.id()),
         }
-        //api.send(message.chat.text("Text to message chat")).await?;
-        //api.send(message.from.text("Private text")).await?;
         Ok(())
     }
 
@@ -113,8 +295,6 @@ impl Execute for Top {
             Ok(_) => debug!("/top command sent to {}", message.chat.id()),
             Err(_) => warn!("/top command sent failed to {}", message.chat.id()),
         }
-        //api.send(message.chat.text("Text to message chat")).await?;
-        //api.send(message.from.text("Private text")).await?;
         Ok(())
     }
 
@@ -387,7 +567,6 @@ impl Execute for Omedeto {
             verbs_i.pop().unwrap_or(placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
             verbs_i.pop().unwrap_or(placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
         );
-        //debug!("{:?}", result);
         match api
             .send(
                 message