Merge pull request #58 from haveibeensquatted/feat/impl-permutation-t…

…ypes Refactor permutations module
haveibeensquatted · Sep 30, 2023 · 33a56d1 · 33a56d1
2 parents 5ebc1d1 + 03d1aa1
commit 33a56d1
Show file tree

Hide file tree

Showing 6 changed files with 165 additions and 98 deletions.
diff --git a/examples/twistrs-cli/src/main.rs b/examples/twistrs-cli/src/main.rs
@@ -3,7 +3,7 @@ use colored::*;
 
 use tokio::sync::mpsc;
 use twistrs::enrich::DomainMetadata;
-use twistrs::permutate::Domain;
+use twistrs::permutate::{Domain, Permutation};
 
 use anyhow::Result;
 use std::collections::HashSet;
@@ -21,15 +21,13 @@ async fn main() -> Result<()> {
 
     let domain = Domain::new(matches.value_of("domain").unwrap()).unwrap();
 
-    let mut domain_permutations = domain.all()?.collect::<HashSet<String>>();
+    let domain_permutations = domain.all()?.collect::<HashSet<Permutation>>();
     let domain_permutation_count = domain_permutations.len();
 
-    domain_permutations.insert(String::from(domain.fqdn));
-
     let (tx, mut rx) = mpsc::channel(5000);
 
     for (i, v) in domain_permutations.into_iter().enumerate() {
-        let domain_metadata = DomainMetadata::new(v.clone());
+        let domain_metadata = DomainMetadata::new(v.domain.fqdn.clone());
         let mut tx = tx.clone();
 
         tokio::spawn(async move {

diff --git a/examples/twistrs-grpc/src/client.rs b/examples/twistrs-grpc/src/client.rs
@@ -3,22 +3,19 @@ use domain_enumeration::Fqdn;
 
 mod domain_enumeration;
 
-
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let channel = tonic::transport::Channel::from_static("http://127.0.0.1:8080")
         .connect()
         .await?;
-    
+
     let mut client = DomainEnumerationClient::new(channel);
 
     println!("[+] Starting DNS resolutions...");
 
-    let request = tonic::Request::new(
-        Fqdn {
-            fqdn:String::from("google.com")
-        },
-    );
+    let request = tonic::Request::new(Fqdn {
+        fqdn: String::from("google.com"),
+    });
 
     let mut response = client.send_dns_resolution(request).await?.into_inner();
 
@@ -28,11 +25,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     println!("[+] Starting MX Checks...");
 
-    let request = tonic::Request::new(
-        Fqdn {
-            fqdn:String::from("google.com")
-        },
-    );
+    let request = tonic::Request::new(Fqdn {
+        fqdn: String::from("google.com"),
+    });
 
     let mut response = client.send_mx_check(request).await?.into_inner();
 
@@ -41,4 +36,4 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     }
 
     Ok(())
-}
+}
diff --git a/examples/twistrs-grpc/src/server.rs b/examples/twistrs-grpc/src/server.rs
@@ -25,7 +25,7 @@ impl DomainEnumeration for DomainEnumerationService {
         let (tx, rx) = mpsc::channel(64);
 
         for permutation in Domain::new(&request.get_ref().fqdn).unwrap().all().unwrap() {
-            let domain_metadata = DomainMetadata::new(permutation.clone());
+            let domain_metadata = DomainMetadata::new(permutation.domain.fqdn.clone());
             let mut tx = tx.clone();
 
             // Spawn DNS Resolution check
@@ -34,7 +34,7 @@ impl DomainEnumeration for DomainEnumerationService {
                     if let Some(ips) = metadata.ips {
                         if tx
                             .send(Ok(DomainEnumerationResponse {
-                                fqdn: permutation.clone().to_string(),
+                                fqdn: permutation.domain.fqdn.to_string(),
                                 ips: ips.into_iter().map(|x| format!("{}", x)).collect(),
                             }))
                             .await
@@ -62,7 +62,7 @@ impl DomainEnumeration for DomainEnumerationService {
         let (tx, rx) = mpsc::channel(64);
 
         for permutation in Domain::new(&request.get_ref().fqdn).unwrap().all().unwrap() {
-            let domain_metadata = DomainMetadata::new(permutation.clone());
+            let domain_metadata = DomainMetadata::new(permutation.domain.fqdn.clone());
             let mut tx = tx.clone();
 
             // Spawn DNS Resolution check
@@ -71,7 +71,7 @@ impl DomainEnumeration for DomainEnumerationService {
                     if let Some(smtp) = metadata.smtp {
                         if tx
                             .send(Ok(MxCheckResponse {
-                                fqdn: permutation.clone().to_string(),
+                                fqdn: permutation.domain.fqdn.to_string(),
                                 is_positive: smtp.is_positive,
                                 message: smtp.message,
                             }))

diff --git a/examples/twistrs-ws/src/main.rs b/examples/twistrs-ws/src/main.rs
@@ -11,7 +11,7 @@ use warp::ws::{Message, WebSocket};
 use warp::Filter;
 
 use twistrs::enrich::DomainMetadata;
-use twistrs::permutate::Domain;
+use twistrs::permutate::{Domain, Permutation};
 
 /// Our global unique user id counter.
 static NEXT_USER_ID: AtomicUsize = AtomicUsize::new(1);
@@ -93,11 +93,10 @@ async fn user_message(my_id: usize, msg: Message, users: &Users) {
             eprintln!("initiating dns resolution checks for user: {}", my_id);
 
             let domain = Domain::new(msg).unwrap();
-            let mut domain_permutations = domain.all().unwrap().collect::<HashSet<String>>();
-            domain_permutations.insert(String::from(domain.fqdn));
+            let domain_permutations = domain.all().unwrap().collect::<HashSet<Permutation>>();
 
             for v in domain_permutations.into_iter() {
-                let domain_metadata = DomainMetadata::new(v.clone());
+                let domain_metadata = DomainMetadata::new(v.domain.fqdn.clone());
                 let tx = tx.clone();
 
                 tokio::spawn(async move {

diff --git a/twistrs/src/enrich.rs b/twistrs/src/enrich.rs
@@ -175,7 +175,7 @@ impl DomainMetadata {
         );
 
         let stream = BufStream::new(
-            TcpStream::connect(&format!("{}:25", self.fqdn.to_string()))
+            TcpStream::connect(&format!("{}:25", self.fqdn))
                 .await
                 .map_err(|e| EnrichmentError::SmtpLookupError {
                     domain: self.fqdn.clone(),