Skip to content

Commit

Permalink
Add headers discovery
Browse files Browse the repository at this point in the history
  • Loading branch information
Sh1Yo committed Jul 22, 2021
1 parent 7aa3ab4 commit 598519d
Show file tree
Hide file tree
Showing 5 changed files with 89 additions and 35 deletions.
18 changes: 12 additions & 6 deletions src/args.rs
Original file line number Diff line number Diff line change
Expand Up @@ -111,8 +111,9 @@ pub fn get_config() -> (Config, usize) {
.arg(
Arg::with_name("headers-discovery")
.long("headers")
.help("Switch to header discovery mode")
.help("Switch to header discovery mode.\nForbidden chars would be automatically removed from headers' names")
.conflicts_with("as-body")
.conflicts_with("param-template")
)
.arg(
Arg::with_name("force")
Expand Down Expand Up @@ -225,7 +226,7 @@ pub fn get_config() -> (Config, usize) {
Arg::with_name("max")
.short("m")
.long("max")
.help("Change the maximum number of parameters. (default is 128/192/256 for query and 512 for body)")
.help("Change the maximum number of parameters. (default is 128/192/256 for query/headers and 512 for body)")
.takes_value(true)
)
.arg(
Expand Down Expand Up @@ -311,6 +312,7 @@ pub fn get_config() -> (Config, usize) {
};

let mut headers: HashMap<String, String> = HashMap::new();
let mut within_headers: bool = false;
if let Some(val) = args.values_of("headers") {
for header in val {
let mut k_v = header.split(':');
Expand All @@ -332,6 +334,10 @@ pub fn get_config() -> (Config, usize) {
k_v.map(|x| ":".to_owned() + x).collect(),
].concat();

if value.contains("%s") {
within_headers = true;
}

headers.insert(key.to_string(), value);
}
};
Expand Down Expand Up @@ -421,10 +427,10 @@ pub fn get_config() -> (Config, usize) {
.unwrap_or("https://something.something")
.to_string();

if !args.is_present("as-body") && !args.is_present("headers-discovery") && url.contains('?') && url.contains('=') && !url.contains("%s") {
if !args.is_present("as-body") && !within_headers && !args.is_present("headers-discovery") && url.contains('?') && url.contains('=') && !url.contains("%s") {
url.push_str("&%s");
path.push_str("&%s");
} else if !args.is_present("as-body") && !args.is_present("headers-discovery") && !url.contains("%s") {
} else if !args.is_present("as-body") && !within_headers &&!args.is_present("headers-discovery") && !url.contains("%s") {
url.push_str("?%s");
path.push_str("?%s");
}
Expand All @@ -445,14 +451,13 @@ pub fn get_config() -> (Config, usize) {
if parameter_template.is_empty() {
if body_type.contains("json") && args.is_present("as-body") {
parameter_template = "\"%k\":\"%v\", ";
} else if args.is_present("headers-discovery") {
} else if within_headers {
parameter_template = "%k=%v; ";
} else {
parameter_template = "%k=%v&";
}
}


let custom_keys: Vec<String> = match args.values_of("custom-parameters") {
Some(val) => {
val.map(|x| x.to_string()).collect()
Expand Down Expand Up @@ -518,6 +523,7 @@ pub fn get_config() -> (Config, usize) {
output_format: args.value_of("output-format").unwrap_or("").to_string(),
as_body: args.is_present("as-body"),
headers_discovery: args.is_present("headers-discovery"),
within_headers,
force: args.is_present("force"),
disable_response_correction: args.is_present("disable-response-correction"),
disable_custom_parameters: args.is_present("disable-custom-parameters"),
Expand Down
11 changes: 7 additions & 4 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,7 @@ async fn run() {

//generate random query for the first request
let query = make_hashmap(
&(0..max).map(|_| random_line(config.value_size)).collect::<Vec<String>>(),
&(0..max).map(|_| random_line(config.value_size*2)).collect::<Vec<String>>(),
config.value_size,
);

Expand All @@ -138,9 +138,12 @@ async fn run() {
std::process::exit(1)
}

for param in heuristic(&initial_response.text) {
if !params.contains(&param) {
params.push(param)

if !config.headers_discovery {
for param in heuristic(&initial_response.text) {
if !params.contains(&param) {
params.push(param)
}
}
}

Expand Down
27 changes: 20 additions & 7 deletions src/requests.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use crate::{
structs::{Config, ResponseData, Stable},
utils::{compare, beautify_html, beautify_json, make_body, make_query, make_header_value, make_hashmap, random_line},
utils::{compare, beautify_html, beautify_json, make_body, make_query, make_header_value, make_hashmap, fix_headers, random_line},
};
use colored::*;
use reqwest::Client;
Expand Down Expand Up @@ -97,7 +97,7 @@ pub async fn random_request(
&config,
&client,
&make_hashmap(
&(0..max).map(|_| random_line(config.value_size)).collect::<Vec<String>>(),
&(0..max).map(|_| random_line(config.value_size*2)).collect::<Vec<String>>(),
config.value_size,
),
reflections
Expand All @@ -107,6 +107,7 @@ pub async fn random_request(
fn create_request(
config: &Config,
query: String,
hashmap_query: &HashMap<String, String>,
client: &Client
) -> reqwest::RequestBuilder {
let url: String = if config.url.contains("%s") {
Expand Down Expand Up @@ -156,13 +157,23 @@ fn create_request(
};

for (key, value) in config.headers.iter() {
if value.contains("%s") && config.headers_discovery {
if value.contains("%s") && config.within_headers {
client = client.header(key, value.replace("%s", &query).replace("{{random}}", &random_line(config.value_size)));
} else {
client = client.header(key, value.replace("{{random}}", &random_line(config.value_size)));
};
}

if config.headers_discovery && !config.within_headers {
for (key, value) in hashmap_query.iter() {

client = match fix_headers(key) {
Some(val) => client.header(&val, value.replace("{{random}}", &random_line(config.value_size))),
None => client.header(key, value.replace("{{random}}", &random_line(config.value_size)))
};
}
}

client
}

Expand All @@ -180,8 +191,10 @@ pub async fn request(
let query: String = if !hashmap_query.is_empty() {
if config.as_body {
make_body(&config, &hashmap_query)
} else if config.headers_discovery {
} else if config.within_headers {
make_header_value(&config, &hashmap_query)
} else if config.headers_discovery {
String::new()
} else {
make_query(&config, &hashmap_query)
}
Expand All @@ -193,7 +206,7 @@ pub async fn request(

let url: &str = &config.url;

let res = match create_request(config, query, client).send().await {
let res = match create_request(config, query, &hashmap_query, client).send().await {
Ok(val) => val,
Err(_) => {
//Try to make a random request instead
Expand All @@ -216,7 +229,7 @@ pub async fn request(
String::new()
};

match create_request(config, random_query.clone(), client).send().await {
match create_request(config, random_query.clone(), &hashmap_query, client).send().await {
Ok(_) => return ResponseData {
text: String::new(),
code: 0,
Expand All @@ -233,7 +246,7 @@ pub async fn request(
};
writeln!(io::stderr(), "[~] error at the {} observed. Wait 50 sec and repeat.", config.url).ok();
std::thread::sleep(Duration::from_secs(50));
match create_request(config, random_query, client).send().await {
match create_request(config, random_query, &hashmap_query, client).send().await {
Ok(_) => return ResponseData {
text: String::new(),
code: 0,
Expand Down
1 change: 1 addition & 0 deletions src/structs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ pub struct Config {
pub test: bool,
pub as_body: bool,
pub headers_discovery: bool,
pub within_headers: bool,
pub verbose: u8,
pub is_json: bool,
pub disable_cachebuster: bool,
Expand Down
67 changes: 49 additions & 18 deletions src/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,19 @@ pub fn heuristic(body: &str) -> Vec<String> {
found
}

//remove forbidden characters from header name, otherwise reqwest throws errors
pub fn fix_headers<'a>(header: &'a str) -> Option<String> {
lazy_static! {
static ref RE: Regex = Regex::new(r"[^!-'*+\-\.0-9a-zA-Z^-`|~]").unwrap();
}

if RE.is_match(header) {
Some(RE.replace_all(header, "").to_string())
} else {
None
}
}

pub fn generate_request(config: &Config, initial_query: &HashMap<String, String>) -> String {
let mut hashmap_query: HashMap<String, String> = HashMap::with_capacity(initial_query.len());
for (k, v) in initial_query.iter() {
Expand All @@ -117,16 +130,18 @@ pub fn generate_request(config: &Config, initial_query: &HashMap<String, String>
let query: String = if !hashmap_query.is_empty() {
if config.as_body {
make_body(&config, &hashmap_query)
} else if config.headers_discovery {
} else if config.within_headers {
make_header_value(&config, &hashmap_query)
} else if config.headers_discovery {
String::new()
} else {
make_query(&config, &hashmap_query)
}
} else {
String::new()
};

let mut req: String = String::with_capacity(1024);
let mut req: String = String::with_capacity(4096);
req.push_str(&config.url);
req.push('\n');
req.push_str(&config.method);
Expand All @@ -147,14 +162,23 @@ pub fn generate_request(config: &Config, initial_query: &HashMap<String, String>
for (key, value) in config.headers.iter() {
req.push_str(key);
req.push_str(": ");
if value.contains("%s") && config.headers_discovery {
if value.contains("%s") && config.headers_discovery && config.within_headers {
req.push_str(&value.replace("%s", &query).replace("{{random}}", &random_line(config.value_size)));
} else {
req.push_str(&value.replace("{{random}}", &random_line(config.value_size)));
}
req.push('\n');
}

if config.headers_discovery && !config.within_headers {
for (key, value) in hashmap_query.iter() {
req.push_str(key);
req.push_str(": ");
req.push_str(&value.replace("{{random}}", &random_line(config.value_size)));
req.push('\n');
}
}

if config.as_body && !query.is_empty() {
req.push('\n');
req.push_str(&query);
Expand Down Expand Up @@ -276,25 +300,13 @@ pub fn parse_request(config: Config, proto: &str, request: &str, custom_paramete
let mut host = String::new();
let mut content_type = String::new();
let mut headers: HashMap<String, String> = config.headers.clone();
let mut within_headers: bool = config.within_headers;
let mut firstline = lines.next()?.split(' ');
let method = firstline.next()?.to_string();
let mut path = firstline.next()?.to_string();

let http2: bool = firstline.next()?.to_string().contains("HTTP/2");

let mut parameter_template = if !custom_parameter_template {
if config.headers_discovery {
String::from("%k=%v; ")
} else {
match config.body_type == "json" {
true => String::from("\"%k\":\"%v\", "),
false => String::from("%k=%v&")
}
}
} else {
config.parameter_template.clone()
};

//read headers
while let Some(line) = lines.next() {
if line.is_empty() {
Expand All @@ -308,6 +320,10 @@ pub fn parse_request(config: Config, proto: &str, request: &str, custom_paramete
k_v.map(|x| ":".to_owned() + x).collect(),
].concat();

if value.contains("%s") {
within_headers = true;
}

match key.to_lowercase().as_str() {
"content-type" => content_type = value.clone(),
"host" => {
Expand All @@ -323,6 +339,19 @@ pub fn parse_request(config: Config, proto: &str, request: &str, custom_paramete
headers.insert(key.to_string(), value);
}

let mut parameter_template = if !custom_parameter_template {
if config.within_headers {
String::from("%k=%v; ")
} else {
match config.body_type == "json" {
true => String::from("\"%k\":\"%v\", "),
false => String::from("%k=%v&")
}
}
} else {
config.parameter_template.clone()
};

let mut body = lines.next().unwrap_or("").to_string();
while let Some(part) = lines.next() {
if !part.is_empty() {
Expand Down Expand Up @@ -351,10 +380,11 @@ pub fn parse_request(config: Config, proto: &str, request: &str, custom_paramete

let mut url = [proto,"://", &host, &path].concat();
let initial_url = url.clone();
if !config.as_body && url.contains('?') && url.contains('=') && !url.contains("%s") {

if !config.as_body && url.contains('?') && !within_headers && !config.headers_discovery && url.contains('=') && !url.contains("%s") {
url.push_str("&%s");
path.push_str("&%s");
} else if !config.as_body {
} else if !config.as_body && !within_headers && !config.headers_discovery {
url.push_str("?%s");
path.push_str("?%s");
}
Expand All @@ -365,6 +395,7 @@ pub fn parse_request(config: Config, proto: &str, request: &str, custom_paramete
host,
path,
headers,
within_headers,
body,
body_type,
parameter_template,
Expand Down

0 comments on commit 598519d

Please sign in to comment.