feat: rag-embedding-ai-chat (#1)

Co-authored-by: Sharang Parnerkar <parnerkarsharang@gmail.com>
Reviewed-on: #1

compliance-dast/src/agents/api_fuzzer.rs

@@ -234,10 +234,7 @@ impl ApiFuzzerAgent {
             .ok()?;
 
         let headers = response.headers();
-        let acao = headers
-            .get("access-control-allow-origin")?
-            .to_str()
-            .ok()?;
+        let acao = headers.get("access-control-allow-origin")?.to_str().ok()?;
 
         if acao == "*" || acao == "https://evil.com" {
             let acac = headers
@@ -265,12 +262,9 @@ impl ApiFuzzerAgent {
                 request_body: None,
                 response_status: response.status().as_u16(),
                 response_headers: Some(
-                    [(
-                        "Access-Control-Allow-Origin".to_string(),
-                        acao.to_string(),
-                    )]
-                    .into_iter()
-                    .collect(),
+                    [("Access-Control-Allow-Origin".to_string(), acao.to_string())]
+                        .into_iter()
+                        .collect(),
                 ),
                 response_snippet: None,
                 screenshot_path: None,

compliance-dast/src/agents/auth_bypass.rs

@@ -132,7 +132,10 @@ impl DastAgent for AuthBypassAgent {
                         String::new(),
                         target_id.clone(),
                         DastVulnType::AuthBypass,
-                        format!("HTTP method tampering: {} accepted on {}", method, endpoint.url),
+                        format!(
+                            "HTTP method tampering: {} accepted on {}",
+                            method, endpoint.url
+                        ),
                         format!(
                             "Endpoint {} accepts {} requests which may bypass access controls.",
                             endpoint.url, method

compliance-dast/src/agents/ssrf.rs

@@ -20,10 +20,7 @@ impl SsrfAgent {
             ("http://[::1]", "localhost IPv6"),
             ("http://0.0.0.0", "zero address"),
             ("http://169.254.169.254/latest/meta-data/", "AWS metadata"),
-            (
-                "http://metadata.google.internal/",
-                "GCP metadata",
-            ),
+            ("http://metadata.google.internal/", "GCP metadata"),
             ("http://127.0.0.1:22", "SSH port probe"),
             ("http://127.0.0.1:3306", "MySQL port probe"),
             ("http://localhost/admin", "localhost admin"),
@@ -91,10 +88,7 @@ impl DastAgent for SsrfAgent {
                             .post(&endpoint.url)
                             .form(&[(param.name.as_str(), payload)])
                     } else {
-                        let test_url = format!(
-                            "{}?{}={}",
-                            endpoint.url, param.name, payload
-                        );
+                        let test_url = format!("{}?{}={}", endpoint.url, param.name, payload);
                         self.http.get(&test_url)
                     };
 
@@ -133,10 +127,7 @@ impl DastAgent for SsrfAgent {
                             String::new(),
                             target_id.clone(),
                             DastVulnType::Ssrf,
-                            format!(
-                                "SSRF ({technique}) via parameter '{}'",
-                                param.name
-                            ),
+                            format!("SSRF ({technique}) via parameter '{}'", param.name),
                             format!(
                                 "Server-side request forgery detected in parameter '{}' at {}. \
                                  The application made a request to an internal resource ({}).",

compliance-dast/src/agents/xss.rs

@@ -17,26 +17,11 @@ impl XssAgent {
     fn payloads(&self) -> Vec<(&str, &str)> {
         vec![
             ("<script>alert(1)</script>", "basic script injection"),
-            (
-                "<img src=x onerror=alert(1)>",
-                "event handler injection",
-            ),
-            (
-                "<svg/onload=alert(1)>",
-                "svg event handler",
-            ),
-            (
-                "javascript:alert(1)",
-                "javascript protocol",
-            ),
-            (
-                "'\"><script>alert(1)</script>",
-                "attribute breakout",
-            ),
-            (
-                "<body onload=alert(1)>",
-                "body event handler",
-            ),
+            ("<img src=x onerror=alert(1)>", "event handler injection"),
+            ("<svg/onload=alert(1)>", "svg event handler"),
+            ("javascript:alert(1)", "javascript protocol"),
+            ("'\"><script>alert(1)</script>", "attribute breakout"),
+            ("<body onload=alert(1)>", "body event handler"),
         ]
     }
 }
@@ -65,10 +50,7 @@ impl DastAgent for XssAgent {
             for param in &endpoint.parameters {
                 for (payload, technique) in self.payloads() {
                     let test_url = if endpoint.method == "GET" {
-                        format!(
-                            "{}?{}={}",
-                            endpoint.url, param.name, payload
-                        )
+                        format!("{}?{}={}", endpoint.url, param.name, payload)
                     } else {
                         endpoint.url.clone()
                     };

compliance-dast/src/crawler/mod.rs

@@ -28,8 +28,8 @@ impl WebCrawler {
         base_url: &str,
         excluded_paths: &[String],
     ) -> Result<Vec<DiscoveredEndpoint>, CoreError> {
-        let base = Url::parse(base_url)
-            .map_err(|e| CoreError::Dast(format!("Invalid base URL: {e}")))?;
+        let base =
+            Url::parse(base_url).map_err(|e| CoreError::Dast(format!("Invalid base URL: {e}")))?;
 
         let mut visited: HashSet<String> = HashSet::new();
         let mut endpoints: Vec<DiscoveredEndpoint> = Vec::new();
@@ -95,12 +95,15 @@ impl WebCrawler {
             let document = Html::parse_document(&body);
 
             // Extract links
-            let link_selector =
-                Selector::parse("a[href]").unwrap_or_else(|_| Selector::parse("a").expect("valid selector"));
+            let link_selector = match Selector::parse("a[href]") {
+                Ok(s) => s,
+                Err(_) => continue,
+            };
             for element in document.select(&link_selector) {
                 if let Some(href) = element.value().attr("href") {
                     if let Some(absolute_url) = self.resolve_url(&base, &url, href) {
-                        if self.is_same_origin(&base, &absolute_url) && !visited.contains(&absolute_url)
+                        if self.is_same_origin(&base, &absolute_url)
+                            && !visited.contains(&absolute_url)
                         {
                             queue.push((absolute_url, depth + 1));
                         }
@@ -109,18 +112,18 @@ impl WebCrawler {
             }
 
             // Extract forms
-            let form_selector = Selector::parse("form")
-                .unwrap_or_else(|_| Selector::parse("form").expect("valid selector"));
-            let input_selector = Selector::parse("input, select, textarea")
-                .unwrap_or_else(|_| Selector::parse("input").expect("valid selector"));
+            let form_selector = match Selector::parse("form") {
+                Ok(s) => s,
+                Err(_) => continue,
+            };
+            let input_selector = match Selector::parse("input, select, textarea") {
+                Ok(s) => s,
+                Err(_) => continue,
+            };
 
             for form in document.select(&form_selector) {
                 let action = form.value().attr("action").unwrap_or("");
-                let method = form
-                    .value()
-                    .attr("method")
-                    .unwrap_or("GET")
-                    .to_uppercase();
+                let method = form.value().attr("method").unwrap_or("GET").to_uppercase();
 
                 let form_url = self
                     .resolve_url(&base, &url, action)
@@ -128,20 +131,12 @@ impl WebCrawler {
 
                 let mut params = Vec::new();
                 for input in form.select(&input_selector) {
-                    let name = input
-                        .value()
-                        .attr("name")
-                        .unwrap_or("")
-                        .to_string();
+                    let name = input.value().attr("name").unwrap_or("").to_string();
                     if name.is_empty() {
                         continue;
                     }
 
-                    let input_type = input
-                        .value()
-                        .attr("type")
-                        .unwrap_or("text")
-                        .to_string();
+                    let input_type = input.value().attr("type").unwrap_or("text").to_string();
 
                     let location = if method == "GET" {
                         "query".to_string()

compliance-dast/src/orchestrator/state_machine.rs

@@ -149,11 +149,8 @@ impl DastOrchestrator {
         let t2 = target.clone();
         let c2 = context.clone();
         let h2 = http.clone();
-        let xss_handle = tokio::spawn(async move {
-            crate::agents::xss::XssAgent::new(h2)
-                .run(&t2, &c2)
-                .await
-        });
+        let xss_handle =
+            tokio::spawn(async move { crate::agents::xss::XssAgent::new(h2).run(&t2, &c2).await });
 
         let t3 = target.clone();
         let c3 = context.clone();
@@ -167,11 +164,10 @@ impl DastOrchestrator {
         let t4 = target.clone();
         let c4 = context.clone();
         let h4 = http.clone();
-        let ssrf_handle = tokio::spawn(async move {
-            crate::agents::ssrf::SsrfAgent::new(h4)
-                .run(&t4, &c4)
-                .await
-        });
+        let ssrf_handle =
+            tokio::spawn(
+                async move { crate::agents::ssrf::SsrfAgent::new(h4).run(&t4, &c4).await },
+            );
 
         let t5 = target.clone();
         let c5 = context.clone();
@@ -182,8 +178,13 @@ impl DastOrchestrator {
                 .await
         });
 
-        let handles: Vec<tokio::task::JoinHandle<Result<Vec<DastFinding>, CoreError>>> =
-            vec![sqli_handle, xss_handle, auth_handle, ssrf_handle, api_handle];
+        let handles: Vec<tokio::task::JoinHandle<Result<Vec<DastFinding>, CoreError>>> = vec![
+            sqli_handle,
+            xss_handle,
+            auth_handle,
+            ssrf_handle,
+            api_handle,
+        ];
 
         let mut all_findings = Vec::new();
         for handle in handles {

compliance-dast/src/recon/mod.rs

@@ -81,10 +81,9 @@ impl ReconAgent {
         ];
         for header in &missing_security {
             if !headers.contains_key(*header) {
-                result.interesting_headers.insert(
-                    format!("missing:{header}"),
-                    "Not present".to_string(),
-                );
+                result
+                    .interesting_headers
+                    .insert(format!("missing:{header}"), "Not present".to_string());
             }
         }
 
@@ -122,10 +121,10 @@ impl ReconAgent {
 
         let body_lower = body.to_lowercase();
         for (tech, pattern) in &patterns {
-            if body_lower.contains(&pattern.to_lowercase()) {
-                if !result.technologies.contains(&tech.to_string()) {
-                    result.technologies.push(tech.to_string());
-                }
+            if body_lower.contains(&pattern.to_lowercase())
+                && !result.technologies.contains(&tech.to_string())
+            {
+                result.technologies.push(tech.to_string());
             }
         }
     }