Hexor/desubot
1
0
Fork 0
mirror of https://github.com/house-of-vanity/desubot.git synced 2025-07-08 21:04:07 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request #6 from house-of-vanity/sql

Browse Source 
Simplify SQL command. Add limit.
This commit is contained in:
House of Vanity
2021-01-05 03:46:03 +03:00
committed by GitHub
parent ea6d9b55a1 7a66034381
commit d30cf10c29
3 changed files with 135 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files

284
src/commands.rs
View File

@ -1,6 +1,6 @@
use crate::db; use crate::db;
use crate::errors::Error; use crate::errors::Error;
use crate::errors::Error::SQLInvalidCommand; use crate::errors::Error::{SQLInvalidCommand, SQLITE3Error};
use async_trait::async_trait; use async_trait::async_trait;
use html_escape::encode_text; use html_escape::encode_text;
use markov::Chain; use markov::Chain;
@ -39,18 +39,18 @@ pub struct Sql {
#[async_trait] #[async_trait]
pub trait Execute { pub trait Execute {
async fn run(&self, api: Api, message: Message) -> Result<(), Error>; async fn run(&self, api: &Api, message: &Message) -> Result<(), Error>;
async fn run_mystem( async fn run_mystem(
&self, &self,
api: Api, api: &Api,
message: Message, message: &Message,
mystem: &mut MyStem, mystem: &mut MyStem,
) -> Result<(), Error>; ) -> Result<(), Error>;
} }
#[async_trait] #[async_trait]
impl Execute for Sql { impl Execute for Sql {
async fn run(&self, api: Api, message: Message) -> Result<(), Error> { async fn run(&self, api: &Api, message: &Message) -> Result<(), Error> {
let mut sql = self.data.to_uppercase(); let mut sql = self.data.to_uppercase();
let is_head = if sql.starts_with('-') { let is_head = if sql.starts_with('-') {
sql = sql.replacen("-", "", 1); sql = sql.replacen("-", "", 1);
@ -59,173 +59,111 @@ impl Execute for Sql {
true true
}; };
let dialect = GenericDialect {}; let dialect = GenericDialect {};
let ast: Result<Vec<Statement>, Error> = match Parser::parse_sql(&dialect, &sql) { let ast: Vec<Statement> = match Parser::parse_sql(&dialect, &sql) {
Ok(ast) => Ok(ast), Ok(ast) => ast,
Err(_) => { Err(_) => {
warn!("Invalid SQL - {}", sql); warn!("Invalid SQL - {}", sql);
Err(SQLInvalidCommand)
}
};
let ast = match ast {
Err(_) => {
let _ = api
.send(
message
.text_reply(format!("❌ Invalid SQL. Syntax error ❌"))
.parse_mode(ParseMode::Html),
)
.await;
return Err(SQLInvalidCommand); return Err(SQLInvalidCommand);
} }
Ok(ast) => ast,
}; };
let msg: Result<String, Error> = match ast.len() { match ast.len() {
l if l > 1 => { l if l > 1 => {
//Max 1 request per message allowed only. return Err(Error::SQLBannedCommand(
Err(Error::SQLBannedCommand) "🚫 One statement per message allowed 🚫".into(),
))
} }
_ => match ast[0] { _ => (),
sqlparser::ast::Statement::Query { .. } => {
let conn = db::open()?;
let x = match conn.prepare_cached(&sql) {
Ok(mut stmt) => {
let query = match stmt.query(rusqlite::NO_PARAMS) {
Err(_) => 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(_)) => {
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 occurred {}❌", 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;
}
_ => {}
},
} }
match ast[0] {
sqlparser::ast::Statement::Query { .. } => {}
_ => {
return Err(Error::SQLBannedCommand(
"🚫 SELECT requests allowed only 🚫".into(),
))
}
}
let conn = db::open()?;
let mut stmt = conn.prepare_cached(&sql)?;
let mut rows = match stmt.query(rusqlite::NO_PARAMS) {
Err(e) => return Err(SQLITE3Error(e)),
Ok(mut rows) => 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(_)) => 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);
}
if res.len() > 100 {
return Err(Error::SQLResultTooLong(
"SQL result too long. Lines limit is 100. Use LIMIT".to_string(),
));
}
// 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(()) Ok(())
} }
#[allow(unused_variables)] #[allow(unused_variables)]
async fn run_mystem( async fn run_mystem(
&self, &self,
api: Api, api: &Api,
message: Message, message: &Message,
mystem: &mut MyStem, mystem: &mut MyStem,
) -> Result<(), Error> { ) -> Result<(), Error> {
unimplemented!() unimplemented!()
@ -234,7 +172,7 @@ impl Execute for Sql {
#[async_trait] #[async_trait]
impl Execute for Here { impl Execute for Here {
async fn run(&self, api: Api, message: Message) -> Result<(), Error> { async fn run(&self, api: &Api, message: &Message) -> Result<(), Error> {
let members: Vec<telegram_bot::User> = db::get_members(message.chat.id()).unwrap(); let members: Vec<telegram_bot::User> = db::get_members(message.chat.id()).unwrap();
for u in &members { for u in &members {
debug!("Found user {:?} in chat {}", u, message.chat.id()); debug!("Found user {:?} in chat {}", u, message.chat.id());
@ -265,8 +203,8 @@ impl Execute for Here {
#[allow(unused_variables)] #[allow(unused_variables)]
async fn run_mystem( async fn run_mystem(
&self, &self,
api: Api, api: &Api,
message: Message, message: &Message,
mystem: &mut MyStem, mystem: &mut MyStem,
) -> Result<(), Error> { ) -> Result<(), Error> {
unimplemented!() unimplemented!()
@ -275,7 +213,7 @@ impl Execute for Here {
#[async_trait] #[async_trait]
impl Execute for Top { impl Execute for Top {
async fn run(&self, api: Api, message: Message) -> Result<(), Error> { async fn run(&self, api: &Api, message: &Message) -> Result<(), Error> {
let top = db::get_top(&message).await?; let top = db::get_top(&message).await?;
let mut msg = "<b>Your top using words:</b>\n<pre>".to_string(); let mut msg = "<b>Your top using words:</b>\n<pre>".to_string();
let mut counter = 1; let mut counter = 1;
@ -300,8 +238,8 @@ impl Execute for Top {
#[allow(unused_variables)] #[allow(unused_variables)]
async fn run_mystem( async fn run_mystem(
&self, &self,
api: Api, api: &Api,
message: Message, message: &Message,
mystem: &mut MyStem, mystem: &mut MyStem,
) -> Result<(), Error> { ) -> Result<(), Error> {
unimplemented!() unimplemented!()
@ -310,7 +248,7 @@ impl Execute for Top {
#[async_trait] #[async_trait]
impl Execute for MarkovAll { impl Execute for MarkovAll {
async fn run(&self, api: Api, message: Message) -> Result<(), Error> { async fn run(&self, api: &Api, message: &Message) -> Result<(), Error> {
let messages = db::get_messages_random_all().await?; let messages = db::get_messages_random_all().await?;
let mut chain = Chain::new(); let mut chain = Chain::new();
chain.feed(messages); chain.feed(messages);
@ -334,8 +272,8 @@ impl Execute for MarkovAll {
#[allow(unused_variables)] #[allow(unused_variables)]
async fn run_mystem( async fn run_mystem(
&self, &self,
api: Api, api: &Api,
message: Message, message: &Message,
mystem: &mut MyStem, mystem: &mut MyStem,
) -> Result<(), Error> { ) -> Result<(), Error> {
unimplemented!() unimplemented!()
@ -344,7 +282,7 @@ impl Execute for MarkovAll {
#[async_trait] #[async_trait]
impl Execute for Markov { impl Execute for Markov {
async fn run(&self, api: Api, message: Message) -> Result<(), Error> { async fn run(&self, api: &Api, message: &Message) -> Result<(), Error> {
let messages = db::get_messages_random_group(&message).await?; let messages = db::get_messages_random_group(&message).await?;
let mut chain = Chain::new(); let mut chain = Chain::new();
chain.feed(messages); chain.feed(messages);
@ -368,8 +306,8 @@ impl Execute for Markov {
#[allow(unused_variables)] #[allow(unused_variables)]
async fn run_mystem( async fn run_mystem(
&self, &self,
api: Api, api: &Api,
message: Message, message: &Message,
mystem: &mut MyStem, mystem: &mut MyStem,
) -> Result<(), Error> { ) -> Result<(), Error> {
unimplemented!() unimplemented!()
@ -379,15 +317,15 @@ impl Execute for Markov {
#[async_trait] #[async_trait]
impl Execute for Omedeto { impl Execute for Omedeto {
#[allow(unused_variables)] #[allow(unused_variables)]
async fn run(&self, api: Api, message: Message) -> Result<(), Error> { async fn run(&self, api: &Api, message: &Message) -> Result<(), Error> {
unimplemented!() unimplemented!()
} }
#[warn(unused_must_use)] #[warn(unused_must_use)]
async fn run_mystem( async fn run_mystem(
&self, &self,
api: Api, api: &Api,
message: Message, message: &Message,
mystem: &mut MyStem, mystem: &mut MyStem,
) -> Result<(), Error> { ) -> Result<(), Error> {
let all_msg = db::get_messages_user_all(&message).await?; let all_msg = db::get_messages_user_all(&message).await?;

4
src/errors.rs
View File

@ -20,9 +20,11 @@ pub enum Error {
JsonParseError(serde_error), JsonParseError(serde_error),
PopenError(popen_error), PopenError(popen_error),
MystemError(mystem_error), MystemError(mystem_error),
SQLBannedCommand, SQLBannedCommand(String),
SQLInvalidCommand, SQLInvalidCommand,
SQLResultTooLong(String),
} }
impl fmt::Display for Error { impl fmt::Display for Error {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "An error occurred.") write!(f, "An error occurred.")

31
src/handlers.rs
View File

@ -31,49 +31,60 @@ pub async fn handler(
Here { Here {
data: "".to_string(), data: "".to_string(),
} }
.run(api, message) .run(&api, &message)
.await? .await?
} }
s if s.to_string().starts_with("/sql") => { s if s.to_string().starts_with("/sql") => match {
Sql { Sql {
data: s.replace("/sql ", ""), data: s.replace("/sql ", ""),
} }
.run(api, message) .run(&api, &message)
.await? .await
} } {
Ok(_) => debug!("/sql command sent to {}", message.chat.id()),
Err(e) => {
api.send(
message
.text_reply(format!("Error: {:#?}", e))
.parse_mode(ParseMode::Html),
)
.await?;
()
}
},
"/top" => { "/top" => {
Top { Top {
data: "".to_string(), data: "".to_string(),
} }
.run(api, message) .run(&api, &message)
.await? .await?
} }
"/stat" => { "/stat" => {
Top { Top {
data: "".to_string(), data: "".to_string(),
} }
.run(api, message) .run(&api, &message)
.await? .await?
} }
"/markov_all" => { "/markov_all" => {
MarkovAll { MarkovAll {
data: "".to_string(), data: "".to_string(),
} }
.run(api, message) .run(&api, &message)
.await? .await?
} }
"/markov" => { "/markov" => {
Markov { Markov {
data: "".to_string(), data: "".to_string(),
} }
.run(api, message) .run(&api, &message)
.await? .await?
} }
"/omedeto" => { "/omedeto" => {
Omedeto { Omedeto {
data: "".to_string(), data: "".to_string(),
} }
.run_mystem(api, message, mystem) .run_mystem(&api, &message, mystem)
.await? .await?
} }
_ => (), _ => (),