Upload files to "src"

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<Self, Box<dyn std::error::Error>> {
+        // 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)
+    }
+}

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<mongodb::bson::Document>,
+    pub partial_coll: Collection<mongodb::bson::Document>,
+}
+
+impl Database {
+    pub async fn connect(config: &Config) -> Result<Self, mongodb::error::Error> {
+        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<mongodb::bson::Document>, full: bool) -> Result<(), mongodb::error::Error> {
+        let collection = if full { &self.full_coll } else { &self.partial_coll };
+        collection.insert_many(records).await?;
+        Ok(())
+    }
+}

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<dyn std::error::Error>>{
+    // Read command-line arguments
+    let args: Vec<String> = env::args().collect();
+    if args.len() != 2 {
+        eprintln!("Usage: ./pwnWriter <input_file>");
+        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(())
+}

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)
+    }
+}

src/utils.rs

@@ -0,0 +1,22 @@
+use std::process::Command;
+use std::str::FromStr;
+
+pub fn get_line_count(filename: &str) -> Result<usize, Box<dyn std::error::Error>> {
+    // 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)
+}