refactor: generic api integration

eredotpkfr · Sep 22, 2024 · 00840cd · 00840cd
1 parent d2a4a7c
commit 00840cd
Show file tree

Hide file tree

Showing 8 changed files with 44 additions and 35 deletions.
diff --git 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<dyn Fn(String) -> String + Sync + Send>,
+    pub url: Box<dyn Fn(&str) -> 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<String> {
-        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
@@ -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
@@ -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
@@ -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
@@ -59,9 +59,9 @@ pub mod mocks {
         }
     }
 
-    pub fn wrap_url_with_mock_func(url: &str) -> Box<dyn Fn(String) -> String + Sync + Send> {
-        let url = Url::parse(url).unwrap().to_string();
+    pub fn wrap_url_with_mock_func(url: &str) -> Box<dyn Fn(&str) -> 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
@@ -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
@@ -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
@@ -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);