From 00840cd19a170795bc02ea6cb3f047ba1243e8be Mon Sep 17 00:00:00 2001 From: Erdogan Yoksul Date: Sun, 22 Sep 2024 17:22:34 +0300 Subject: [PATCH] refactor: generic api integration --- src/modules/generics/api_integration.rs | 49 +++++++++++-------- src/modules/integrations/alienvault.rs | 6 +-- src/modules/integrations/anubis.rs | 6 +-- src/modules/integrations/bevigil.rs | 6 +-- tests/modules/common.rs | 6 +-- tests/modules/integrations/alienvault_test.rs | 2 +- tests/modules/integrations/anubis_test.rs | 2 +- tests/modules/integrations/bevigil_test.rs | 2 +- 8 files changed, 44 insertions(+), 35 deletions(-) diff --git a/src/modules/generics/api_integration.rs b/src/modules/generics/api_integration.rs index 9321d5e5..4cf065c1 100644 --- a/src/modules/generics/api_integration.rs +++ b/src/modules/generics/api_integration.rs @@ -6,10 +6,8 @@ use crate::{ }, }; use async_trait::async_trait; -use reqwest::{ - header::{HeaderName, HeaderValue}, - Url, -}; +use reqwest::header::{HeaderName, HeaderValue}; +use reqwest::Url; use std::{collections::BTreeSet, str::FromStr}; use tokio::sync::Mutex; @@ -27,7 +25,7 @@ pub struct GenericAPIIntegrationModule { pub name: String, /// Simple function field that gets query URL /// by given domain address - pub url: Box String + Sync + Send>, + pub url: Box String + Sync + Send>, /// Set authentication method, see [`AuthMethod`] enum /// for details pub auth: AuthMethod, @@ -37,6 +35,30 @@ pub struct GenericAPIIntegrationModule { pub extractor: SubdomainExtractorDispatcher, } +impl GenericAPIIntegrationModule { + async fn authenticate(&self, domain: &str) -> Url { + let url: Url = (self.url)(domain).parse().unwrap(); + let apikey = self.fetch_apikey().await; + + match &self.auth { + AuthMethod::APIKeyInHeader(key) => { + if let Ok(apikey) = apikey { + let mut requester = self.requester.lock().await; + + let (name, value) = (HeaderName::from_str(key), HeaderValue::from_str(&apikey)); + + if let (Ok(name), Ok(value)) = (name, value) { + requester.config().await.add_header(name, value); + } + } + } + AuthMethod::APIKeyInURL | AuthMethod::NoAuth => {} + } + + url + } +} + #[async_trait(?Send)] impl SubscanModuleInterface for GenericAPIIntegrationModule { async fn name(&self) -> &str { @@ -52,22 +74,9 @@ impl SubscanModuleInterface for GenericAPIIntegrationModule { } async fn run(&mut self, domain: String) -> BTreeSet { - let mut requester = self.requester.lock().await; - let url = Url::parse(&(self.url)(domain.clone())).unwrap(); - - match &self.auth { - AuthMethod::APIKeyInHeader(key) => { - if let Ok(apikey) = self.fetch_apikey().await { - let name = HeaderName::from_str(key.as_str()).unwrap(); - let value = HeaderValue::from_str(apikey.as_str()).unwrap(); - - requester.config().await.add_header(name, value); - } - } - AuthMethod::APIKeyInURL => {} - AuthMethod::NoAuth => {} - } + let url = self.authenticate(&domain).await; + let requester = self.requester.lock().await; let content = requester.get_content(url).await.unwrap_or_default(); self.extractor.extract(content, domain).await diff --git a/src/modules/integrations/alienvault.rs b/src/modules/integrations/alienvault.rs index 6f4e784e..4e288e24 100644 --- a/src/modules/integrations/alienvault.rs +++ b/src/modules/integrations/alienvault.rs @@ -56,13 +56,13 @@ impl AlienVault { /// /// #[tokio::main] /// async fn main() { - /// let domain = "foo.com".to_string(); - /// let url = alienvault::AlienVault::get_query_url(domain.clone()); + /// let domain = "foo.com"; + /// let url = alienvault::AlienVault::get_query_url(&domain); /// /// assert_eq!(url, format!("{ALIENVAULT_URL}/{domain}/passive_dns")); /// } /// ``` - pub fn get_query_url(domain: String) -> String { + pub fn get_query_url(domain: &str) -> String { format!("{ALIENVAULT_URL}/{domain}/passive_dns") } diff --git a/src/modules/integrations/anubis.rs b/src/modules/integrations/anubis.rs index 3f10f77e..63849eaa 100644 --- a/src/modules/integrations/anubis.rs +++ b/src/modules/integrations/anubis.rs @@ -55,13 +55,13 @@ impl Anubis { /// /// #[tokio::main] /// async fn main() { - /// let domain = "foo.com".to_string(); - /// let url = anubis::Anubis::get_query_url(domain.clone()); + /// let domain = "foo.com"; + /// let url = anubis::Anubis::get_query_url(&domain); /// /// assert_eq!(url, format!("{ANUBIS_URL}/{domain}")); /// } /// ``` - pub fn get_query_url(domain: String) -> String { + pub fn get_query_url(domain: &str) -> String { format!("{ANUBIS_URL}/{domain}") } diff --git a/src/modules/integrations/bevigil.rs b/src/modules/integrations/bevigil.rs index f511e2eb..8e167af8 100644 --- a/src/modules/integrations/bevigil.rs +++ b/src/modules/integrations/bevigil.rs @@ -55,13 +55,13 @@ impl Bevigil { /// /// #[tokio::main] /// async fn main() { - /// let domain = "foo.com".to_string(); - /// let url = bevigil::Bevigil::get_query_url(domain.clone()); + /// let domain = "foo.com"; + /// let url = bevigil::Bevigil::get_query_url(&domain); /// /// assert_eq!(url, format!("{BEVIGIL_URL}/{domain}/subdomains")); /// } /// ``` - pub fn get_query_url(domain: String) -> String { + pub fn get_query_url(domain: &str) -> String { format!("{BEVIGIL_URL}/{domain}/subdomains") } diff --git a/tests/modules/common.rs b/tests/modules/common.rs index adc9fe58..a3a9c5d1 100644 --- a/tests/modules/common.rs +++ b/tests/modules/common.rs @@ -59,9 +59,9 @@ pub mod mocks { } } - pub fn wrap_url_with_mock_func(url: &str) -> Box String + Sync + Send> { - let url = Url::parse(url).unwrap().to_string(); + pub fn wrap_url_with_mock_func(url: &str) -> Box String + Sync + Send> { + let url: Url = url.parse().unwrap(); - Box::new(move |_| url.clone()) + Box::new(move |_| url.to_string().clone()) } } diff --git a/tests/modules/integrations/alienvault_test.rs b/tests/modules/integrations/alienvault_test.rs index 64584c39..2155d5e3 100644 --- a/tests/modules/integrations/alienvault_test.rs +++ b/tests/modules/integrations/alienvault_test.rs @@ -32,7 +32,7 @@ async fn alienvault_run_test() { #[tokio::test] async fn get_query_url_test() { - let url = alienvault::AlienVault::get_query_url(TEST_DOMAIN.to_string()); + let url = alienvault::AlienVault::get_query_url(TEST_DOMAIN); let expected = format!("{ALIENVAULT_URL}/{TEST_DOMAIN}/passive_dns"); assert_eq!(url, expected); diff --git a/tests/modules/integrations/anubis_test.rs b/tests/modules/integrations/anubis_test.rs index 04ae4452..6dac28ac 100644 --- a/tests/modules/integrations/anubis_test.rs +++ b/tests/modules/integrations/anubis_test.rs @@ -25,7 +25,7 @@ async fn anubis_run_test() { #[tokio::test] async fn get_query_url_test() { - let url = anubis::Anubis::get_query_url(TEST_DOMAIN.to_string()); + let url = anubis::Anubis::get_query_url(TEST_DOMAIN); let expected = format!("{ANUBIS_URL}/{TEST_DOMAIN}"); assert_eq!(url, expected); diff --git a/tests/modules/integrations/bevigil_test.rs b/tests/modules/integrations/bevigil_test.rs index 1736fbe5..6c784a3a 100644 --- a/tests/modules/integrations/bevigil_test.rs +++ b/tests/modules/integrations/bevigil_test.rs @@ -33,7 +33,7 @@ async fn bevigil_run_test() { #[tokio::test] async fn get_query_url_test() { - let url = bevigil::Bevigil::get_query_url(TEST_DOMAIN.to_string()); + let url = bevigil::Bevigil::get_query_url(TEST_DOMAIN); let expected = format!("{BEVIGIL_URL}/{TEST_DOMAIN}/subdomains"); assert_eq!(url, expected);