feat: Add Context enum and semaphore field (#5)

kkrypt0nn · Nov 2, 2024 · f995807 · f995807
1 parent ec844a4
commit f995807
Show file tree

Hide file tree

Showing 9 changed files with 69 additions and 27 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -24,7 +24,8 @@ flume = "0.11.1"
 heck = "0.5.0"
 lazy_static = "1.5.0"
 rand = "0.8.5"
-reqwest = { version = "0.12.8", features = ["blocking", "json"] }
-serde = { version = "1.0.210", features = ["derive"] }
+reqwest = { version = "0.12.9", features = ["blocking", "json"] }
+serde = { version = "1.0.214", features = ["derive"] }
+simple-semaphore = "0.1.0"
 tangra = "1.3.0"
 toml = "0.8.19"
diff --git a/src/events.rs b/src/events.rs
@@ -1,6 +1,6 @@
 use std::fmt;
 
-#[derive(Debug, Eq, Hash, PartialEq)]
+#[derive(Clone, Debug, Eq, Hash, PartialEq)]
 pub enum Type {
     Ready,
     FinishedTask,

diff --git a/src/modules/enumerate_files/mod.rs b/src/modules/enumerate_files/mod.rs
@@ -5,7 +5,7 @@ use std::io::{BufRead, BufReader};
 
 use reqwest::header::USER_AGENT;
 
-use crate::modules::Module;
+use crate::modules::{Context, Module};
 use crate::session::Session;
 use crate::{events, logger};
 
@@ -36,9 +36,15 @@ impl Module for ModuleEnumerateFiles {
         vec![events::Type::DiscoveredDomain(String::new())]
     }
 
-    fn execute(&self, session: &Session) {
+    fn execute(&self, session: &Session, context: Context) {
+        let domain = match context {
+            Context::Domain(domain) => domain,
+            Context::None => {
+                logger::error(self.name(), "Received wrong context, exiting module");
+                return;
+            }
+        };
         let args = session.get_args();
-        let domain = &args.domain;
         let config = session.get_config();
         let wordlist = config
             .enumerate_files

diff --git a/src/modules/enumerate_subdomains/mod.rs b/src/modules/enumerate_subdomains/mod.rs
@@ -4,7 +4,7 @@ use std::io::{BufRead, BufReader};
 
 use reqwest::header::USER_AGENT;
 
-use crate::modules::Module;
+use crate::modules::{Context, Module};
 use crate::session::Session;
 use crate::{events, logger};
 
@@ -37,9 +37,15 @@ impl Module for ModuleEnumerateSubdomains {
         vec![events::Type::DiscoveredDomain(String::new())]
     }
 
-    fn execute(&self, session: &Session) {
+    fn execute(&self, session: &Session, context: Context) {
+        let domain = match context {
+            Context::Domain(domain) => domain,
+            Context::None => {
+                logger::error(self.name(), "Received wrong context, exiting module");
+                return;
+            }
+        };
         let args = session.get_args();
-        let domain = &args.domain;
         let config = session.get_config();
         let wordlist = config
             .enumerate_subdomains

diff --git a/src/modules/mod.rs b/src/modules/mod.rs
@@ -6,10 +6,22 @@ pub mod enumerate_subdomains;
 pub mod passive_dns;
 pub mod ready;
 
+pub enum Context {
+    Domain(String),
+    None,
+}
+
+pub fn get_context_for_event(event: &events::Type) -> Context {
+    match event {
+        events::Type::DiscoveredDomain(domain) => Context::Domain(domain.clone()),
+        _ => Context::None,
+    }
+}
+
 pub trait Module: Send + Sync {
     fn name(&self) -> String;
     #[allow(dead_code)]
     fn description(&self) -> String;
     fn subscribers(&self) -> Vec<events::Type>;
-    fn execute(&self, session: &Session);
+    fn execute(&self, session: &Session, context: Context);
 }
diff --git a/src/modules/passive_dns/mod.rs b/src/modules/passive_dns/mod.rs
@@ -5,7 +5,7 @@ use std::sync::Mutex;
 use reqwest::header::USER_AGENT;
 
 use crate::modules::passive_dns::crt_sh::CrtShItem;
-use crate::modules::Module;
+use crate::modules::{Context, Module};
 use crate::session::Session;
 use crate::{events, logger};
 
@@ -52,9 +52,14 @@ impl Module for ModulePassiveDNS {
         vec![events::Type::DiscoveredDomain(String::new())]
     }
 
-    fn execute(&self, session: &Session) {
-        let args = session.get_args();
-        let domain = &args.domain;
+    fn execute(&self, session: &Session, context: Context) {
+        let domain = match context {
+            Context::Domain(domain) => domain,
+            Context::None => {
+                logger::error(self.name(), "Received wrong context, exiting module");
+                return;
+            }
+        };
         let config = session.get_config();
         let ignore_expired = config.passive_dns.ignore_expired.unwrap_or_else(|| false);
         let recent_only = config.passive_dns.recent_only.unwrap_or_else(|| false);

diff --git a/src/modules/ready/mod.rs b/src/modules/ready/mod.rs
@@ -1,6 +1,6 @@
 use crate::events;
 use crate::logger;
-use crate::modules::Module;
+use crate::modules::{Context, Module};
 use crate::session::Session;
 
 pub struct ModuleReady {}
@@ -30,7 +30,7 @@ impl Module for ModuleReady {
         vec![events::Type::Ready]
     }
 
-    fn execute(&self, _: &Session) {
+    fn execute(&self, _: &Session, _: Context) {
         logger::println(
             "ready",
             "Project Serpens is now ready and will start doing its magic!",

diff --git a/src/session.rs b/src/session.rs
@@ -76,7 +76,6 @@ impl Session {
     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() {
             if event == events::Type::FinishedTask {
@@ -108,12 +107,17 @@ impl Session {
                         _ => false,
                     })
                 {
-                    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);
+                    thread::spawn({
+                        self.get_state().increment_tasks();
+                        let permit = self.get_state().get_semaphore_permit();
+                        let context = modules::get_context_for_event(&event);
+                        let module_clone = Arc::clone(module);
+                        let session_clone = Arc::clone(&self);
+                        move || {
+                            module_clone.execute(&session_clone, context);
+                            session_clone.emit(events::Type::FinishedTask);
+                            drop(permit);
+                        }
                     });
                 }
             }

diff --git a/src/state.rs b/src/state.rs
@@ -3,10 +3,13 @@ use std::sync::{
     Arc, Mutex,
 };
 
+use simple_semaphore::Permit;
+
 use crate::logger;
 
 pub struct State {
     active_tasks: Arc<AtomicUsize>,
+    semaphore: Arc<simple_semaphore::Semaphore>,
     verbose: bool,
 
     discovered_subdomains: Mutex<Vec<String>>,
@@ -16,16 +19,13 @@ impl State {
     pub fn new(verbose: bool) -> Self {
         State {
             active_tasks: Arc::new(AtomicUsize::new(0)),
+            semaphore: simple_semaphore::Semaphore::default(),
             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() {
@@ -56,6 +56,14 @@ impl State {
         self.active_tasks.load(Ordering::SeqCst)
     }
 
+    pub fn get_semaphore_permit(&self) -> Permit {
+        self.semaphore.acquire()
+    }
+
+    pub fn is_verbose(&self) -> bool {
+        self.verbose
+    }
+
     pub fn discover_subdomain(&self, subdomain: String) {
         self.discovered_subdomains.lock().unwrap().push(subdomain)
     }