Improve data type guessing

Sh1Yo · Apr 23, 2023 · c73b92f · c73b92f
1 parent 95b9c62
commit c73b92f
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 17 deletions.
diff --git a/src/config/args.rs b/src/config/args.rs
@@ -3,7 +3,7 @@ use crate::{
         structs::Config,
         utils::{convert_to_string_if_some, parse_request},
     },
-    network::utils::DataType,
+    network::utils::{DataType, Headers},
 };
 use clap::{crate_version, App, AppSettings, Arg};
 use std::{collections::HashMap, error::Error, fs, io::{self, Write}};
@@ -442,7 +442,11 @@ Increase the amount of workers to remove the error or use --force.")?;
                     Err("Incorrect --data-type specified")?
                 }
             }
-            None => None,
+            None => if headers.get_value_case_insensitive("content-type") == Some("application/json".to_string()) {
+                Some(DataType::ProbablyJson)
+            } else {
+                None
+            },
         };
 
         let http_version = if args.value_of("http").is_some() {

diff --git a/src/config/utils.rs b/src/config/utils.rs
@@ -83,7 +83,7 @@ pub(super) fn parse_request<'a>(
         match key.to_lowercase().as_str() {
             "content-type" => {
                 if value.contains("json") {
-                    data_type = Some(DataType::Json)
+                    data_type = Some(DataType::ProbablyJson)
                 } else if value.contains("urlencoded") {
                     data_type = Some(DataType::Urlencoded)
                 }

diff --git a/src/network/request.rs b/src/network/request.rs
@@ -444,7 +444,7 @@ impl<'a> RequestDefaults {
         template: Option<S>,
         joiner: Option<S>,
         encode: bool,
-        data_type: Option<DataType>,
+        mut data_type: Option<DataType>,
         invert: bool,
         headers_discovery: bool,
         body: &str,
@@ -461,23 +461,24 @@ impl<'a> RequestDefaults {
             InjectionPlace::Path
         };
 
-        if headers_discovery && custom_headers.iter().any(|x| x.1.contains("%s")) {
-            injection_place = InjectionPlace::HeaderValue;
+        if headers_discovery {
+            data_type = Some(DataType::Headers);
+
+            if custom_headers.iter().any(|x| x.1.contains("%s")) {
+                injection_place = InjectionPlace::HeaderValue;
+            }
         }
 
-        let data_type = if data_type.is_none()
-            || data_type.is_some() && data_type != Some(DataType::ProbablyJson)
-        {
+        let data_type = if data_type != Some(DataType::ProbablyJson) {
             data_type
-
         // explained in DataType enum comments
         // tl.dr. data_type was taken from a parsed request's content-type so we are not 100% sure what did a user mean
-        } else if data_type == Some(DataType::ProbablyJson)
-            && injection_place == InjectionPlace::Body
-        {
+        } else if injection_place == InjectionPlace::Body {
             Some(DataType::Json)
-        } else {
+        } else if injection_place == InjectionPlace::Path {
             Some(DataType::Urlencoded)
+        } else {
+            unreachable!()
         };
 
         let (guessed_template, guessed_joiner, is_json, data_type) =
@@ -540,11 +541,11 @@ impl<'a> RequestDefaults {
         injection_place: &InjectionPlace,
         data_type: Option<DataType>,
     ) -> (&'a str, &'a str, bool, Option<DataType>) {
-        if let Some(data_type) = data_type {
+        if data_type.is_some() && data_type != Some(DataType::Headers) {
             match data_type {
                 // %v isn't within quotes because not every json value needs to be in quotes
-                DataType::Json => ("\"%k\": %v", ", ", true, Some(DataType::Json)),
-                DataType::Urlencoded => ("%k=%v", "&", false, Some(DataType::Urlencoded)),
+                Some(DataType::Json) => ("\"%k\": %v", ", ", true, Some(DataType::Json)),
+                Some(DataType::Urlencoded) => ("%k=%v", "&", false, Some(DataType::Urlencoded)),
                 _ => unreachable!(),
             }
         } else {

diff --git a/src/network/utils.rs b/src/network/utils.rs
@@ -39,6 +39,7 @@ pub enum DataType {
     /// to exclude false positive /?{"ZXxZPLN":"ons9XDZ", ..} or Cookie: {"ZXxZPLN":"ons9XDZ", ..} queries
     // it still can be bypassed with the correct --data-type argument
     ProbablyJson,
+    Headers
 }
 
 /// where to insert parameters