From 074ba2bc10e504ae9fe129d370a308e2e8a3bf06 Mon Sep 17 00:00:00 2001
From: kirbylife <kirbylife@protonmail.com>
Date: Mon, 30 Aug 2021 14:29:46 -0500
Subject: [PATCH] Add custom user agent possibility

---
 build.rs            | 13 +++++++++++++
 src/request.rs      | 21 ++++++++++++++++++++-
 tests/unit_tests.rs | 27 +++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 1 deletion(-)
 create mode 100644 build.rs

diff --git a/build.rs b/build.rs
new file mode 100644
index 0000000..b13f103
--- /dev/null
+++ b/build.rs
@@ -0,0 +1,13 @@
+use std::env;
+
+fn main() {
+    println!("cargo:rerun-if-env-changed={}", "DEFAULT_USER_AGENT");
+
+    let user_agent = env::var("DEFAULT_USER_AGENT").unwrap_or(format!(
+        "Mozilla/5.0 ({}) Raspa/{}",
+        std::env::consts::OS,
+        env!("CARGO_PKG_VERSION")
+    ));
+
+    println!("cargo:rustc-env=DEFAULT_USER_AGENT={}", user_agent);
+}
diff --git a/src/request.rs b/src/request.rs
index c153633..0c546a8 100644
--- a/src/request.rs
+++ b/src/request.rs
@@ -1,9 +1,14 @@
 use http::Method;
+use std::collections::HashMap;
 use url::ParseError;
 
 use crate::into_url::IntoUrl;
 use crate::response;
 
+// The DEFAULT_USER_AGENT is setted in the build.rs script
+// and can be overwritten in compiled time
+static DEFAULT_USER_AGENT: &str = env!("DEFAULT_USER_AGENT");
+
 #[derive(Clone)]
 pub enum ParameterType {
     Number(isize),
@@ -22,21 +27,30 @@ pub trait RequestBase {
 #[derive(Debug)]
 pub struct Request {
     pub url: url::Url,
+    pub user_agent: String,
     method: Method,
     params: Vec<(String, String)>,
+    headers: HashMap<String, String>,
 }
 
 impl RequestBase for Request {
     fn new<T: IntoUrl>(url: T) -> Result<Request, ParseError> {
         url.into_url().map(|url_parsed| Request {
             url: url_parsed,
+            user_agent: DEFAULT_USER_AGENT.to_string(),
             method: Method::GET,
             params: Vec::new(),
+            headers: HashMap::new(),
         })
     }
 
     fn launch(mut self) -> response::Response {
-        let client = reqwest::blocking::Client::new();
+        let client = reqwest::blocking::Client::builder()
+            .user_agent(self.user_agent)
+            .build()
+            .unwrap();
+
+        // Set GET params
         for (key, value) in &self.params {
             self.url.query_pairs_mut().append_pair(key, value);
         }
@@ -82,6 +96,11 @@ impl Request {
         self
     }
 
+    pub fn set_user_agent<S: AsRef<str>>(mut self, user_agent: S) -> Self {
+        self.user_agent = user_agent.as_ref().to_string();
+        self
+    }
+
     pub fn add_params<S: AsRef<str>, T: ToString>(mut self, params: Vec<(S, T)>) -> Self {
         self.params = params
             .iter()
diff --git a/tests/unit_tests.rs b/tests/unit_tests.rs
index 2e111cc..b05c9ea 100644
--- a/tests/unit_tests.rs
+++ b/tests/unit_tests.rs
@@ -145,3 +145,30 @@ fn parameter_get_request() {
     assert_eq!(resp["args"]["param_bool"].as_str(), Some("true"));
     assert_eq!(resp["args"]["param_float"].as_str(), Some("3.14159"));
 }
+
+#[test]
+fn default_user_agent() {
+    let resp: Value = Request::new("https://httpbin.org/get")
+        .unwrap()
+        .launch()
+        .to_json()
+        .expect("Weird JSON parsed fail");
+    assert_eq!(
+        resp["headers"]["User-Agent"].as_str().unwrap(),
+        std::env::var("DEFAULT_USER_AGENT").unwrap()
+    );
+}
+
+#[test]
+fn custom_user_agent() {
+    let resp: Value = Request::new("https://httpbin.org/get")
+        .unwrap()
+        .set_user_agent("test user agent")
+        .launch()
+        .to_json()
+        .expect("Weird JSON parsed fail");
+    assert_eq!(
+        resp["headers"]["User-Agent"].as_str(),
+        Some("test user agent")
+    );
+}