From ccf2361a61b6e43bb37a1944bd18ef46da4b560c Mon Sep 17 00:00:00 2001
From: Simone Margaritelli <evilsocket@gmail.com>
Date: Thu, 15 Feb 2024 15:59:27 +0100
Subject: [PATCH] fix: using --tcp-ports-http-headers

---
 src/plugins/tcp_ports/grabbers/http.rs | 15 ++++++++++-----
 src/plugins/tcp_ports/grabbers/mod.rs  |  2 +-
 src/plugins/tcp_ports/options.rs       |  2 +-
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/src/plugins/tcp_ports/grabbers/http.rs b/src/plugins/tcp_ports/grabbers/http.rs
index 6e82ac3..d453266 100644
--- a/src/plugins/tcp_ports/grabbers/http.rs
+++ b/src/plugins/tcp_ports/grabbers/http.rs
@@ -6,9 +6,6 @@ use regex::Regex;
 
 use super::Banner;
 
-// TODO: read from args
-static HTTP_HEADERS_OF_INTEREST: &[&str] = &["server", "x-powered-by", "location", "content-type"];
-
 lazy_static! {
     static ref HTML_TITLE_PARSER: Regex =
         Regex::new(r"(?i)<\s*title\s*>([^<]+)<\s*/\s*title\s*>").unwrap();
@@ -41,6 +38,7 @@ pub(crate) fn is_http_port(opts: &options::Options, port: u16) -> (bool, bool) {
 }
 
 pub(crate) async fn http_grabber(
+    opts: &options::Options,
     address: &str,
     port: u16,
     stream: Box<dyn StreamLike>,
@@ -85,6 +83,12 @@ pub(crate) async fn http_grabber(
     if let Ok(resp) = resp {
         // TODO: find a way to collect certificate information if ssl
 
+        let headers_of_interest: Vec<&str> = opts
+            .tcp_ports_http_headers
+            .split(",")
+            .map(|s| s.trim())
+            .filter(|s| !s.is_empty())
+            .collect();
         let mut content_type = String::from("text/html");
 
         // collect headers
@@ -95,7 +99,7 @@ pub(crate) async fn http_grabber(
             if name == "content-type" {
                 content_type = value.to_owned();
             }
-            if HTTP_HEADERS_OF_INTEREST.contains(&name.as_str()) {
+            if headers_of_interest.contains(&name.as_str()) {
                 banner.insert(name, value.to_owned());
             }
         }
@@ -110,7 +114,8 @@ pub(crate) async fn http_grabber(
                         caps.get(1).unwrap().as_str().to_owned(),
                     );
                 }
-            } else if content_type == "application/json" {
+            } else if content_type.starts_with("application/") || content_type.starts_with("text/")
+            {
                 banner.insert("body".to_owned(), body.to_owned());
             }
         } else {
diff --git a/src/plugins/tcp_ports/grabbers/mod.rs b/src/plugins/tcp_ports/grabbers/mod.rs
index 053c24f..1337e35 100644
--- a/src/plugins/tcp_ports/grabbers/mod.rs
+++ b/src/plugins/tcp_ports/grabbers/mod.rs
@@ -19,7 +19,7 @@ pub(crate) async fn grab_banner(
 ) -> Banner {
     let (is_http, with_ssl) = http::is_http_port(opts, port);
     if is_http {
-        return http::http_grabber(address, port, stream, with_ssl, timeout).await;
+        return http::http_grabber(opts, address, port, stream, with_ssl, timeout).await;
     }
 
     // default to an attempt at line grabbing
diff --git a/src/plugins/tcp_ports/options.rs b/src/plugins/tcp_ports/options.rs
index 1d4ba7d..ae3ba82 100644
--- a/src/plugins/tcp_ports/options.rs
+++ b/src/plugins/tcp_ports/options.rs
@@ -19,7 +19,7 @@ pub(crate) struct Options {
     #[clap(long, default_value = "443, 8443")]
     /// Comma separated list of ports for HTTPS grabbing.
     pub tcp_ports_https: String,
-    #[clap(long, default_value = "server, x-powered-by, location")]
+    #[clap(long, default_value = "server, x-powered-by, location, content-type")]
     /// Comma separated list lowercase header names for HTTP/HTTPS grabbing.
     pub tcp_ports_http_headers: String,
 }