StockMAPI/project/src/base/database.rs

// 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<PooledConn>
}

// Functions
impl Database {
    // Constructors
    pub fn init() -> Database {
        // Returning a database
        return Database {
            connection: None
        };
    }

    // Functions
    fn check_injection(&mut self, query: &str) -> Result<Option<String>, Box<dyn Error>> {
        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> = 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<dyn Error>> {
        // 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<dyn Error>> {
        // 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<bool, Box<dyn Error>> {
        // 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<bool, Box<dyn Error>> {
        // Preparing sql query
        let response: Vec<Row> = 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<dyn Error>> {
        // 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<dyn Error>> {
        // 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<Vec<Row>, Box<dyn Error>> {
        // 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<Row> = 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);
    }
}