new: added cookies persistency for http namespace

Author: evilsocket

Cargo.lock

@@ -35,7 +35,7 @@ regex = "1.10.5"
 serde = { version = "1.0.203", features = ["derive", "serde_derive"] }
 serde_trim = "1.1.0"
 serde_yaml = "0.9.34"
-simple-home-dir = "0.3.5"
+simple-home-dir = "0.4.0"
 tokio = "1.38.0"
 xml-rs = "0.8.20"
 duration-string = { version = "0.4.0", optional = true }
@@ -53,6 +53,7 @@ reqwest = { version = "0.12.5", default-features = false, features = [
     "rustls-tls",
 ] }
 url = "2.5.2"
+reqwest_cookie_store = "0.8.0"
 
 [features]
 default = ["ollama", "groq", "openai", "fireworks"]

Cargo.toml

@@ -1,4 +1,4 @@
-use std::{collections::HashMap, sync::Arc, time::Duration};
+use std::{collections::HashMap, path::PathBuf, sync::Arc, time::Duration};
 
 use anyhow::Result;
 use mini_rag::Embedder;
@@ -18,6 +18,21 @@ pub mod serialization;
 pub mod state;
 pub mod task;
 
+pub fn data_path(path: &str) -> Result<PathBuf> {
+    let user_home = match simple_home_dir::home_dir() {
+        Some(path) => path,
+        None => return Err(anyhow!("can't get user home folder")),
+    };
+
+    let inner_path = user_home.join(".nerve").join(path);
+    if !inner_path.exists() {
+        log::info!("creating {} ...", inner_path.display());
+        std::fs::create_dir_all(&inner_path)?;
+    }
+
+    Ok(inner_path)
+}
+
 #[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct Invocation {
     pub action: String,

src/agent/mod.rs

@@ -1,12 +1,16 @@
 use std::{
     collections::HashMap,
+    fs::OpenOptions,
     str::FromStr,
+    sync::Arc,
     time::{Duration, Instant},
 };
 
 use anyhow::Result;
 use async_trait::async_trait;
 use colored::Colorize;
+use lazy_static::lazy_static;
+use reqwest_cookie_store::CookieStoreMutex;
 use url::Url;
 
 use crate::agent::state::SharedState;
@@ -15,6 +19,28 @@ use super::{Action, Namespace, StorageDescriptor};
 
 const DEFAULT_HTTP_SCHEMA: &str = "https";
 
+lazy_static! {
+    static ref COOKIE_STORE: Arc<CookieStoreMutex> = {
+        let cookies_file = crate::agent::data_path("http")
+            .unwrap()
+            .join("cookies.json");
+
+        let file = OpenOptions::new()
+            .read(true)
+            .write(true)
+            .create(true)
+            .truncate(true)
+            .open(&cookies_file)
+            .map(std::io::BufReader::new)
+            .unwrap_or_else(|_| panic!("can't open {}", cookies_file.display()));
+
+        let cookie_store = reqwest_cookie_store::CookieStore::load_json(file)
+            .unwrap_or_else(|_| panic!("can't load {}", cookies_file.display()));
+
+        Arc::new(reqwest_cookie_store::CookieStoreMutex::new(cookie_store))
+    };
+}
+
 #[derive(Debug, Default, Clone)]
 struct ClearHeaders {}
 
@@ -153,9 +179,14 @@ impl Request {
 
     fn create_request(method: &str, target_url: Url) -> Result<reqwest::RequestBuilder> {
         let method = reqwest::Method::from_str(method)?;
-        let mut request = reqwest::Client::new().request(method.clone(), target_url.clone());
-        let query_str = target_url.query().unwrap_or("").to_string();
 
+        let mut request = reqwest::Client::builder()
+            .cookie_provider(COOKIE_STORE.clone())
+            .build()?
+            .request(method.clone(), target_url.clone());
+
+        // get query string if any
+        let query_str = target_url.query().unwrap_or("").to_string();
         // if there're parameters and we're not in GET, set them as the body
         if !query_str.is_empty() && !matches!(method, reqwest::Method::GET) {
             request = request.header(
@@ -212,8 +243,6 @@ impl Action for Request {
         let target_url_str = target_url.to_string();
         let mut request = Self::create_request(method, target_url)?;
 
-        // TODO: handle cookie/session persistency
-
         // add defined headers
         for (key, value) in state.lock().await.get_storage("http-headers")?.iter() {
             request = request.header(key, &value.data);

src/agent/namespaces/http/mod.rs

@@ -9,7 +9,6 @@ use colored::Colorize;
 use duration_string::DurationString;
 use serde::Deserialize;
 use serde_trim::*;
-use simple_home_dir::home_dir;
 
 use super::{variables::interpolate_variables, Task};
 use crate::{
@@ -251,25 +250,22 @@ pub struct Tasklet {
 }
 
 impl Tasklet {
-    pub fn from_path(path: &str, defines: &Vec<String>) -> Result<Self> {
+    pub fn from_path(tasklet_path: &str, defines: &Vec<String>) -> Result<Self> {
         parse_pre_defined_values(defines)?;
 
-        let mut ppath = PathBuf::from_str(path)?;
-
+        let mut tasklet_path = PathBuf::from_str(tasklet_path)?;
         // try to look it up in ~/.nerve/tasklets
-        if !ppath.exists() {
-            let in_home = home_dir()
-                .unwrap()
-                .join(PathBuf::from_str(".nerve/tasklets")?.join(&ppath));
+        if !tasklet_path.exists() {
+            let in_home = crate::agent::data_path("tasklets")?.join(&tasklet_path);
             if in_home.exists() {
-                ppath = in_home;
+                tasklet_path = in_home;
             }
         }
 
-        if ppath.is_dir() {
-            Self::from_folder(ppath.to_str().unwrap())
+        if tasklet_path.is_dir() {
+            Self::from_folder(tasklet_path.to_str().unwrap())
         } else {
-            Self::from_yaml_file(ppath.to_str().unwrap())
+            Self::from_yaml_file(tasklet_path.to_str().unwrap())
         }
     }