From edc7582de04210ed336b51983becc555550228ab Mon Sep 17 00:00:00 2001
From: kirbylife <kirbylife@protonmail.com>
Date: Sun, 6 Aug 2023 00:36:50 -0600
Subject: [PATCH] add support to add cookies to the request

---
 src/request.rs      | 38 +++++++++++++++++++++++++++++++++++++-
 tests/unit_tests.rs | 28 ++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/src/request.rs b/src/request.rs
index 0c546a8..582c212 100644
--- a/src/request.rs
+++ b/src/request.rs
@@ -1,5 +1,7 @@
-use http::Method;
+use http::header::{HeaderName, HeaderValue};
+use http::{HeaderMap, Method};
 use std::collections::HashMap;
+use std::str::FromStr;
 use url::ParseError;
 
 use crate::into_url::IntoUrl;
@@ -31,6 +33,7 @@ pub struct Request {
     method: Method,
     params: Vec<(String, String)>,
     headers: HashMap<String, String>,
+    cookies: HashMap<String, String>,
 }
 
 impl RequestBase for Request {
@@ -41,6 +44,7 @@ impl RequestBase for Request {
             method: Method::GET,
             params: Vec::new(),
             headers: HashMap::new(),
+            cookies: HashMap::new(),
         })
     }
 
@@ -55,6 +59,24 @@ impl RequestBase for Request {
             self.url.query_pairs_mut().append_pair(key, value);
         }
 
+        let cookies = self
+            .cookies
+            .iter()
+            .map(|item| format!("{}={}", item.0, item.1))
+            .collect::<Vec<_>>()
+            .join(";");
+
+        if !self.cookies.is_empty() {
+            self.headers.insert("Cookie".to_string(), cookies);
+        }
+
+        let mut headers = HeaderMap::new();
+        for (name, header) in self.headers {
+            let header_name = HeaderName::from_str(name.as_str()).unwrap();
+            let header_value = HeaderValue::from_str(header.as_str()).unwrap();
+            headers.insert(header_name, header_value);
+        }
+
         let resp = (match self.method {
             Method::GET => client.get(self.url.as_str()),
             Method::POST => client.post(self.url.as_str()),
@@ -63,6 +85,7 @@ impl RequestBase for Request {
             Method::DELETE => client.delete(self.url.as_str()),
             _ => unimplemented!(),
         })
+        .headers(headers)
         .send()
         .unwrap();
 
@@ -113,4 +136,17 @@ impl Request {
         self.params
             .push((key.as_ref().to_string(), value.to_string()));
     }
+
+    pub fn add_cookies<S: AsRef<str>, T: ToString>(mut self, cookies: Vec<(S, T)>) -> Self {
+        self.cookies = cookies
+            .iter()
+            .map(|item| (item.0.as_ref().to_string(), item.1.to_string()))
+            .collect();
+        self
+    }
+
+    pub fn add_cookie<S: AsRef<str>, T: ToString>(&mut self, key: S, value: T) {
+        self.cookies
+            .insert(key.as_ref().to_string(), value.to_string());
+    }
 }
diff --git a/tests/unit_tests.rs b/tests/unit_tests.rs
index b05c9ea..77ac7ef 100644
--- a/tests/unit_tests.rs
+++ b/tests/unit_tests.rs
@@ -1,6 +1,7 @@
 use http::Method;
 use http::StatusCode;
 use raspa::request::{ParameterType, Request, RequestBase};
+use raspa::response::ResponseBase;
 use raspa::selector::{Selector, SelectorBase};
 use serde_json::Value;
 // use std::collections::HashMap;
@@ -172,3 +173,30 @@ fn custom_user_agent() {
         Some("test user agent")
     );
 }
+
+#[test]
+fn simple_cookie_test() {
+    let mut req = Request::new("https://httpbin.org/cookies").unwrap();
+    req.add_cookie("test_cookie", "100");
+
+    let resp: Value = req.launch().to_json().expect("Weird JSON parsed fail");
+    assert_eq!(resp["cookies"]["test_cookie"].as_str(), Some("100"));
+}
+
+#[test]
+fn multiple_cookies_test() {
+    let resp: Value = Request::new("https://httpbin.org/cookies")
+        .unwrap()
+        .add_cookies(vec![
+            ("cookie_1", 100),
+            ("cookie_2", 200),
+            ("cookie_3", 300),
+        ])
+        .launch()
+        .to_json()
+        .expect("Weird JSON parsed fail");
+
+    assert_eq!(resp["cookies"]["cookie_1"].as_str(), Some("100"));
+    assert_eq!(resp["cookies"]["cookie_2"].as_str(), Some("200"));
+    assert_eq!(resp["cookies"]["cookie_3"].as_str(), Some("300"));
+}