diff --git a/Cargo.toml b/Cargo.toml index 1d81b05..12f8b6b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "dumpsync" -version = "0.1.17" +version = "0.1.18" edition = "2021" license = "MIT" authors = ["kremilly"] diff --git a/src/engine/export.rs b/src/engine/export.rs index e9a7037..dfbd185 100644 --- a/src/engine/export.rs +++ b/src/engine/export.rs @@ -16,6 +16,7 @@ use crate::{ helpers::{ configs::Configs, export_handlers::ExportHandlers, + queries_builders::QueriesBuilders, }, }; @@ -78,7 +79,7 @@ impl Export { export_handlers.comments_header(writer.as_write())?; export_handlers.write_create_new_database(writer.as_write())?; - let tables: Vec = conn.query("SHOW TABLES")?; + let tables: Vec = conn.query(QueriesBuilders.show_tables())?; let ignore_tables = Configs.list("exports", "ignore_tables").unwrap_or_default(); for table in tables { diff --git a/src/engine/scan_xss.rs b/src/engine/scan_xss.rs index 5ec7c42..9c76013 100644 --- a/src/engine/scan_xss.rs +++ b/src/engine/scan_xss.rs @@ -12,7 +12,11 @@ use crate::{ consts::global::Global, ui::scan_alerts::ScanAlerts, engine::connection::Connection, - helpers::scan_handlers::ScanHandlers, + + helpers::{ + scan_handlers::ScanHandlers, + queries_builders::QueriesBuilders, + }, }; pub struct ScanXSS { @@ -82,7 +86,7 @@ impl ScanXSS { } }; - let query = ScanHandlers.build_query(&self.table, self.offset.map(|o| o as usize), self.limit.map(|l| l as usize)); + let query = QueriesBuilders.select(&self.table, self.offset.map(|o| o as usize), self.limit.map(|l| l as usize)); let rows: Vec = conn.query(query)?; for (row_index, row) in rows.iter().enumerate() { diff --git a/src/helpers/export_handlers.rs b/src/helpers/export_handlers.rs index 90a5b76..18e42d3 100644 --- a/src/helpers/export_handlers.rs +++ b/src/helpers/export_handlers.rs @@ -21,7 +21,11 @@ use mysql::{ use crate::{ utils::date::Date, - helpers::configs::Configs, + + helpers::{ + configs::Configs, + queries_builders::QueriesBuilders, + }, }; pub enum Writer { @@ -43,6 +47,7 @@ impl Writer { pub struct ExportHandlers { file: File, dbname: String, + dump_data: bool, compress_data: bool, insert_ignore_into: bool, @@ -67,10 +72,21 @@ impl ExportHandlers { pub fn create_writer(&self) -> Result { if self.compress_data { - let encoder = GzEncoder::new(self.file.try_clone()?, Compression::default()); - Ok(Writer::Compressed(BufWriter::new(encoder))) + let encoder = GzEncoder::new( + self.file.try_clone()?, Compression::default() + ); + + Ok( + Writer::Compressed( + BufWriter::new(encoder) + ) + ) } else { - Ok(Writer::Uncompressed(BufWriter::new(self.file.try_clone()?))) + Ok( + Writer::Uncompressed( + BufWriter::new(self.file.try_clone()?) + ) + ) } } @@ -85,8 +101,10 @@ impl ExportHandlers { pub fn write_create_new_database(&self, writer: &mut dyn Write) -> Result<(), Box> { if self.database_if_not_exists { - writeln!(writer, "CREATE DATABASE IF NOT EXISTS `{}`;", self.dbname)?; - writeln!(writer, "USE `{}`;", self.dbname)?; + let queries = QueriesBuilders.create_database(&self.dbname)?; + + write!(writer, "{}", queries.0)?; + writeln!(writer, "{}", queries.1)?; writeln!(writer, "-- ---------------------------------------------------\n")?; } @@ -95,7 +113,9 @@ impl ExportHandlers { pub fn write_inserts_for_table(&self, table: &str, conn: &mut PooledConn, writer: &mut dyn Write) -> Result<(), Box> { if self.dump_data { - let rows: Vec = conn.query(format!("SELECT * FROM `{}`", table))?; + let rows: Vec = conn.query( + QueriesBuilders.select(table, None, None) + )?; if rows.is_empty() { writeln!(writer, "-- Table `{}` contains no data.", table)?; @@ -111,9 +131,9 @@ impl ExportHandlers { }).collect(); let line = if self.insert_ignore_into { - format!("INSERT IGNORE INTO `{}` VALUES ({});", table, values.join(", ")) + QueriesBuilders.insert_into(table, values, true) } else { - format!("INSERT INTO `{}` VALUES ({});", table, values.join(", ")) + QueriesBuilders.insert_into(table, values, false) }; writeln!(writer, "{}", line)?; @@ -128,10 +148,10 @@ impl ExportHandlers { writeln!(writer, "-- Exporting the table: `{}`", table)?; if self.drop_table_if_exists { - writeln!(writer, "DROP TABLE IF EXISTS `{}`;", table)?; + writeln!(writer, "{}", QueriesBuilders.drop_table(table))?; } - let row: Row = conn.query_first(format!("SHOW CREATE TABLE `{}`", table))?.unwrap(); + let row: Row = conn.query_first(QueriesBuilders.show_create_table(table))?.unwrap(); let create_table: String = row.get(1).expect("Error retrieving CREATE TABLE"); writeln!(writer, "{};\n", create_table)?; diff --git a/src/helpers/mod.rs b/src/helpers/mod.rs index 4254072..7f41897 100644 --- a/src/helpers/mod.rs +++ b/src/helpers/mod.rs @@ -3,4 +3,5 @@ pub mod configs; pub mod dump_handlers; pub mod scan_handlers; pub mod export_handlers; -pub mod import_handlers; \ No newline at end of file +pub mod import_handlers; +pub mod queries_builders; \ No newline at end of file diff --git a/src/helpers/queries_builders.rs b/src/helpers/queries_builders.rs new file mode 100644 index 0000000..41cfd8c --- /dev/null +++ b/src/helpers/queries_builders.rs @@ -0,0 +1,46 @@ +pub struct QueriesBuilders; + +impl QueriesBuilders { + + pub fn show_tables(&self) -> String { + "SHOW TABLES".to_string() + } + + pub fn show_create_table(&self, table: &str) -> String { + format!("SHOW CREATE TABLE `{}`;", table) + } + + pub fn drop_table(&self, table: &str) -> String { + format!("DROP TABLE IF EXISTS `{}`;", table) + } + + pub fn create_database(&self, dbname: &str) -> Result<(String, String), String> { + let create_db = format!("CREATE DATABASE IF NOT EXISTS `{}`;\n", dbname); + let use_db = format!("USE `{}`;", dbname); + + Ok((create_db, use_db)) + } + + pub fn insert_into(&self, table: &str, values: Vec, ignore: bool) -> String { + if ignore { + format!("INSERT IGNORE INTO `{}` VALUES ({});", table, values.join(", ")) + } else { + format!("INSERT INTO `{}` VALUES ({});", table, values.join(", ")) + } + } + + pub fn select(&self, table: &str, offset: Option, limit: Option) -> String { + let mut query = format!("SELECT * FROM `{}`", table); + + if let Some(l) = limit { + query.push_str(&format!(" LIMIT {}", l)); + } + + if let Some(o) = offset { + query.push_str(&format!(" OFFSET {}", o)); + } + + query + } + +} diff --git a/src/helpers/scan_handlers.rs b/src/helpers/scan_handlers.rs index 47d3f89..d2312de 100644 --- a/src/helpers/scan_handlers.rs +++ b/src/helpers/scan_handlers.rs @@ -14,21 +14,7 @@ use std::{ pub struct ScanHandlers; -impl ScanHandlers { - - pub fn build_query(&self, table: &str, offset: Option, limit: Option) -> String { - let mut query = format!("SELECT * FROM `{}`", table); - - if let Some(l) = limit { - query.push_str(&format!(" LIMIT {}", l)); - } - - if let Some(o) = offset { - query.push_str(&format!(" OFFSET {}", o)); - } - - query - } +impl ScanHandlers { pub fn load_patterns_from_file(&self, path: &str) -> Result, Box> { let path = Path::new(path);