From aa8f63f0f3cdd970181b7dd7fb58952ade29ea90 Mon Sep 17 00:00:00 2001
From: Subhash Dasyam <subhashdasyam@users.noreply.github.com>
Date: Fri, 5 Jul 2024 21:17:37 +0400
Subject: [PATCH] Update mod.rs add handle the response type properly

---
 src/agent/namespaces/http/mod.rs | 50 +++++++++++++++++++++++---------
 1 file changed, 36 insertions(+), 14 deletions(-)

diff --git a/src/agent/namespaces/http/mod.rs b/src/agent/namespaces/http/mod.rs
index d8d8f59..0e267c9 100644
--- a/src/agent/namespaces/http/mod.rs
+++ b/src/agent/namespaces/http/mod.rs
@@ -147,59 +147,81 @@ impl Action for Request {
         let method = reqwest::Method::from_str(attrs.get("method").unwrap())?;
         let parsed = Self::create_url_from(&state, payload.clone()).await?;
         let query_str = parsed.query().unwrap_or("").to_string();
-
+    
         let mut client = reqwest::Client::new().request(method.clone(), parsed.clone());
         let lock = state.lock().await;
         let headers = lock.get_storage("http-headers")?;
-
+    
         for (key, value) in headers.iter() {
             client = client.header(key, &value.data);
         }
-
+    
         // if there're parameters and we're not in GET, set them as the body
         if !query_str.is_empty() {
             if !matches!(method, reqwest::Method::GET) {
                 client = client.body(query_str);
             }
         }
-
+    
         log::info!(
             "{}.{} {} ...",
             "http".bold(),
             method.to_string().yellow(),
             parsed.to_string(),
         );
-
+    
         let start = Instant::now();
         let res = client.send().await?;
         let elaps = start.elapsed();
-
+    
         return if res.status().is_success() {
             let reason = res.status().canonical_reason().unwrap();
             let mut resp = format!("{} {}\n", res.status().as_u16(), &reason);
-
+    
             for (key, val) in res.headers() {
                 resp += &format!("{}: {}\n", key, val.to_str().unwrap());
             }
-
-            // TODO: handle the response type properly
+    
+            // Handle the response type properly
             resp += "\n\n";
-            resp += &res.text().await?;
-
+            match res.headers().get("content-type") {
+                Some(content_type) if content_type.to_str().unwrap().starts_with("application/") || content_type.to_str().unwrap().starts_with("text/") => {
+                    if content_type.to_str().unwrap() == "application/octet-stream" {
+                        // Download the first few bytes to determine if it's binary or text
+                        let partial_content = res.bytes().await?;
+                        let is_binary = partial_content.iter().any(|&byte| byte == 0 || (byte < 32 && byte != 9 && byte != 10 && byte != 13));
+                        
+                        if is_binary {
+                            log::warn!("Ignoring binary data with content type: application/octet-stream");
+                        } else {
+                            resp += std::str::from_utf8(&partial_content).unwrap_or("");
+                        }
+                    } else {
+                        resp += &res.text().await?;
+                    }
+                },
+                Some(content_type) => {
+                    log::warn!("Ignoring non-textual content type: {}", content_type.to_str().unwrap());
+                },
+                None => {
+                    log::warn!("No content type specified in the response");
+                }
+            }
+    
             log::info!(
                 "   {} {} -> {} bytes",
                 reason.green(),
                 format!("({:?})", elaps).dimmed(),
                 resp.len()
             );
-
+    
             Ok(Some(resp))
         } else {
             let reason = res.status().canonical_reason().unwrap();
             let resp = format!("{} {}", res.status().as_u16(), &reason);
-
+    
             log::error!("   {} {}", reason.red(), format!("({:?})", elaps).dimmed(),);
-
+    
             Err(anyhow!(resp))
         };
     }