From e967136d675dfb7b833c885bc410dabdc8078d27 Mon Sep 17 00:00:00 2001 From: Simone Margaritelli Date: Wed, 21 Aug 2024 12:14:12 +0200 Subject: [PATCH] =?UTF-8?q?chore:=20refactored=20to=20register=20plugins?= =?UTF-8?q?=20with=20a=20macro=C2=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/plugins/amqp/mod.rs | 6 +----- src/plugins/cmd/mod.rs | 6 +----- src/plugins/dns/mod.rs | 5 +---- src/plugins/ftp/mod.rs | 6 +----- src/plugins/http/mod.rs | 25 ++++++++++++++++--------- src/plugins/imap/mod.rs | 6 +----- src/plugins/kerberos/mod.rs | 6 +----- src/plugins/ldap/mod.rs | 6 +----- src/plugins/manager.rs | 12 ++++++++++++ src/plugins/mongodb/mod.rs | 6 +----- src/plugins/mqtt/mod.rs | 5 +---- src/plugins/mssql/mod.rs | 6 +----- src/plugins/oracle/mod.rs | 6 +----- src/plugins/pop3/mod.rs | 6 +----- src/plugins/port_scanner/mod.rs | 5 +---- src/plugins/rdp/mod.rs | 6 +----- src/plugins/redis/mod.rs | 5 +---- src/plugins/samba/mod.rs | 6 +----- src/plugins/scylla/mod.rs | 6 +----- src/plugins/smtp/mod.rs | 6 +----- src/plugins/socks5/mod.rs | 6 +----- src/plugins/sql/mod.rs | 12 ++++++------ src/plugins/ssh/mod.rs | 8 +------- src/plugins/stomp/mod.rs | 6 +----- src/plugins/telnet/mod.rs | 6 +----- src/plugins/vnc/mod.rs | 5 +---- 26 files changed, 57 insertions(+), 127 deletions(-) diff --git a/src/plugins/amqp/mod.rs b/src/plugins/amqp/mod.rs index a8f89ea..0a050ff 100644 --- 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 index 4588708..02b7ada 100644 --- 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 index 7653b62..6c04be8 100644 --- 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 index afc0691..90f260d 100644 --- 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 index 2408951..bc68ff7 100644 --- 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 index c261b18..11b7bf7 100644 --- 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 index d7774d9..4517942 100644 --- 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 index 89eb3c4..6593aba 100644 --- 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 index 6d0e13a..3210c79 100644 --- a/src/plugins/manager.rs +++ b/src/plugins/manager.rs @@ -16,6 +16,18 @@ use super::plugin::PayloadStrategy; type Inventory = BTreeMap<&'static str, Box>; +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(&mut self, name: &'static str, plugin: P); } diff --git a/src/plugins/mongodb/mod.rs b/src/plugins/mongodb/mod.rs index d722709..62fdfd5 100644 --- 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 index fdb7c10..0cc69fe 100644 --- 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 index c06929b..613062a 100644 --- 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 index 6d681c9..36e999a 100644 --- 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 index cc71c2e..c992647 100644 --- 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 index 74dc133..4d91825 100644 --- 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 index 2b28495..dfa1983 100644 --- 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 index aa18d1d..3116547 100644 --- 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 index 297e2bb..71ed2cc 100644 --- a/src/plugins/samba/mod.rs +++ b/src/plugins/samba/mod.rs @@ -12,8 +12,6 @@ use crate::{utils, Options}; use lazy_static::lazy_static; -use super::manager::PluginRegistrar; - pub(crate) mod options; lazy_static! { @@ -21,9 +19,7 @@ lazy_static! { static ref PAVAO_LOCK: Mutex = 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 index f953d51..42d6488 100644 --- 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 index e3b2c12..3244de1 100644 --- 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 index 2252966..1c2d875 100644 --- 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 index 4e3e8e8..a04afa4 100644 --- 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 = { 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 index 7a4fcbf..c7bd2e1 100644 --- 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 index 04126e8..b3c87c7 100644 --- 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 {} diff --git a/src/plugins/telnet/mod.rs b/src/plugins/telnet/mod.rs index 5f75569..2ec3362 100644 --- a/src/plugins/telnet/mod.rs +++ b/src/plugins/telnet/mod.rs @@ -8,13 +8,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) { - registrar.register("telnet", Telnet::new()); -} +super::manager::register_plugin!("telnet", Telnet::new()); #[derive(Clone)] pub(crate) struct Telnet { diff --git a/src/plugins/vnc/mod.rs b/src/plugins/vnc/mod.rs index b07c75e..61b28ac 100644 --- a/src/plugins/vnc/mod.rs +++ b/src/plugins/vnc/mod.rs @@ -9,12 +9,9 @@ use crate::{utils, Options}; use crate::creds::Credentials; -use super::manager::PluginRegistrar; use super::plugin::PayloadStrategy; -pub(super) fn register(registrar: &mut impl PluginRegistrar) { - registrar.register("vnc", VNC::new()); -} +super::manager::register_plugin!("vnc", VNC::new()); #[derive(Clone)] pub(crate) struct VNC {}