feat: Modules now run in separate threads (#4)

To be honest it was hell of a mess, not sure if it's an ideal solution but hey; it works :D

src/events.rs

@@ -3,6 +3,7 @@ use std::fmt;
 #[derive(Debug, Eq, Hash, PartialEq)]
 pub enum Type {
     Ready,
+    FinishedTask,
     /// String: the discovered domain
     DiscoveredDomain(String),
 }
@@ -13,6 +14,9 @@ impl fmt::Display for Type {
             Type::Ready => {
                 write!(formatter, "ready")
             }
+            Type::FinishedTask => {
+                write!(formatter, "finished:task")
+            }
             Type::DiscoveredDomain(_) => {
                 write!(formatter, "discovered:domain")
             }

src/main.rs

@@ -9,6 +9,7 @@ mod events;
 mod logger;
 mod modules;
 mod session;
+mod state;
 
 fn main() {
     config::write_default_config_if_not_existing();
@@ -24,7 +25,7 @@ fn main() {
     };
 
     let (tx, rx) = flume::bounded::<events::Type>(100);
-    let mut session = session::Session::new(args, config, tx, rx);
+    let session = session::Session::new(args, config, tx, rx);
     session.register_default_modules();
     session.start();
 }

src/modules/enumerate_subdomains/mod.rs

@@ -60,13 +60,13 @@ impl Module for ModuleEnumerateSubdomains {
                 .header(USER_AGENT, random_user_agent)
                 .send()
                 .is_ok()
-                && !session.has_discovered_subdomain(uri.clone())
+                && !session.get_state().has_discovered_subdomain(uri.clone())
             {
                 logger::println(
                     self.name(),
                     format!("Discovered '{}' as a new subdomain", uri),
                 );
-                session.discover_subdomain(uri);
+                session.get_state().discover_subdomain(uri);
             }
         }
     }

src/modules/mod.rs

@@ -6,7 +6,7 @@ pub mod enumerate_subdomains;
 pub mod passive_dns;
 pub mod ready;
 
-pub trait Module {
+pub trait Module: Send + Sync {
     fn name(&self) -> String;
     #[allow(dead_code)]
     fn description(&self) -> String;

src/modules/passive_dns/mod.rs

@@ -85,7 +85,10 @@ impl Module for ModulePassiveDNS {
                         if name_value == &domain.to_string() {
                             continue;
                         }
-                        if !session.has_discovered_subdomain(name_value.to_string()) {
+                        if !session
+                            .get_state()
+                            .has_discovered_subdomain(name_value.to_string())
+                        {
                             let now = Utc::now();
 
                             // Check if the certificate has expired
@@ -118,7 +121,9 @@ impl Module for ModulePassiveDNS {
                                     }
                                 ),
                             );
-                            session.discover_subdomain(name_value.to_string());
+                            session
+                                .get_state()
+                                .discover_subdomain(name_value.to_string());
                         }
                     }
                 }

src/session.rs

@@ -1,20 +1,20 @@
-use std::sync::Mutex;
+use std::sync::{Arc, Mutex};
+use std::thread;
 
 use flume::{Receiver, Sender};
 
 use crate::modules::Module;
-use crate::{args, config, events, logger, modules};
+use crate::{args, config, events, logger, modules, state};
 
 pub struct Session {
     args: args::Args,
     config: config::Config,
+    state: Arc<state::State>,
+
     sender: Sender<events::Type>,
     receiver: Receiver<events::Type>,
 
-    modules: Vec<Box<dyn Module>>,
-    discovered_subdomains: Mutex<Vec<String>>,
-
-    dev_mode: bool,
+    modules: Mutex<Vec<Arc<Box<dyn Module>>>>,
 }
 
 impl Session {
@@ -23,47 +23,40 @@ impl Session {
         config: config::Config,
         sender: Sender<events::Type>,
         receiver: Receiver<events::Type>,
-    ) -> Self {
-        Session {
+    ) -> Arc<Self> {
+        Arc::new(Session {
             args,
             config,
+            state: Arc::new(state::State::new(false)),
+
             sender,
             receiver,
 
-            modules: Vec::new(),
-            discovered_subdomains: Mutex::new(Vec::new()),
-
-            dev_mode: false,
-        }
-    }
-
-    pub fn get_config(&self) -> &config::Config {
-        &self.config
+            modules: Mutex::new(Vec::new()),
+        })
     }
 
     pub fn get_args(&self) -> &args::Args {
         &self.args
     }
 
-    pub fn discover_subdomain(&self, subdomain: String) {
-        self.discovered_subdomains.lock().unwrap().push(subdomain)
+    pub fn get_config(&self) -> &config::Config {
+        &self.config
     }
 
-    pub fn has_discovered_subdomain(&self, subdomain: String) -> bool {
-        self.discovered_subdomains
-            .lock()
-            .unwrap()
-            .contains(&subdomain)
+    pub fn get_state(&self) -> Arc<state::State> {
+        Arc::clone(&self.state)
     }
 
-    pub fn register_module<T: Module + 'static>(&mut self, module: T) {
-        if self.dev_mode {
+    pub fn register_module<T: Module + Send + Sync + 'static>(&self, module: T) {
+        if self.get_state().is_verbose() {
             logger::debug("", format!("Registered module {}", module.name()))
         }
-        self.modules.push(Box::new(module));
+        let mut modules = self.modules.lock().unwrap();
+        modules.push(Arc::new(Box::new(module)));
     }
 
-    pub fn register_default_modules(&mut self) {
+    pub fn register_default_modules(&self) {
         self.register_module(modules::ready::ModuleReady::new());
         // self.register_module(modules::enumerate_files::ModuleEnumerateFiles::new());
         // self.register_module(modules::enumerate_subdomains::ModuleEnumerateSubdomains::new());
@@ -80,12 +73,30 @@ impl Session {
         };
     }
 
-    pub fn start(&mut self) {
+    pub fn start(self: Arc<Self>) {
         self.emit(events::Type::Ready);
         self.emit(events::Type::DiscoveredDomain(self.args.domain.clone()));
+        self.emit(events::Type::DiscoveredDomain(self.args.domain.clone()));
 
         while let Ok(event) = self.receiver.recv() {
-            for module in &self.modules {
+            if event == events::Type::FinishedTask {
+                self.get_state().decrement_tasks();
+                if self.get_state().is_verbose() {
+                    logger::debug(
+                        event.to_string(),
+                        format!(
+                            "A task has finished execution, current running tasks: {}",
+                            self.get_state().active_tasks_count()
+                        ),
+                    );
+                }
+                if self.get_state().active_tasks_count() == 0 {
+                    logger::println("finished", "All modules have finished execution, leaving!");
+                    break;
+                }
+            }
+            let modules = self.modules.lock().unwrap();
+            for module in &*modules {
                 if module
                     .subscribers()
                     .iter()
@@ -97,7 +108,13 @@ impl Session {
                         _ => false,
                     })
                 {
-                    module.execute(&self);
+                    self.get_state().increment_tasks();
+                    let module_clone = Arc::clone(module);
+                    let session_clone = Arc::clone(&self);
+                    thread::spawn(move || {
+                        module_clone.execute(&session_clone);
+                        session_clone.emit(events::Type::FinishedTask);
+                    });
                 }
             }
         }

src/state.rs

@@ -0,0 +1,69 @@
+use std::sync::{
+    atomic::{AtomicUsize, Ordering},
+    Arc, Mutex,
+};
+
+use crate::logger;
+
+pub struct State {
+    active_tasks: Arc<AtomicUsize>,
+    verbose: bool,
+
+    discovered_subdomains: Mutex<Vec<String>>,
+}
+
+impl State {
+    pub fn new(verbose: bool) -> Self {
+        State {
+            active_tasks: Arc::new(AtomicUsize::new(0)),
+            verbose,
+
+            discovered_subdomains: Mutex::new(Vec::new()),
+        }
+    }
+
+    pub fn is_verbose(&self) -> bool {
+        self.verbose
+    }
+
+    pub fn increment_tasks(&self) {
+        self.active_tasks.fetch_add(1, Ordering::SeqCst);
+        if self.is_verbose() {
+            logger::debug(
+                "state",
+                format!(
+                    "The running task counter has incremented, current running tasks: {}",
+                    self.active_tasks_count()
+                ),
+            );
+        }
+    }
+
+    pub fn decrement_tasks(&self) {
+        self.active_tasks.fetch_sub(1, Ordering::SeqCst);
+        if self.is_verbose() {
+            logger::debug(
+                "state",
+                format!(
+                    "The running task counter has decremented, current running tasks: {}",
+                    self.active_tasks_count()
+                ),
+            );
+        }
+    }
+
+    pub fn active_tasks_count(&self) -> usize {
+        self.active_tasks.load(Ordering::SeqCst)
+    }
+
+    pub fn discover_subdomain(&self, subdomain: String) {
+        self.discovered_subdomains.lock().unwrap().push(subdomain)
+    }
+
+    pub fn has_discovered_subdomain(&self, subdomain: String) -> bool {
+        self.discovered_subdomains
+            .lock()
+            .unwrap()
+            .contains(&subdomain)
+    }
+}