From 050900b4ce152dc8f8d9cb7a69487d5b269ef932 Mon Sep 17 00:00:00 2001
From: Khaled Nassar <knassar702@gmail.com>
Date: Wed, 16 Nov 2022 20:08:16 +0200
Subject: [PATCH] adding default headers

---
 src/cli/args.rs           |  8 ++++++++
 src/cli/args/validator.rs | 14 ++++++++++++++
 src/cli/logo.rs           |  1 +
 src/cli/main.rs           |  9 +++++++++
 src/cli/txt/logo.txt      |  8 ++++++++
 src/core/mod.rs           |  4 ++++
 src/core/utils/http.rs    |  5 ++++-
 7 files changed, 48 insertions(+), 1 deletion(-)
 create mode 100644 src/cli/logo.rs
 create mode 100644 src/cli/txt/logo.txt

diff --git a/src/cli/args.rs b/src/cli/args.rs
index 60d241c..07d5dac 100644
--- a/src/cli/args.rs
+++ b/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")
diff --git a/src/cli/args/validator.rs b/src/cli/args/validator.rs
index 56dcd2b..5121ddb 100644
--- a/src/cli/args/validator.rs
+++ b/src/cli/args/validator.rs
@@ -1,4 +1,5 @@
 use std::path::Path;
+use std::collections::HashMap;
 
 pub(crate) fn file_exists(file_path: &str) -> Result<(), String> {
     match Path::new(file_path).exists() {
@@ -6,3 +7,16 @@ pub(crate) fn file_exists(file_path: &str) -> Result<(), String> {
         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")
+                )
+        }
+    }
+}
diff --git a/src/cli/logo.rs b/src/cli/logo.rs
new file mode 100644
index 0000000..89910e8
--- /dev/null
+++ b/src/cli/logo.rs
@@ -0,0 +1 @@
+pub const LOGO: &str = include_str!("./txt/logo.txt");
diff --git a/src/cli/main.rs b/src/cli/main.rs
index 18aea0c..b7ffa24 100644
--- a/src/cli/main.rs
+++ b/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,
diff --git a/src/cli/txt/logo.txt b/src/cli/txt/logo.txt
new file mode 100644
index 0000000..636eea4
--- /dev/null
+++ b/src/cli/txt/logo.txt
@@ -0,0 +1,8 @@
+    .'`'.'`'.
+.''.`.  :  .`.''.
+'.    '. .'    .'
+.```  .' '.  ```.
+'..',`  :  `,'..'
+     `-'`'-`))
+            ((    ldb
+             \|
diff --git a/src/core/mod.rs b/src/core/mod.rs
index fa6ae6b..93de668 100644
--- a/src/core/mod.rs
+++ b/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,
diff --git a/src/core/utils/http.rs b/src/core/utils/http.rs
index eb101ec..812ccf4 100644
--- a/src/core/utils/http.rs
+++ b/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(