use cookie::SameSite;
use serde::{Deserialize, Serialize};
use std::convert::{TryFrom, TryInto};
use time::OffsetDateTime;
use webdriver::command::{AddCookieParameters, WebDriverCommand};
use webdriver::common::Date;
use crate::client::Client;
use crate::error;
pub type Cookie<'a> = cookie::Cookie<'a>;
#[derive(Debug, Serialize)]
pub(crate) struct AddCookieParametersWrapper<'a> {
#[serde(with = "AddCookieParameters")]
pub(crate) cookie: &'a AddCookieParameters,
}
#[derive(Debug, Deserialize, Serialize)]
pub(crate) struct WebDriverCookie {
name: String,
value: String,
#[serde(skip_serializing_if = "Option::is_none")]
path: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
domain: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
secure: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none", rename = "httpOnly")]
http_only: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
expiry: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none", rename = "sameSite")]
same_site: Option<String>,
}
impl WebDriverCookie {
fn into_params(self) -> AddCookieParameters {
AddCookieParameters {
name: self.name,
value: self.value,
path: self.path,
domain: self.domain,
secure: self.secure.unwrap_or_default(),
httpOnly: self.http_only.unwrap_or_default(),
expiry: self.expiry.map(Date),
sameSite: self.same_site,
}
}
}
impl TryFrom<WebDriverCookie> for Cookie<'static> {
type Error = error::CmdError;
fn try_from(webdriver_cookie: WebDriverCookie) -> Result<Self, Self::Error> {
let mut cookie = cookie::Cookie::new(webdriver_cookie.name, webdriver_cookie.value);
if let Some(path) = webdriver_cookie.path {
cookie.set_path(path);
}
if let Some(domain) = webdriver_cookie.domain {
cookie.set_domain(domain);
}
if let Some(secure) = webdriver_cookie.secure {
cookie.set_secure(secure);
}
if let Some(http_only) = webdriver_cookie.http_only {
cookie.set_http_only(http_only);
}
if let Some(expiry) = webdriver_cookie.expiry {
let dt = OffsetDateTime::from_unix_timestamp(expiry as i64).ok();
cookie.set_expires(dt);
}
if let Some(same_site) = webdriver_cookie.same_site {
cookie.set_same_site(match &same_site {
x if x.eq_ignore_ascii_case("strict") => SameSite::Strict,
x if x.eq_ignore_ascii_case("lax") => SameSite::Lax,
x if x.eq_ignore_ascii_case("none") => SameSite::None,
_ => {
return Err(error::CmdError::InvalidArgument(
"same_site".to_string(),
same_site,
))
}
});
}
Ok(cookie)
}
}
impl<'a> From<Cookie<'a>> for WebDriverCookie {
fn from(cookie: Cookie<'a>) -> Self {
let name = cookie.name().to_string();
let value = cookie.value().to_string();
let path = cookie.path().map(String::from);
let domain = cookie.domain().map(String::from);
let secure = cookie.secure();
let http_only = cookie.http_only();
let expiry = cookie
.expires()
.and_then(|e| e.datetime().map(|dt| dt.unix_timestamp() as u64));
let same_site = cookie.same_site().map(|x| match x {
SameSite::Strict => "Strict".to_string(),
SameSite::Lax => "Lax".to_string(),
SameSite::None => "None".to_string(),
});
Self {
name,
value,
path,
domain,
secure,
http_only,
expiry,
same_site,
}
}
}
impl Client {
pub async fn get_all_cookies(&self) -> Result<Vec<Cookie<'static>>, error::CmdError> {
let resp = self.issue(WebDriverCommand::GetCookies).await?;
let webdriver_cookies: Vec<WebDriverCookie> = serde_json::from_value(resp)?;
webdriver_cookies
.into_iter()
.map(|raw_cookie| raw_cookie.try_into())
.collect()
}
pub async fn get_named_cookie(&self, name: &str) -> Result<Cookie<'static>, error::CmdError> {
let resp = self
.issue(WebDriverCommand::GetNamedCookie(name.to_string()))
.await?;
let webdriver_cookie: WebDriverCookie = serde_json::from_value(resp)?;
webdriver_cookie.try_into()
}
pub async fn add_cookie(&self, cookie: Cookie<'static>) -> Result<(), error::CmdError> {
let webdriver_cookie: WebDriverCookie = cookie.into();
self.issue(WebDriverCommand::AddCookie(webdriver_cookie.into_params()))
.await?;
Ok(())
}
pub async fn delete_cookie(&self, name: &str) -> Result<(), error::CmdError> {
self.issue(WebDriverCommand::DeleteCookie(name.to_string()))
.await
.map(|_| ())
}
pub async fn delete_all_cookies(&self) -> Result<(), error::CmdError> {
self.issue(WebDriverCommand::DeleteCookies)
.await
.map(|_| ())
}
}