Merge pull request #54 from rusty-sec/default_headers

adding default headers

src/cli/args.rs

@@ -32,6 +32,14 @@ pub fn cmd_args() -> ArgMatches {
                 .takes_value(true)
                 .required_unless_present("output"),
         )
+        .arg(
+            Arg::with_name("default_headers")
+                .help("Default Request Headers")
+                .validator(validator::valid_json)
+                .takes_value(true)
+                .default_value("{}")
+                .long("headers")
+            )
         .arg(
             Arg::with_name("redirects")
                 .help("Set limit of http redirects")

src/cli/args/validator.rs

@@ -1,8 +1,22 @@
 use std::path::Path;
+use std::collections::HashMap;
 
 pub(crate) fn file_exists(file_path: &str) -> Result<(), String> {
     match Path::new(file_path).exists() {
         true => Ok(()),
         false => Err(format!("the lua Report File doesnt exists: {}", file_path)),
     }
 }
+
+pub(crate) fn valid_json(json_value: &str) -> Result<(), String> {
+    match serde_json::from_str::<HashMap<String, String>>(json_value) {
+        Ok(_json_data) => {
+            Ok(())
+        },
+        Err(_err) => {
+            Err(
+                format!("Headers Value is not a Valid Json data")
+                )
+        }
+    }
+}

src/cli/logo.rs

@@ -0,0 +1 @@
+pub const LOGO: &str = include_str!("./txt/logo.txt");

src/cli/main.rs

@@ -18,6 +18,8 @@
 
 use lotus::Lotus;
 use lotus::RequestOpts;
+use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
+use std::collections::HashMap;
 mod args;
 mod logger;
 
@@ -27,7 +29,14 @@ async fn main() -> Result<(), std::io::Error> {
         logger::init_log(args::cmd_args().value_of("log").unwrap()).unwrap();
     }
     let lottas = Lotus::init(args::cmd_args().value_of("scripts").unwrap().to_string());
+    let parsed_headers: HashMap<String,String> = serde_json::from_str(args::cmd_args().value_of("default_headers").unwrap()).unwrap();
+    let mut user_headers = HeaderMap::new();
+    parsed_headers.iter().for_each(|(headername,headervalue)| {
+        user_headers.insert( HeaderName::from_bytes(headername.as_bytes()).unwrap() , HeaderValue::from_bytes(headervalue.as_bytes()).unwrap());
+    });
+    drop(parsed_headers);
     let request_opts = RequestOpts {
+        headers: user_headers,
         proxy: match args::cmd_args().value_of("proxy") {
             Some(proxy) => Some(proxy.to_string()),
             None => None,

src/cli/txt/logo.txt

@@ -0,0 +1,8 @@
+    .'`'.'`'.
+.''.`.  :  .`.''.
+'.    '. .'    .'
+.```  .' '.  ```.
+'..',`  :  `,'..'
+     `-'`'-`))
+            ((    ldb
+             \|

src/core/mod.rs

@@ -22,6 +22,7 @@ use futures::lock::Mutex;
 use log::{debug, error, info, warn};
 use mlua::Lua;
 use thirtyfour::prelude::*;
+use reqwest::header::HeaderMap;
 use url::Url;
 
 use utils::files::filename_to_string;
@@ -39,8 +40,10 @@ use std::io::Write;
 use std::path::Path;
 use std::sync::Arc;
 
+
 #[derive(Clone)]
 pub struct RequestOpts {
+    pub headers: HeaderMap,
     pub proxy: Option<String>,
     pub timeout: u64,
     pub redirects: u32,
@@ -285,6 +288,7 @@ impl<'a> LuaLoader<'a> {
             .set(
                 "http",
                 http_sender::Sender::init(
+                    self.request.headers.clone(),
                     self.request.proxy.clone(),
                     self.request.timeout,
                     self.request.redirects,

src/core/utils/http.rs

@@ -26,6 +26,7 @@ use tealr::{mlu::FromToLua, TypeName};
 
 #[derive(Clone)]
 pub struct Sender {
+    headers: HeaderMap,
     proxy: Option<String>,
     timeout: u64,
     redirects: u32,
@@ -98,8 +99,9 @@ impl UserData for Sender {
 impl Sender {
     /// Build your own http request module with user option
     ///
-    pub fn init(proxy: Option<String>, timeout: u64, redirects: u32) -> Sender {
+    pub fn init(headers: HeaderMap, proxy: Option<String>, timeout: u64, redirects: u32) -> Sender {
         Sender {
+            headers,
             timeout,
             redirects,
             proxy,
@@ -112,6 +114,7 @@ impl Sender {
                 Client::builder()
                     .timeout(Duration::from_secs(self.timeout))
                     .redirect(redirect::Policy::limited(self.redirects as usize))
+                    .default_headers(self.headers.clone())
                     .proxy(Proxy::all(the_proxy).unwrap())
                     .no_trust_dns()
                     .user_agent(