Merge pull request #52 from shurizzle/register

refactor: remove ctors
evilsocket · Aug 21, 2024 · 6c5baed · 6c5baed
2 parents d41fd51 + 78f206c
commit 6c5baed
Show file tree

Hide file tree

Showing 29 changed files with 193 additions and 184 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -22,7 +22,6 @@ async-trait = "0.1.73"
 chrono = { version = "0.4.31", features = ["serde"] }
 clap = { version = "4.4.4", features = ["derive"] }
 clap_complete = "4.4.6"
-ctor = "0.2.4"
 ctrlc = "3.4.1"
 indexmap = { version = "2.0.1", features = ["serde"] }
 lazy_static = "1.4.0"

diff --git a/src/plugins/amqp/mod.rs b/src/plugins/amqp/mod.rs
@@ -1,7 +1,6 @@
 use std::time::Duration;
 
 use async_trait::async_trait;
-use ctor::ctor;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 
 use crate::session::{Error, Loot};
@@ -11,13 +10,14 @@ use crate::Plugin;
 
 use crate::creds::Credentials;
 
+use super::manager::PluginRegistrar;
+
 pub(crate) mod options;
 
 const PROTOCOL_HEADER_091: &[u8] = &[b'A', b'M', b'Q', b'P', 0, 0, 9, 1];
 
