From f9958074f06974ab341285312ac3612eba727e35 Mon Sep 17 00:00:00 2001 From: Krypton Date: Sat, 2 Nov 2024 21:17:38 +0100 Subject: [PATCH] feat: Add `Context` enum and `semaphore` field (#5) --- Cargo.toml | 5 +++-- src/events.rs | 2 +- src/modules/enumerate_files/mod.rs | 12 +++++++++--- src/modules/enumerate_subdomains/mod.rs | 12 +++++++++--- src/modules/mod.rs | 14 +++++++++++++- src/modules/passive_dns/mod.rs | 13 +++++++++---- src/modules/ready/mod.rs | 4 ++-- src/session.rs | 18 +++++++++++------- src/state.rs | 16 ++++++++++++---- 9 files changed, 69 insertions(+), 27 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 21edf81..c8a03bc 100644 --- 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 index fb013ce..15ebc40 100644 --- 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 index 6139716..864cb30 100644 --- 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 index 43a7c9e..64711c7 100644 --- 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 index 8c9c104..b319f9b 100644 --- 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; - 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 index 5719adc..90e4d29 100644 --- 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 index bfe0fcf..af2628b 100644 --- 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 index f8a48df..8584227 100644 --- a/src/session.rs +++ b/src/session.rs @@ -76,7 +76,6 @@ impl Session { pub fn start(self: Arc) { 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 index 6f794d6..34e3681 100644 --- 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, + semaphore: Arc, verbose: bool, discovered_subdomains: Mutex>, @@ -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) }