chore: refactored to register plugins with a macro£

evilsocket · Aug 21, 2024 · e967136 · e967136
1 parent 6c5baed
commit e967136
Show file tree

Hide file tree

Showing 26 changed files with 57 additions and 127 deletions.
diff --git a/src/plugins/amqp/mod.rs b/src/plugins/amqp/mod.rs
@@ -10,15 +10,11 @@ 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];
 
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("amqp", AMQP::new());
-}
+super::manager::register_plugin!("amqp", AMQP::new());
 
 #[derive(Clone)]
 pub(crate) struct AMQP {

diff --git a/src/plugins/cmd/mod.rs b/src/plugins/cmd/mod.rs
@@ -9,13 +9,9 @@ use crate::Plugin;
 
 use crate::creds::Credentials;
 
-use super::manager::PluginRegistrar;
-
 pub(crate) mod options;
 
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("cmd", Command::new());
-}
+super::manager::register_plugin!("cmd", Command::new());
 
 #[derive(Clone)]
 pub(crate) struct Command {

diff --git a/src/plugins/dns/mod.rs b/src/plugins/dns/mod.rs
@@ -15,14 +15,11 @@ use crate::Plugin;
 
 use crate::creds::Credentials;
 
-use super::manager::PluginRegistrar;
 use super::plugin::PayloadStrategy;
 
 pub(crate) mod options;
 
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("dns", DNS::new());
-}
+super::manager::register_plugin!("dns", DNS::new());
 
 #[derive(Clone)]
 pub(crate) struct DNS {

diff --git a/src/plugins/ftp/mod.rs b/src/plugins/ftp/mod.rs
@@ -11,11 +11,7 @@ use crate::Plugin;
 
 use crate::creds::Credentials;
 
-use super::manager::PluginRegistrar;
-
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("ftp", FTP::new());
-}
+super::manager::register_plugin!("ftp", FTP::new());
 
 #[derive(Clone)]
 pub(crate) struct FTP {}

diff --git a/src/plugins/http/mod.rs b/src/plugins/http/mod.rs
@@ -28,15 +28,22 @@ const HTTP_USERNAME_VAR: &str = "{$username}";
 const HTTP_PASSWORD_VAR: &str = "{$password}";
 const HTTP_PAYLOAD_VAR: &str = "{$payload}";
 
-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));
-}
+super::manager::register_plugin!(
+    "http",
+    HTTP::new(Strategy::Request),
+    "http.form",
+    HTTP::new(Strategy::Form),
+    "http.basic",
+    HTTP::new(Strategy::BasicAuth),
+    "http.ntlm1",
+    HTTP::new(Strategy::NLTMv1),
+    "http.ntlm2",
+    HTTP::new(Strategy::NLTMv2),
+    "http.enum",
+    HTTP::new(Strategy::Enumeration),
+    "http.vhost",
+    HTTP::new(Strategy::VHostEnum)
+);
 
 fn method_requires_payload(method: &Method) -> bool {
     matches!(method, &Method::POST | &Method::PUT | &Method::PATCH)

diff --git a/src/plugins/imap/mod.rs b/src/plugins/imap/mod.rs
@@ -9,11 +9,7 @@ use crate::Plugin;
 use crate::creds::Credentials;
 use crate::utils;
 
-use super::manager::PluginRegistrar;
-
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("imap", IMAP::new());
-}
+super::manager::register_plugin!("imap", IMAP::new());
 
 #[derive(Clone)]
 pub(crate) struct IMAP {}

diff --git a/src/plugins/kerberos/mod.rs b/src/plugins/kerberos/mod.rs
@@ -15,16 +15,12 @@ use crate::creds::Credentials;
 use crate::utils;
 use transport::Protocol;
 
-use super::manager::PluginRegistrar;
-
 mod builder;
 mod transport;
 
 pub(crate) mod options;
 
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("kerberos", Kerberos::new());
-}
+super::manager::register_plugin!("kerberos", Kerberos::new());
 
 #[derive(Clone)]
 pub(crate) struct Kerberos {

diff --git a/src/plugins/ldap/mod.rs b/src/plugins/ldap/mod.rs
@@ -10,13 +10,9 @@ use crate::Plugin;
 use crate::creds::Credentials;
 use crate::utils;
 
-use super::manager::PluginRegistrar;
-
 pub(crate) mod options;
 
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("ldap", LDAP::new());
-}
+super::manager::register_plugin!("ldap", LDAP::new());
 
 #[derive(Clone)]
 pub(crate) struct LDAP {

diff --git a/src/plugins/manager.rs b/src/plugins/manager.rs
@@ -16,6 +16,18 @@ use super::plugin::PayloadStrategy;
 
 type Inventory = BTreeMap<&'static str, Box<dyn Plugin>>;
 
+macro_rules! register_plugin {
+    ($($name:literal, $instance:expr),+) => {
+        pub(super) fn register(registrar: &mut impl $crate::plugins::manager::PluginRegistrar) {
+            $(
+                registrar.register($name, $instance);
+            )*
+        }
+    };
+}
+
+pub(crate) use register_plugin;
+
 pub(crate) trait PluginRegistrar {
     fn register<P: Plugin + 'static>(&mut self, name: &'static str, plugin: P);
 }

