commit c1aecbe8b27e752fdce556d844ed565f89ec6fe6 Author: CatWithAHat Date: Sat Mar 29 07:09:42 2025 +0000 Upload files to "src" diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..3daf292 --- /dev/null +++ b/src/config.rs @@ -0,0 +1,30 @@ +use serde::Deserialize; +use std::fs; +use toml; + +#[derive(Debug, Deserialize)] +pub struct Config { + pub mongodb: MongoConfig, +} + +#[derive(Debug, Deserialize)] +pub struct MongoConfig { + pub uri: String, + pub database: String, + pub full_collection: String, + pub partial_collection: String, + pub batch_size: usize, +} + +impl Config { + pub fn load() -> Result> { + // Construct the path to ~/.config/pwnwriter/config.toml + let mut config_path = dirs::home_dir().ok_or("Cannot determine home directory")?; + config_path.push(".config/pwnwriter/config.toml"); + + // Read the configuration file + let contents = fs::read_to_string(config_path)?; + let config: Config = toml::from_str(&contents)?; + Ok(config) + } +} diff --git a/src/database.rs b/src/database.rs new file mode 100644 index 0000000..9476973 --- /dev/null +++ b/src/database.rs @@ -0,0 +1,27 @@ +use mongodb::{options::ClientOptions, Client, Collection}; +use crate::config::Config; + +#[derive(Clone)] +pub struct Database { + pub full_coll: Collection, + pub partial_coll: Collection, +} + +impl Database { + pub async fn connect(config: &Config) -> Result { + let client_options = ClientOptions::parse(&config.mongodb.uri).await?; + let client = Client::with_options(client_options)?; + let db = client.database(&config.mongodb.database); + + Ok(Database { + full_coll: db.collection(&config.mongodb.full_collection), + partial_coll: db.collection(&config.mongodb.partial_collection), + }) + } + + pub async fn insert_many(&self, records: Vec, full: bool) -> Result<(), mongodb::error::Error> { + let collection = if full { &self.full_coll } else { &self.partial_coll }; + collection.insert_many(records).await?; + Ok(()) + } +} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..a7ac920 --- /dev/null +++ b/src/main.rs @@ -0,0 +1,56 @@ +mod config; +mod parser; +mod database; +mod file_handler; +mod progress; +mod utils; + +use std::env; +//use std::sync::Arc; +use tokio; +//use tokio::sync::Mutex; +use crate::config::Config; +use crate::database::Database; +use crate::file_handler::process_file; +//use crate::progress::ProgressTracker; +use crate::utils::get_line_count; + +#[tokio::main] +async fn main() -> Result<(), Box>{ + // Read command-line arguments + let args: Vec = env::args().collect(); + if args.len() != 2 { + eprintln!("Usage: ./pwnWriter "); + std::process::exit(1); + } + let input_file = &args[1]; + let total_lines = get_line_count(input_file)?; + //let progress = Arc::new(Mutex::new(ProgressTracker::new(Some(total_lines)))); + println!("Total lines in {}: {}", input_file, total_lines); + + // Load configuration + let config = match Config::load() { + Ok(cfg) => cfg, + Err(e) => { + eprintln!("Error loading configuration: {}", e); + std::process::exit(1); + } + }; + + // Initialize MongoDB + let db = match Database::connect(&config).await { + Ok(database) => database, + Err(e) => { + eprintln!("Error connecting to MongoDB: {}", e); + std::process::exit(1); + } + }; + + // Process the input file + if let Err(e) = process_file(input_file, &db, &config, total_lines).await { + eprintln!("Error processing file: {}", e); + std::process::exit(1); + } + + Ok(()) +} diff --git a/src/parser.rs b/src/parser.rs new file mode 100644 index 0000000..932b4d8 --- /dev/null +++ b/src/parser.rs @@ -0,0 +1,30 @@ +#[derive(Debug)] +pub enum ParsedRecord { + Full { url: String, username: String, password: String }, + Partial { username: String, password: String }, + Malformed(String), +} + +pub fn parse_line(line: &str) -> ParsedRecord { + let parts: Vec<&str> = line.split(':').collect(); + match parts.len() { + 3 => ParsedRecord::Full { + url: normalize_url(parts[0]), + username: parts[1].to_string(), + password: parts[2].to_string(), + }, + 2 => ParsedRecord::Partial { + username: parts[0].to_string(), + password: parts[1].to_string(), + }, + _ => ParsedRecord::Malformed(line.to_string()), + } +} + +fn normalize_url(url: &str) -> String { + if url.starts_with("http://") || url.starts_with("https://") { + url.to_string() + } else { + format!("https://{}", url) + } +} diff --git a/src/utils.rs b/src/utils.rs new file mode 100644 index 0000000..c756e56 --- /dev/null +++ b/src/utils.rs @@ -0,0 +1,22 @@ +use std::process::Command; +use std::str::FromStr; + +pub fn get_line_count(filename: &str) -> Result> { + // Run "wc -l" on the file. + let output = Command::new("wc") + .arg("-l") + .arg(filename) + .output()?; + + if !output.status.success() { + return Err(format!("wc command failed with status: {:?}", output.status).into()); + } + + // Convert output to a String. + let stdout = String::from_utf8(output.stdout)?; + // The output format is typically "1234 filename", so split and parse the first token. + let count_str = stdout.split_whitespace().next().ok_or("Invalid output from wc")?; + let count = usize::from_str(count_str)?; + + Ok(count) +} \ No newline at end of file