use super::HttpHeaders;
use crate::cookie::Cookie;
use crate::error::{Error, FileNotCreatedError};
use std::collections::HashMap;
use std::fs;
use std::fs::File;
use std::io::prelude::*;
use std::path::Path;
use std::time::{SystemTime, UNIX_EPOCH};
use url::Url;
#[derive(Clone, Debug)]
pub struct CookieJar {
jar_file: String,
auto_update: bool,
cookies: HashMap<String, Cookie>,
}
impl CookieJar {
pub fn new() -> Self {
Self {
jar_file: String::new(),
auto_update: false,
cookies: HashMap::new(),
}
}
pub fn from_file(jar_file: &str, _auto_update: bool) -> Result<Self, Error> {
if !Path::new(jar_file).exists() {
return Err(Error::FileNotExists(jar_file.to_string()));
}
let contents = fs::read_to_string(jar_file).unwrap();
let mut jar = Self::from_string(&contents);
jar.jar_file = jar_file.to_string();
jar.auto_update = true;
Ok(jar)
}
pub fn from_string(contents: &String) -> Self {
let mut cookies: HashMap<String, Cookie> = HashMap::new();
for line in contents.split('\n') {
if line.starts_with('#') {
continue;
}
if let Some(cookie) = Cookie::from_line(line) {
cookies.insert(cookie.name.clone(), cookie.clone());
}
}
Self {
jar_file: String::new(),
auto_update: false,
cookies,
}
}
pub fn set_jar_file(&mut self, jar_file: &str) {
self.jar_file = jar_file.to_string();
}
pub fn set_auto_update(&mut self, auto_update: bool) {
self.auto_update = auto_update;
}
pub fn get(&self, name: &str) -> Option<Cookie> {
if let Some(cookie) = self.cookies.get(name) {
return Some(cookie.clone());
}
None
}
pub fn set(&mut self, cookie: &Cookie) {
let name = cookie.name.clone();
*self.cookies.entry(name.clone()).or_insert(cookie.clone()) = cookie.clone();
}
pub fn delete(&mut self, name: &str) {
self.cookies.remove(name);
}
pub fn clear(&mut self) {
self.cookies.clear();
}
pub fn get_http_header(&self, uri: &Url) -> Option<String> {
let mut pairs = Vec::new();
let host = uri.host_str().unwrap();
let host_chk = format!(".{}", host);
for (_name, cookie) in self.iter() {
if (cookie.host != host && cookie.host != host_chk)
|| (!uri.path().starts_with(&cookie.path))
|| (cookie.secure && uri.scheme() != "https")
{
continue;
}
let line = format!("{}={}", cookie.name, cookie.value);
pairs.push(line);
}
if pairs.is_empty() {
return None;
}
Some(pairs.join("; ").to_string())
}
pub fn iter(&self) -> Box<dyn Iterator<Item = (String, Cookie)>> {
Box::new(self.cookies.clone().into_iter())
}
pub fn update_jar(&mut self, headers: &HttpHeaders) {
for line in headers.get_lower_vec("set-cookie") {
let eq_index = line.find('=').unwrap_or(0);
let sc_index = line.find(';').unwrap_or(0);
if eq_index == 0 || sc_index == 0 || eq_index >= sc_index {
continue;
}
let name = line[..eq_index].to_string();
let value = line[eq_index + 1..sc_index].trim().to_string();
if value.is_empty() {
self.delete(name.as_str());
continue;
}
let elem: HashMap<String, String> = line[sc_index + 1..]
.split(';')
.map(|e| {
let (mut ekey, mut evalue) = (e.to_string(), "".to_string());
if let Some(eindex) = e.find('=') {
ekey = e[..eindex].to_lowercase().trim().to_string();
evalue = e[eindex + 1..].trim().to_string();
}
(ekey, evalue)
})
.collect();
let expires: u64 = 0;
if let Some(_max_age) = elem.get(&"max-age".to_string()) {
let _secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
}
let cookie = Cookie {
host: elem
.get(&"domain".to_string())
.unwrap_or(&"".to_string())
.clone(),
path: elem
.get(&"path".to_string())
.unwrap_or(&"/".to_string())
.clone(),
http_only: elem.contains_key(&"httponly".to_string()),
secure: elem.contains_key(&"secure".to_string()),
expires,
name: name.to_string(),
value: line[eq_index + 1..sc_index].trim().to_string(),
};
self.set(&cookie);
}
if self.auto_update {
self.save_jar();
}
}
pub fn save_jar(&mut self) -> Result<(), Error> {
if self.jar_file.is_empty() {
return Ok(());
}
let mut file = match File::create(&self.jar_file) {
Ok(r) => r,
Err(e) => {
return Err(Error::FileNotCreated(FileNotCreatedError {
filename: self.jar_file.clone(),
error: e.to_string(),
}));
}
};
writeln!(
file,
"# Auto-generated by atlas-http (https://crates.io/crates/atlas-http)\n"
);
for (_name, cookie) in self.iter() {
writeln!(file, "{}", Cookie::to_line(&cookie));
}
Ok(())
}
}