diff --git a/src/plugins/mongodb/mod.rs b/src/plugins/mongodb/mod.rs
@@ -9,11 +9,7 @@ use crate::{utils, Options};
 
 use crate::creds::Credentials;
 
-use super::manager::PluginRegistrar;
-
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("mongodb", MongoDB::new());
-}
+super::manager::register_plugin!("mongodb", MongoDB::new());
 
 #[derive(Clone)]
 pub(crate) struct MongoDB {}

diff --git a/src/plugins/mqtt/mod.rs b/src/plugins/mqtt/mod.rs
@@ -10,12 +10,9 @@ use crate::Plugin;
 
 use crate::creds::Credentials;
 
-use super::manager::PluginRegistrar;
 pub(crate) mod options;
 
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("mqtt", Mqtt::new());
-}
+super::manager::register_plugin!("mqtt", Mqtt::new());
 
 #[derive(Clone)]
 pub(crate) struct Mqtt {

diff --git a/src/plugins/mssql/mod.rs b/src/plugins/mssql/mod.rs
@@ -10,8 +10,6 @@ use crate::Plugin;
 use crate::creds::Credentials;
 use crate::utils;
 
-use super::manager::PluginRegistrar;
-
 // ripped from medusa mssql.c
 const MS_PACKET_HEADER: &[u8] = &[
     0x02, 0x00, 0x02, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
@@ -59,9 +57,7 @@ const MS_PACKET_LANGP: &[u8] = &[
 
 const MS_MAX_LEN: usize = 30;
 
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("mssql", MSSQL::new());
-}
+super::manager::register_plugin!("mssql", MSSQL::new());
 
 #[derive(Clone)]
 pub(crate) struct MSSQL {}

diff --git a/src/plugins/oracle/mod.rs b/src/plugins/oracle/mod.rs
@@ -11,11 +11,7 @@ use crate::Plugin;
 
 pub(crate) mod options;
 
-use super::manager::PluginRegistrar;
-
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("oracle", Oracle::new());
-}
+super::manager::register_plugin!("oracle", Oracle::new());
 
 #[derive(Clone)]
 pub(crate) struct Oracle {

diff --git a/src/plugins/pop3/mod.rs b/src/plugins/pop3/mod.rs
@@ -9,13 +9,9 @@ use crate::Plugin;
 use crate::creds::Credentials;
 use crate::utils;
 
-use super::manager::PluginRegistrar;
-
 pub(crate) mod options;
 
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("pop3", POP3::new());
-}
+super::manager::register_plugin!("pop3", POP3::new());
 
 #[derive(Clone)]
 pub(crate) struct POP3 {

diff --git a/src/plugins/port_scanner/mod.rs b/src/plugins/port_scanner/mod.rs
@@ -12,15 +12,12 @@ use crate::{creds, utils};
 
 use crate::creds::{Credentials, Expression};
 
-use super::manager::PluginRegistrar;
 use super::plugin::PayloadStrategy;
 
 mod grabbers;
 pub(crate) mod options;
 
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("port.scanner", PortScanner::new());
-}
+super::manager::register_plugin!("port.scanner", PortScanner::new());
 
 #[derive(Clone)]
 pub(crate) struct PortScanner {

diff --git a/src/plugins/rdp/mod.rs b/src/plugins/rdp/mod.rs
@@ -12,13 +12,9 @@ use crate::{utils, Options};
 
 use crate::creds::Credentials;
 
-use super::manager::PluginRegistrar;
-
 pub(crate) mod options;
 
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("rdp", RDP::new());
-}
+super::manager::register_plugin!("rdp", RDP::new());
 
 #[derive(Clone)]
 pub(crate) struct RDP {

diff --git a/src/plugins/redis/mod.rs b/src/plugins/redis/mod.rs
@@ -9,12 +9,9 @@ use crate::{utils, Options};
 
 use crate::creds::Credentials;
 
-use super::manager::PluginRegistrar;
 pub(crate) mod options;
 
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("redis", Redis::new());
-}
+super::manager::register_plugin!("redis", Redis::new());
 
 #[derive(Clone)]
 pub(crate) struct Redis {

diff --git a/src/plugins/samba/mod.rs b/src/plugins/samba/mod.rs
@@ -12,18 +12,14 @@ use crate::{utils, Options};
 
 use lazy_static::lazy_static;
 
-use super::manager::PluginRegistrar;
-
 pub(crate) mod options;
 
 lazy_static! {
     static ref SHARE_CACHE: Mutex<HashMap<String, String>> = Mutex::new(HashMap::new());
     static ref PAVAO_LOCK: Mutex<bool> = tokio::sync::Mutex::new(true);
 }
 
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("smb", SMB::new());
-}
+super::manager::register_plugin!("smb", SMB::new());
 
 #[derive(Clone)]
 pub(crate) struct SMB {

diff --git a/src/plugins/scylla/mod.rs b/src/plugins/scylla/mod.rs
@@ -9,11 +9,7 @@ use crate::Plugin;
 
 use crate::creds::Credentials;
 
-use super::manager::PluginRegistrar;
-
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("scylla", Scylla::new());
-}
+super::manager::register_plugin!("scylla", Scylla::new());
 
 #[derive(Clone)]
 pub(crate) struct Scylla {}

diff --git a/src/plugins/smtp/mod.rs b/src/plugins/smtp/mod.rs
@@ -11,13 +11,9 @@ use crate::Plugin;
 use crate::creds::Credentials;
 use crate::utils;
 
-use super::manager::PluginRegistrar;
-
 pub(crate) mod options;
 
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("smtp", SMTP::new());
-}
+super::manager::register_plugin!("smtp", SMTP::new());
 
 #[derive(Clone)]
 pub(crate) struct SMTP {

diff --git a/src/plugins/socks5/mod.rs b/src/plugins/socks5/mod.rs
@@ -9,13 +9,9 @@ use crate::Plugin;
 
 use crate::creds::Credentials;
 
-use super::manager::PluginRegistrar;
-
 pub(crate) mod options;
 
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("socks5", Socks5::new());
-}
+super::manager::register_plugin!("socks5", Socks5::new());
 
 #[derive(Clone)]
 pub(crate) struct Socks5 {

diff --git a/src/plugins/sql/mod.rs b/src/plugins/sql/mod.rs
@@ -12,8 +12,6 @@ use crate::utils;
 use crate::Options;
 use crate::Plugin;
 
-use super::manager::PluginRegistrar;
-
 lazy_static! {
     static ref DESCRIPTIONS: HashMap<Flavour, &'static str> = {
         HashMap::from([
@@ -25,10 +23,12 @@ lazy_static! {
         HashMap::from([(Flavour::My, 3306), (Flavour::PG, 5432),]);
 }
 
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("mysql", SQL::new(Flavour::My));
-    registrar.register("pgsql", SQL::new(Flavour::PG));
-}
+super::manager::register_plugin!(
+    "mysql",
+    SQL::new(Flavour::My),
+    "pgsql",
+    SQL::new(Flavour::PG)
+);
 
 #[derive(Clone, PartialEq, Eq, Debug, Hash)]
 pub(crate) enum Flavour {

diff --git a/src/plugins/ssh/mod.rs b/src/plugins/ssh/mod.rs
@@ -10,15 +10,9 @@ use crate::utils;
 use crate::Options;
 use crate::Plugin;
 
-use super::manager::PluginRegistrar;
-
 pub(crate) mod options;
 
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    let ssh = SSH::new();
-    registrar.register("ssh", ssh.clone());
-    registrar.register("sftp", ssh);
-}
+super::manager::register_plugin!("ssh", SSH::new(), "sftp", SSH::new());
 
 #[derive(Clone)]
 pub(crate) struct SSH {

diff --git a/src/plugins/stomp/mod.rs b/src/plugins/stomp/mod.rs
@@ -10,13 +10,9 @@ use crate::Plugin;
 
 use crate::creds::Credentials;
 
-use super::manager::PluginRegistrar;
-
 const CONNECTED_RESPONSE: &[u8] = &[67, 79, 78, 78, 69, 67, 84, 69, 68];
 
-pub(super) fn register(registrar: &mut impl PluginRegistrar) {
-    registrar.register("stomp", STOMP::new());
-}
+super::manager::register_plugin!("stomp", STOMP::new());
 
 #[derive(Clone)]
 pub(crate) struct STOMP {}