Skip to content
This repository has been archived by the owner on Oct 8, 2024. It is now read-only.

Commit

Permalink
feat: add cookie options, fix input screens
Browse files Browse the repository at this point in the history
  • Loading branch information
PThorpe92 committed Apr 5, 2024
1 parent 00e3213 commit ded63ba
Show file tree
Hide file tree
Showing 15 changed files with 263 additions and 79 deletions.
51 changes: 26 additions & 25 deletions src/app.rs
Original file line number Diff line number Diff line change
Expand Up @@ -124,8 +124,8 @@ impl<'a> App<'a> {

pub fn goto_screen(&mut self, screen: &Screen) {
self.input.reset();
self.screen_stack.push(screen.clone());
self.current_screen = screen.clone();
self.screen_stack.push(screen.clone());
self.cursor = 0;
match screen {
Screen::Method => {
Expand Down Expand Up @@ -178,7 +178,9 @@ impl<'a> App<'a> {
pub fn go_back_screen(&mut self) {
let last = self.screen_stack.pop().unwrap_or_default(); // current screen
match self.screen_stack.last().cloned() {
Some(screen) if screen == last => self.go_back_screen(),
Some(screen) if std::mem::discriminant(&screen) == std::mem::discriminant(&last) => {
self.go_back_screen()
}
Some(
Screen::InputMenu(_)
| Screen::CmdMenu(_)
Expand All @@ -187,13 +189,14 @@ impl<'a> App<'a> {
) => self.go_back_screen(),
Some(Screen::RequestBodyInput) => self.goto_screen(&Screen::Method),
Some(Screen::Error(_)) => self.goto_screen(&Screen::Home),
Some(Screen::RequestMenu(ref e)) => {
if !e.as_ref().is_some_and(|err| err.is_error()) {
self.goto_screen(&Screen::RequestMenu(e.clone()));
Some(Screen::RequestMenu(e)) => {
if e.as_ref().is_some() {
self.goto_screen(&Screen::RequestMenu(None));
} else {
self.goto_screen(&Screen::Method);
}
}
Some(Screen::Method) => self.goto_screen(&Screen::Home),
Some(screen) => {
self.goto_screen(&screen);
}
Expand Down Expand Up @@ -298,10 +301,6 @@ impl<'a> App<'a> {
}
}

pub fn get_url(&self) -> &str {
self.command.get_url()
}

pub fn select_item(&mut self) {
if let Some(state) = self.state.as_mut() {
if let Some(selected) = state.selected() {
Expand Down Expand Up @@ -410,18 +409,6 @@ impl<'a> App<'a> {
Ok(())
}

pub fn has_auth(&self) -> bool {
self.command.has_auth()
}

pub fn has_unix_socket(&self) -> bool {
self.command.has_unix_socket()
}

pub fn has_url(&self) -> bool {
!self.command.get_url().is_empty()
}

pub fn delete_saved_key(&mut self, index: i32) -> Result<(), rusqlite::Error> {
self.db.as_ref().delete_key(index)?;
self.goto_screen(&Screen::SavedKeys);
Expand Down Expand Up @@ -474,7 +461,10 @@ impl<'a> App<'a> {
AppOptions::UserAgent(_) => self.command.set_user_agent(""),
AppOptions::Referrer(_) => self.command.set_referrer(""),
AppOptions::RequestBody(_) => self.command.set_request_body(""),
AppOptions::Cookie(_) => self.command.remove_headers(&opt.get_value()),
AppOptions::CookieJar(_) => self.command.set_cookie_jar(&opt.get_value()),
AppOptions::CookiePath(_) => self.command.set_cookie_path(&opt.get_value()),
AppOptions::NewCookie(_) => self.command.add_cookie(&opt.get_value()),
AppOptions::NewCookieSession => self.command.reset_cookie_session(),
AppOptions::Headers(_) => self.command.remove_headers(&opt.get_value()),
AppOptions::Auth(_) => self.command.set_auth(crate::request::curl::AuthKind::None),
AppOptions::EnableHeaders => self.command.enable_response_headers(false),
Expand All @@ -500,6 +490,7 @@ impl<'a> App<'a> {
match opt {
// push headers, reset everything else
AppOptions::Headers(_) => true,
AppOptions::NewCookie(_) => true,
_ => !self.has_app_option(opt),
}
}
Expand Down Expand Up @@ -591,10 +582,14 @@ impl<'a> App<'a> {

AppOptions::Outfile(outfile) => self.command.set_outfile(&outfile),

AppOptions::Cookie(cookie) => self.command.add_cookie(&cookie),
AppOptions::NewCookie(cookie) => self.command.add_cookie(&cookie),

AppOptions::CookieJar(cookie) => self.command.set_cookie_jar(&cookie),

AppOptions::Response(resp) => self.command.set_response(&resp),

AppOptions::CookiePath(cookie) => self.command.set_cookie_path(&cookie),

AppOptions::Referrer(referrer) => self.command.set_referrer(&referrer),

AppOptions::UserAgent(agent) => self.command.set_user_agent(&agent),
Expand Down Expand Up @@ -648,12 +643,18 @@ impl<'a> App<'a> {
self.command.set_referrer(referrer);
}
}
AppOptions::Cookie(_) => {
if let AppOptions::Cookie(ref mut cookie) = opt {
AppOptions::CookiePath(_) => {
if let AppOptions::CookiePath(ref mut cookie) = opt {
option.replace_value(cookie.clone());
self.command.add_cookie(cookie);
}
}
AppOptions::CookieJar(_) => {
if let AppOptions::CookieJar(ref mut cookie) = opt {
option.replace_value(cookie.clone());
self.command.set_cookie_jar(cookie);
}
}
AppOptions::CaPath(_) => {
if let AppOptions::CaPath(ref ca_path) = opt {
option.replace_value(String::from(ca_path));
Expand Down
12 changes: 10 additions & 2 deletions src/display/inputopt.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,11 @@ pub enum InputOpt {
UnixSocket,
UserAgent,
MaxRedirects,
Cookie,
CookiePath,
NewCookie,
CookieJar,
CookieValue(String), // store the name
CookieExpires(String), // store the rest
FtpAccount,
CaPath,
CaCert,
Expand Down Expand Up @@ -52,7 +56,10 @@ impl Display for InputOpt {
InputOpt::UnixSocket => write!(f, "| Unix Socket"),
InputOpt::UserAgent => write!(f, "| User Agent"),
InputOpt::MaxRedirects => write!(f, "| Max Redirects"),
InputOpt::Cookie => write!(f, "| Cookie"),
InputOpt::NewCookie => write!(f, "| Cookie"),
InputOpt::CookiePath => write!(f, "| Cookie"),
InputOpt::CookieValue(_) => write!(f, "| Cookie Val"),
InputOpt::CookieExpires(_) => write!(f, "| Cookie Expires"),
InputOpt::CaPath => write!(f, "| Ca Path"),
InputOpt::CaCert => write!(f, "| Ca Cert"),
InputOpt::VerifyPeer => write!(f, "| Verify Peer DNS-Over-HTTPS"),
Expand All @@ -63,6 +70,7 @@ impl Display for InputOpt {
InputOpt::RenameCollection(_) => write!(f, "| Rename Collection"),
InputOpt::RequestError(ref err) => write!(f, "| Error: {}", err),
InputOpt::Method(method) => write!(f, "| Method: {}", method),
InputOpt::CookieJar => write!(f, "| Cookie Jar"),
}
}
}
18 changes: 14 additions & 4 deletions src/display/menuopts.rs
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,8 @@ pub const DISPLAY_OPT_TCP_KEEPALIVE: &str = " Enable TCP keepalive 󰗶 ";
pub const DISPLAY_OPT_MAX_REC: &str = " Specify recursive depth: ";
pub const DISPLAY_OPT_OUTFILE: &str = " Specify output filepath: ";
pub const DISPLAY_OPT_REFERRER: &str = " Specify Referrer: ";
pub const DISPLAY_OPT_COOKIE: &str = " Add Cookie: ";
pub const DISPLAY_OPT_COOKIE: &str = " Cookie Path: ";
pub const DISPLAY_OPT_COOKIE_JAR: &str = " Cookie Jar: ";
pub const DISPLAY_OPT_USERAGENT: &str = " Specify User-Agent: ";
pub const DISPLAY_OPT_PROXY_TUNNEL: &str = " Enable HTTP Proxy-Tunnel 󱠾 ";
pub const DISPLAY_OPT_URL: &str = " Request URL: ";
Expand Down Expand Up @@ -141,7 +142,7 @@ pub const COLLECTION_ALERT_MENU_OPTS: [&str; 4] = [
pub const REQUEST_MENU_OPTIONS: [&str; 12] = [
"Add a URL 󰖟 ",
"Add a file for uploads  ",
"Add Unix Socket address 󰟩 ",
"Cookie options 󰖟  ",
"Authentication 󰯄 ",
"Header Options  ",
"Enable verbose output [-v]",
Expand All @@ -152,6 +153,14 @@ pub const REQUEST_MENU_OPTIONS: [&str; 12] = [
"More Options  ",
"Clear all options  ",
];

pub const COOKIE_MENU_OPTIONS: [&str; 4] = [
"Set Cookie file path (Use Cookies) 󰖟  ",
"Set Cookie-Jar path (Storage) 󰖟  ",
"Add New Cookie 󰆘 ",
"Reset Cookie Session 󰆘 ",
];

pub const METHOD_MENU_OPTIONS: [&str; 7] = [
"OTHER (custom command)",
"GET",
Expand All @@ -176,10 +185,11 @@ pub const AUTHENTICATION_MENU_OPTIONS: [&str; 6] = [
"Ntlm",
"SPNEGO",
];
pub const MORE_FLAGS_MENU: [&str; 13] = [
pub const MORE_FLAGS_MENU: [&str; 14] = [
"Follow Redirects 󱀀 ",
"Specify Max redirects 󱀀 ",
"Add Cookie 󰆘 ",
"Add New Cookie 󰆘 ",
"Set Cookie-Jar path (Storage) 󰆘 ",
"Enable HTTP Proxy-Tunnel 󱠾 ",
"Unrestricted Auth  ",
"Specify Referrer 󰆽 ",
Expand Down
20 changes: 15 additions & 5 deletions src/display/mod.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use std::fmt::{Display, Formatter};

use crate::{
display::menuopts::{DISPLAY_OPT_MAX_REDIRECTS, DISPLAY_OPT_REFERRER},
display::menuopts::{DISPLAY_OPT_COOKIE_JAR, DISPLAY_OPT_MAX_REDIRECTS, DISPLAY_OPT_REFERRER},
request::curl::AuthKind,
};

Expand Down Expand Up @@ -46,7 +46,6 @@ pub mod menuopts;
#[derive(Debug, Clone, PartialEq)]
pub enum AppOptions {
Verbose,
// TODO: support more headers
Headers(String),
URL(String),
Outfile(String),
Expand All @@ -56,7 +55,8 @@ pub enum AppOptions {
SaveToken,
UnixSocket(String),
FollowRedirects,
Cookie(String),
CookieJar(String),
CookiePath(String),
EnableHeaders,
ContentHeaders(HeaderKind),
ProgressBar,
Expand All @@ -71,7 +71,9 @@ pub enum AppOptions {
UnrestrictedAuth,
MaxRedirects(usize),
UploadFile(String),
NewCookie(String),
RequestBody(String),
NewCookieSession,
}

impl AppOptions {
Expand All @@ -97,7 +99,10 @@ impl AppOptions {
AppOptions::UnixSocket(ref mut socket) => {
*socket = val;
}
AppOptions::Cookie(ref mut cookie) => {
AppOptions::CookiePath(ref mut cookie) => {
*cookie = val;
}
AppOptions::CookieJar(ref mut cookie) => {
*cookie = val;
}
AppOptions::Referrer(ref mut referrer) => {
Expand Down Expand Up @@ -135,6 +140,7 @@ impl AppOptions {
AppOptions::UnixSocket(socket) => {
format!("{}{}", DISPLAY_OPT_UNIX_SOCKET, socket.clone())
}
AppOptions::NewCookie(cookie) => format!("{}{}", DISPLAY_OPT_COOKIE, cookie.clone()),
AppOptions::EnableHeaders => DISPLAY_OPT_HEADERS.to_string(),
AppOptions::ProgressBar => String::from(DISPLAY_OPT_PROGRESS_BAR),
AppOptions::FailOnError => String::from(DISPLAY_OPT_FAIL_ON_ERROR),
Expand All @@ -143,7 +149,11 @@ impl AppOptions {
AppOptions::MaxRedirects(max_redirects) => {
format!("{}{}", DISPLAY_OPT_MAX_REDIRECTS, max_redirects)
}
AppOptions::Cookie(cookie) => format!("{}{}", DISPLAY_OPT_COOKIE, cookie.clone()),
AppOptions::NewCookieSession => String::from("New Cookie Session"),
AppOptions::CookiePath(cookie) => format!("{}{}", DISPLAY_OPT_COOKIE, cookie.clone()),
AppOptions::CookieJar(cookie) => {
format!("{}{}", DISPLAY_OPT_COOKIE_JAR, cookie.clone())
}
AppOptions::Referrer(referrer) => {
format!("{}{}", DISPLAY_OPT_REFERRER, referrer.clone())
}
Expand Down
53 changes: 51 additions & 2 deletions src/request/curl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -284,7 +284,7 @@ impl Display for AuthKind {
AuthKind::None => write!(f, "None"),
AuthKind::Ntlm => write!(f, "NTLM"),
AuthKind::Basic(login) => write!(f, "Basic: {}", login),
AuthKind::Bearer(token) => write!(f, "Authorization: {}", token),
AuthKind::Bearer(token) => write!(f, "Authorization: Bearer {}", token),
AuthKind::Digest(login) => write!(f, "Digest Auth: {}", login),
AuthKind::AwsSigv4 => write!(f, "AWS SignatureV4"),
AuthKind::Spnego => write!(f, "SPNEGO Auth"),
Expand Down Expand Up @@ -377,6 +377,31 @@ impl<'a> Curl<'a> {
self.resp = Some(String::from(response));
}

pub fn get_cookie_path(&self) -> Option<String> {
let flag = &CurlFlag::CookieFile(CurlFlagType::CookieFile.get_value(), None);
self.opts
.iter()
.find(|pos| *pos == flag)
.and_then(|pos| pos.get_arg())
}

pub fn set_cookie_path(&mut self, path: &str) {
let flag = CurlFlag::CookieFile(CurlFlagType::CookieFile.get_value(), None);
self.toggle_flag(&flag);
self.curl.cookie_file(path).unwrap();
}

pub fn set_cookie_jar(&mut self, path: &str) {
let flag = CurlFlag::CookieJar(CurlFlagType::CookieJar.get_value(), None);
self.toggle_flag(&flag);
self.curl.cookie_jar(path).unwrap();
}
pub fn reset_cookie_session(&mut self) {
let flag = CurlFlag::CookieSession(CurlFlagType::CookieSession.get_value(), None);
self.toggle_flag(&flag);
self.curl.cookie_session(true).unwrap();
}

pub fn get_response(&self) -> String {
self.resp.clone().unwrap_or_default()
}
Expand Down Expand Up @@ -484,6 +509,19 @@ impl<'a> Curl<'a> {
let flag = &CurlFlag::UnixSocket(CurlFlagType::UnixSocket.get_value(), None);
self.has_flag(flag)
}

pub fn get_unix_socket(&self) -> Option<String> {
if self.has_unix_socket() {
let flag = &CurlFlag::UnixSocket(CurlFlagType::UnixSocket.get_value(), None);
self.opts
.iter()
.find(|pos| *pos == flag)
.and_then(|pos| pos.get_arg())
} else {
None
}
}

pub fn set_method(&mut self, method: Method) {
match method {
Method::Get => self.set_get_method(),
Expand Down Expand Up @@ -538,6 +576,14 @@ impl<'a> Curl<'a> {
self.curl.progress(on).unwrap();
}

pub fn get_cookie_jar_path(&self) -> Option<String> {
let flag = &CurlFlag::CookieJar(CurlFlagType::CookieJar.get_value(), None);
self.opts
.iter()
.find(|pos| *pos == flag)
.and_then(|pos| pos.get_arg())
}

pub fn set_content_header(&mut self, kind: HeaderKind) {
if kind == HeaderKind::None && self.headers.is_some() {
self.headers
Expand Down Expand Up @@ -902,7 +948,7 @@ impl<'a> Curl<'a> {
pub fn set_bearer_auth(&mut self, token: &str) {
self.add_flag(CurlFlag::Bearer(
CurlFlagType::Bearer.get_value(),
Some(format!("Authorization: {token}")),
Some(format!("Authorization: Bearer {token}")),
));
self.auth = AuthKind::Bearer(String::from(token));
}
Expand Down Expand Up @@ -1026,6 +1072,7 @@ define_curl_flags! {
CertInfo("--certinfo"),
Headers("-H"),
Digest("--digest"),
CookieSession("--junk-session-cookies"),
Basic("-H"),
AnyAuth("--any-auth"),
UnixSocket("--unix-socket"),
Expand All @@ -1049,6 +1096,8 @@ define_curl_flags! {
SpnegoAuth("--negotiate -u:"),
Progress("--progress-bar"),
RequestBody("--data"),
CookieFile("-c"),
CookieJar("-b"),
}

#[cfg(test)]
Expand Down
4 changes: 2 additions & 2 deletions src/screens/auth.rs
Original file line number Diff line number Diff line change
Expand Up @@ -56,14 +56,14 @@ pub fn handle_authentication_screen(app: &mut App, frame: &mut Frame<'_>) {
}
}
4 => {
if app.has_auth() {
if app.command.has_auth() {
app.remove_app_option(&AppOptions::Auth(AuthKind::None));
}
app.add_app_option(AppOptions::Auth(AuthKind::Spnego));
app.goto_screen(&Screen::RequestMenu(None));
}
5 => {
if app.has_auth() {
if app.command.has_auth() {
app.remove_app_option(&AppOptions::Auth(AuthKind::None));
}
app.add_app_option(AppOptions::Auth(AuthKind::Ntlm));
Expand Down
Loading

0 comments on commit ded63ba

Please sign in to comment.