// Libraries use std::error::Error; use mysql::{prelude::Queryable, Opts, Pool, PooledConn, Row}; use chrono::Utc; use regex::Regex; // Structures pub struct Database { connection: Option } // Functions impl Database { // Constructors pub fn init() -> Database { // Returning a database return Database { connection: None }; } // Functions fn check_injection(&mut self, query: &str) -> Result, Box> { let sql_injection_pattern: Regex = Regex::new(r"(?i)(--|;|/\*|\*/|xp_|exec|select|insert|update|delete|drop|union|shutdown|create|alter)")?; if sql_injection_pattern.is_match(query) { // Getting time let now: chrono::DateTime = Utc::now(); let formatted: String = now.format("%Y-%m-%d %H:%M:%S").to_string(); // Add to flag database self.inst_table("Flags", "`occurance` DATETIME")?; self.insert("Flags", "`occurance`", &format!("'{}'", formatted))?; // Return error return Ok(Some(format!("SQL Injection Detected ({})", query))); } else { return Ok(None); } } pub fn connect(&mut self, address: String, user: String, pass: String) -> Result<(), Box> { // Creating database connection url let url: String = format!("mysql://{}:{}@{}/neurostock", user, pass, address); // Connecting via pool let pool: Pool = Pool::new(Opts::from_url(&url)?)?; // Creating the connection self.connection = Some(pool.get_conn()?); // Success return Ok(()); } pub fn inst_table(&mut self, table: &str, columns: &str) -> Result<(), Box> { // Preparing sql query let query: String = format!("CREATE TABLE IF NOT EXISTS {} ({})", table, columns); // Verify we have a connection match &mut self.connection { Some(conn) => { // Creating table if it doesn't exist conn.query_drop(query)?; }, None => { panic!("Database> No database connection to check or create table."); } }; // Successful return Ok(()); } pub fn exists_table(&mut self, table: &str) -> Result> { // Result let result: bool; // Compile SQL query let query: String = format!("SELECT EXISTS (SELECT 1 FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = '{}') AS status;", table); // Sending to get result match &mut self.connection { Some(conn) => { // Getting the first row let row: Row = conn.query_iter(query)?.nth(0).unwrap()?; // Was it successful? result = row.get(0).unwrap(); }, None => { panic!("Database> No database connection to insert data into."); } } // Success return Ok(result); } pub fn exists_row(&mut self, table: &str, element: &str, value: &str) -> Result> { // Preparing sql query let response: Vec = self.get( table, "*", &format!(" WHERE {}={}", element, value) )?; // Success return Ok(response.len() != 0); } pub fn insert(&mut self, table: &str, columns: &str, data: &str) -> Result<(), Box> { // Checking for SQL Injection match self.check_injection(data)? { Some(e) => { return Err(e.into()); }, None => {} } // Preparing sql query let query: String = format!("INSERT INTO {} ({}) VALUES ({})", table, columns, data); // Verify we have a connection match &mut self.connection { Some(conn) => { // Sending query conn.query_drop(query)?; }, None => { panic!("Database> No database connection to insert data into."); } }; // Successful return Ok(()); } pub fn update(&mut self, table: &str, element: &str, value: &str, case: &str) -> Result<(), Box> { // Checking for SQL Injection match self.check_injection(value)? { Some(e) => { return Err(e.into()); }, None => {} } // Checking for SQL Injection match self.check_injection(case)? { Some(e) => { return Err(e.into()); }, None => {} } // Preparing sql query let query: String = format!("UPDATE {} SET {}={} WHERE {}", table, element, value, case); // Verify we have a connection match &mut self.connection { Some(conn) => { // Sending query conn.query_drop(query)?; }, None => { panic!("Database> No database connection to insert data into."); } }; // Successful return Ok(()); } pub fn get(&mut self, table: &str, elements: &str, case: &str) -> Result, Box> { // Is there a case? if case.contains("WHERE") { // Checking for SQL Injection match self.check_injection(&case.split("WHERE").nth(1).unwrap())? { Some(e) => { return Err(e.into()); }, None => {} } } // Preparing sql query let query: String = format!("SELECT {} FROM {}{}", elements, table, case); // Creating vector to store the data in let mut result: Vec = Vec::new(); // Verify we have a connection match &mut self.connection { Some(conn) => { // Going through all rows it sent back for row in conn.query_iter(query)? { result.push(row?); } }, None => { panic!("Database> No database connection to insert data into."); } }; // Successful return Ok(result); } }