-#[ctor]
-fn register() {
-    crate::plugins::manager::register("amqp", Box::new(AMQP::new()));
+pub(super) fn register(registrar: &mut impl PluginRegistrar) {
+    registrar.register("amqp", AMQP::new());
 }
 
 #[derive(Clone)]

diff --git a/src/plugins/cmd/mod.rs b/src/plugins/cmd/mod.rs
@@ -2,19 +2,19 @@ use std::process::Stdio;
 use std::time::Duration;
 
 use async_trait::async_trait;
-use ctor::ctor;
 
 use crate::session::{Error, Loot};
 use crate::Options;
 use crate::Plugin;
 
 use crate::creds::Credentials;
 
+use super::manager::PluginRegistrar;
+
 pub(crate) mod options;
 
-#[ctor]
-fn register() {
-    crate::plugins::manager::register("cmd", Box::new(Command::new()));
+pub(super) fn register(registrar: &mut impl PluginRegistrar) {
+    registrar.register("cmd", Command::new());
 }
 
 #[derive(Clone)]

diff --git a/src/plugins/dns/mod.rs b/src/plugins/dns/mod.rs
@@ -4,7 +4,6 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use async_trait::async_trait;
-use ctor::ctor;
 use tokio::sync::Mutex;
 use trust_dns_resolver::{config::*, AsyncResolver, TokioAsyncResolver};
 use x509_parser::prelude::{FromDer, GeneralName, X509Certificate};
@@ -16,13 +15,13 @@ use crate::Plugin;
 
 use crate::creds::Credentials;
 
+use super::manager::PluginRegistrar;
 use super::plugin::PayloadStrategy;
 
 pub(crate) mod options;
 
-#[ctor]
-fn register() {
-    crate::plugins::manager::register("dns", Box::new(DNS::new()));
+pub(super) fn register(registrar: &mut impl PluginRegistrar) {
+    registrar.register("dns", DNS::new());
 }
 
 #[derive(Clone)]

diff --git a/src/plugins/ftp/mod.rs b/src/plugins/ftp/mod.rs
@@ -3,7 +3,6 @@ use async_ftp::FtpStream;
 use std::time::Duration;
 
 use async_trait::async_trait;
-use ctor::ctor;
 
 use crate::session::{Error, Loot};
 use crate::utils;
@@ -12,9 +11,10 @@ use crate::Plugin;
 
 use crate::creds::Credentials;
 
-#[ctor]
-fn register() {
-    crate::plugins::manager::register("ftp", Box::new(FTP::new()));
+use super::manager::PluginRegistrar;
+
+pub(super) fn register(registrar: &mut impl PluginRegistrar) {
+    registrar.register("ftp", FTP::new());
 }
 
 #[derive(Clone)]

diff --git a/src/plugins/http/mod.rs b/src/plugins/http/mod.rs
@@ -1,7 +1,6 @@
 use std::time::Duration;
 
 use async_trait::async_trait;
-use ctor::ctor;
 use rand::seq::SliceRandom;
 use reqwest::{
     header::{HeaderMap, HeaderName, HeaderValue, CONTENT_TYPE, COOKIE, HOST, USER_AGENT},
@@ -15,7 +14,7 @@ use crate::Options;
 use crate::creds::Credentials;
 use crate::plugins::Plugin;
 
-use super::plugin::PayloadStrategy;
+use super::{manager::PluginRegistrar, plugin::PayloadStrategy};
 
 mod csrf;
 mod ntlm;
@@ -29,15 +28,14 @@ const HTTP_USERNAME_VAR: &str = "{$username}";
 const HTTP_PASSWORD_VAR: &str = "{$password}";
 const HTTP_PAYLOAD_VAR: &str = "{$payload}";
 
-#[ctor]
-fn register() {
-    crate::plugins::manager::register("http", Box::new(HTTP::new(Strategy::Request)));
-    crate::plugins::manager::register("http.form", Box::new(HTTP::new(Strategy::Form)));
-    crate::plugins::manager::register("http.basic", Box::new(HTTP::new(Strategy::BasicAuth)));
-    crate::plugins::manager::register("http.ntlm1", Box::new(HTTP::new(Strategy::NLTMv1)));
-    crate::plugins::manager::register("http.ntlm2", Box::new(HTTP::new(Strategy::NLTMv2)));
-    crate::plugins::manager::register("http.enum", Box::new(HTTP::new(Strategy::Enumeration)));
-    crate::plugins::manager::register("http.vhost", Box::new(HTTP::new(Strategy::VHostEnum)));
+pub(super) fn register(registrar: &mut impl PluginRegistrar) {
+    registrar.register("http", HTTP::new(Strategy::Request));
+    registrar.register("http.form", HTTP::new(Strategy::Form));
+    registrar.register("http.basic", HTTP::new(Strategy::BasicAuth));
+    registrar.register("http.ntlm1", HTTP::new(Strategy::NLTMv1));
+    registrar.register("http.ntlm2", HTTP::new(Strategy::NLTMv2));
+    registrar.register("http.enum", HTTP::new(Strategy::Enumeration));
+    registrar.register("http.vhost", HTTP::new(Strategy::VHostEnum));
 }
 
 fn method_requires_payload(method: &Method) -> bool {

diff --git a/src/plugins/imap/mod.rs b/src/plugins/imap/mod.rs
@@ -1,7 +1,6 @@
 use std::time::Duration;
 
 use async_trait::async_trait;
-use ctor::ctor;
 
 use crate::session::{Error, Loot};
 use crate::Options;
@@ -10,9 +9,10 @@ use crate::Plugin;
 use crate::creds::Credentials;
 use crate::utils;
 
-#[ctor]
-fn register() {
-    crate::plugins::manager::register("imap", Box::new(IMAP::new()));
+use super::manager::PluginRegistrar;
+
+pub(super) fn register(registrar: &mut impl PluginRegistrar) {
+    registrar.register("imap", IMAP::new());
 }
 
 #[derive(Clone)]

diff --git a/src/plugins/kerberos/mod.rs b/src/plugins/kerberos/mod.rs
@@ -4,7 +4,6 @@ use std::time::Duration;
 
 use ahash::HashSet;
 use async_trait::async_trait;
-use ctor::ctor;
 use kerberos_asn1::{AsRep, Asn1Object, KrbError};
 use kerberos_constants::error_codes;
 
@@ -16,14 +15,15 @@ use crate::creds::Credentials;
 use crate::utils;
 use transport::Protocol;
 
+use super::manager::PluginRegistrar;
+
 mod builder;
 mod transport;
 
 pub(crate) mod options;
 
-#[ctor]
-fn register() {
-    crate::plugins::manager::register("kerberos", Box::new(Kerberos::new()));
+pub(super) fn register(registrar: &mut impl PluginRegistrar) {
+    registrar.register("kerberos", Kerberos::new());
 }
 
 #[derive(Clone)]

diff --git a/src/plugins/ldap/mod.rs b/src/plugins/ldap/mod.rs
@@ -1,7 +1,6 @@
 use std::time::Duration;
 
 use async_trait::async_trait;
-use ctor::ctor;
 use ldap3::{LdapConnAsync, LdapConnSettings};
 
 use crate::session::{Error, Loot};
@@ -11,11 +10,12 @@ use crate::Plugin;
 use crate::creds::Credentials;
 use crate::utils;
 
+use super::manager::PluginRegistrar;
+
 pub(crate) mod options;
 
-#[ctor]
-fn register() {
-    crate::plugins::manager::register("ldap", Box::new(LDAP::new()));
+pub(super) fn register(registrar: &mut impl PluginRegistrar) {
+    registrar.register("ldap", LDAP::new());
 }
 
 #[derive(Clone)]

diff --git a/src/plugins/manager.rs b/src/plugins/manager.rs
@@ -1,10 +1,9 @@
 use std::collections::BTreeMap;
-use std::sync::{Mutex, RwLock};
+use std::sync::{LazyLock, Mutex, RwLock};
 use std::time;
 
 use ahash::HashSet;
 use ansi_term::Style;
-use lazy_static::lazy_static;
 use rand::Rng;
 use std::sync::Arc;
 use tokio::task;
@@ -17,14 +16,23 @@ use super::plugin::PayloadStrategy;
 
 type Inventory = BTreeMap<&'static str, Box<dyn Plugin>>;
 
-lazy_static! {
-    pub(crate) static ref INVENTORY: Mutex<Inventory> = Mutex::new(Inventory::new());
+pub(crate) trait PluginRegistrar {
+    fn register<P: Plugin + 'static>(&mut self, name: &'static str, plugin: P);
 }
 
-pub(crate) fn register(name: &'static str, plugin: Box<dyn Plugin>) {
-    INVENTORY.lock().unwrap().insert(name, plugin);
+impl PluginRegistrar for Inventory {
+    #[inline]
+    fn register<P: Plugin + 'static>(&mut self, name: &'static str, plugin: P) {
+        self.insert(name, Box::new(plugin));
+    }
 }
 
+pub(crate) static INVENTORY: LazyLock<Mutex<Inventory>> = LazyLock::new(|| {
+    let mut ps = Inventory::new();
+    super::add_defaults(&mut ps);
+    Mutex::new(ps)
+});
+
 pub(crate) fn list() {
     let bold = Style::new().bold();
 
@@ -49,15 +57,16 @@ pub(crate) fn list() {
 }
 
 pub(crate) fn setup(options: &Options) -> Result<&'static mut dyn Plugin, Error> {
-    let plugin_name = if let Some(value) = options.plugin.as_ref() {
-        value.to_string()
-    } else {
+    let Some(plugin_name) = options.plugin.as_ref() else {
         return Err("no plugin selected".to_owned());
     };
-
-    let plugin = match INVENTORY.lock().unwrap().remove(plugin_name.as_str()) {
-        Some(p) => Box::leak(p), // makes the plugin &'static
-        None => return Err(format!("{} is not a valid plugin name, run with --list-plugins to see the list of available plugins", plugin_name)),
+    let Some(plugin) = INVENTORY
+        .lock()
+        .unwrap()
+        .remove(plugin_name.as_str())
+        .map(Box::leak)
+    else {
+        return Err(format!("{} is not a valid plugin name, run with --list-plugins to see the list of available plugins", plugin_name));
     };
 
     plugin.setup(options)?;