From 0c3d3c20f08171d9039ae775eeea911e4daef8f0 Mon Sep 17 00:00:00 2001 From: Jordan Doyle Date: Wed, 12 Jul 2023 23:54:41 +0100 Subject: [PATCH] Unify command implementations behind common interface --- pisshoff-server/src/command.rs | 202 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++----------------------------------------------------------------------------------------------------------------------- pisshoff-server/src/command/echo.rs | 34 ++++++++++++++++++++++++++++++++++ pisshoff-server/src/command/exit.rs | 38 ++++++++++++++++++++++++++++++++++++++ pisshoff-server/src/command/ls.rs | 58 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ pisshoff-server/src/command/pwd.rs | 36 ++++++++++++++++++++++++++++++++++++ pisshoff-server/src/command/scp.rs | 28 ++++++++++++++++------------ pisshoff-server/src/command/snapshots/pisshoff_server__command__uname__test__.snap | 6 ++++++ pisshoff-server/src/command/uname.rs | 86 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++------------------- pisshoff-server/src/command/whoami.rs | 32 ++++++++++++++++++++++++++++++++ pisshoff-server/src/subsystem/shell.rs | 79 +++++++++++++++++++++++++++++++++++++++++++++++++++---------------------------- 10 files changed, 421 insertions(+), 178 deletions(-) create mode 100644 pisshoff-server/src/command/echo.rs create mode 100644 pisshoff-server/src/command/exit.rs create mode 100644 pisshoff-server/src/command/ls.rs create mode 100644 pisshoff-server/src/command/pwd.rs create mode 100644 pisshoff-server/src/command/snapshots/pisshoff_server__command__uname__test__.snap create mode 100644 pisshoff-server/src/command/whoami.rs diff --git a/pisshoff-server/src/command.rs b/pisshoff-server/src/command.rs index 8446ae6..04ecea1 100644 --- a/pisshoff-server/src/command.rs +++ b/pisshoff-server/src/command.rs @@ -1,147 +1,111 @@ -pub mod scp; -pub mod uname; - -use crate::{command::scp::Scp, server::Connection}; +mod echo; +mod exit; +mod ls; +mod pwd; +mod scp; +mod uname; +mod whoami; + +use crate::server::Connection; use async_trait::async_trait; -use itertools::{Either, Itertools}; -use std::{f32, fmt::Write, str::FromStr, time::Duration}; +use itertools::Either; use thrussh::{server::Session, ChannelId}; -pub async fn run_command( - args: &[String], - channel: ChannelId, - session: &mut Session, - conn: &mut Connection, -) -> Option { - let Some(command) = args.get(0) else { - return None; - }; - - match command.as_str() { - "echo" => { - session.data( - channel, - format!("{}\n", args.iter().skip(1).join(" ")).into(), - ); - } - "whoami" => { - session.data(channel, format!("{}\n", conn.username()).into()); - } - "pwd" => { - session.data( - channel, - format!("{}\n", conn.file_system().pwd().display()).into(), - ); - } - "ls" => { - let resp = if args.len() == 1 { - conn.file_system().ls(None).join(" ") - } else if args.len() == 2 { - conn.file_system().ls(Some(args.get(1).unwrap())).join(" ") - } else { - let mut out = String::new(); - - for dir in args.iter().skip(1) { - if !out.is_empty() { - out.push_str("\n\n"); - } - - write!(out, "{dir}:").unwrap(); - out.push_str(&conn.file_system().ls(Some(dir)).join(" ")); - } - - out - }; - - if !resp.is_empty() { - session.data(channel, format!("{resp}\n").into()); - } - } - "cd" => { - if args.len() > 2 { - session.data( - channel, - "-bash: cd: too many arguments\n".to_string().into(), - ); - return None; - } - - conn.file_system().cd(args.get(1).map(String::as_str)); - } - "exit" => { - let exit_status = args - .get(1) - .map(String::as_str) - .map_or(Ok(0), u32::from_str) - .unwrap_or(2); +pub enum CommandResult { + ReadStdin(T), + Exit(u32), + Close(u32), +} - session.exit_status_request(channel, exit_status); - session.close(channel); - } - "sleep" => { - if let Some(Ok(secs)) = args.get(1).map(String::as_str).map(f32::from_str) { - tokio::time::sleep(Duration::from_secs_f32(secs)).await; - } - } - "uname" => { - let out = uname::execute(&args[1..]); - session.data(channel, out.into()); - } - "scp" => match Scp::new(&args[1..], channel, session) { - Ok(v) => return Some(ConcreteLongRunningCommand::Scp(v)), - Err(e) => session.data(channel, e.to_string().into()), - }, - other => { - // TODO: fix stderr displaying out of order - session.data( - channel, - format!("bash: {other}: command not found\n").into(), - ); +impl CommandResult { + fn map(self, f: fn(T) -> N) -> CommandResult { + match self { + Self::ReadStdin(val) => CommandResult::ReadStdin(f(val)), + Self::Exit(v) => CommandResult::Exit(v), + Self::Close(v) => CommandResult::Close(v), } } - - None } #[async_trait] -pub trait LongRunningCommand: Sized { - fn new( +pub trait Command: Sized { + async fn new( + connection: &mut Connection, params: &[String], channel: ChannelId, session: &mut Session, - ) -> Result; + ) -> CommandResult; - async fn data( + async fn stdin( self, connection: &mut Connection, channel: ChannelId, data: &[u8], session: &mut Session, - ) -> Option; + ) -> CommandResult; } -#[derive(Debug, Clone)] -pub enum ConcreteLongRunningCommand { - Scp(Scp), -} +macro_rules! define_commands { + ($($name:ident($ty:ty) = $command:expr),*) => { + #[derive(Debug, Clone)] + pub enum ConcreteCommand { + $($name($ty)),* + } -impl ConcreteLongRunningCommand { - pub async fn data( - self, - connection: &mut Connection, - channel: ChannelId, - data: &[u8], - session: &mut Session, - ) -> Option { - match self { - Self::Scp(cmd) => cmd - .data(connection, channel, data, session) - .await - .map(Self::Scp), + impl ConcreteCommand { + pub async fn new( + connection: &mut Connection, + params: &[String], + channel: ChannelId, + session: &mut Session, + ) -> CommandResult { + let Some(command) = params.get(0) else { + return CommandResult::Exit(0); + }; + + match command.as_str() { + $($command => <$ty as Command>::new(connection, ¶ms[1..], channel, session).await.map(Self::$name),)* + other => { + // TODO: fix stderr displaying out of order + session.data( + channel, + format!("bash: {other}: command not found\n").into(), + ); + CommandResult::Exit(1) + } + } + } + + pub async fn stdin( + self, + connection: &mut Connection, + channel: ChannelId, + data: &[u8], + session: &mut Session, + ) -> CommandResult { + match self { + $(Self::$name(cmd) => { + cmd + .stdin(connection, channel, data, session) + .await + .map(Self::$name) + }),* + } + } } } } +define_commands! { + Echo(echo::Echo) = "echo", + Exit(exit::Exit) = "exit", + Ls(ls::Ls) = "ls", + Pwd(pwd::Pwd) = "pwd", + Scp(scp::Scp) = "scp", + Uname(uname::Uname) = "uname", + Whoami(whoami::Whoami) = "whoami" +} + #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub enum Arg<'a> { Operand(&'a str), diff --git a/pisshoff-server/src/command/echo.rs b/pisshoff-server/src/command/echo.rs new file mode 100644 index 0000000..777c1f2 --- /dev/null +++ b/pisshoff-server/src/command/echo.rs @@ -0,0 +1,34 @@ +use crate::{ + command::{Command, CommandResult}, + server::Connection, +}; +use async_trait::async_trait; +use itertools::Itertools; +use thrussh::{server::Session, ChannelId}; + +#[derive(Debug, Clone)] +pub struct Echo {} + +#[async_trait] +impl Command for Echo { + async fn new( + _connection: &mut Connection, + params: &[String], + channel: ChannelId, + session: &mut Session, + ) -> CommandResult { + session.data(channel, format!("{}\n", params.iter().join(" ")).into()); + + CommandResult::Exit(0) + } + + async fn stdin( + self, + _connection: &mut Connection, + _channel: ChannelId, + _data: &[u8], + _session: &mut Session, + ) -> CommandResult { + CommandResult::Exit(0) + } +} diff --git a/pisshoff-server/src/command/exit.rs b/pisshoff-server/src/command/exit.rs new file mode 100644 index 0000000..af4fa2f --- /dev/null +++ b/pisshoff-server/src/command/exit.rs @@ -0,0 +1,38 @@ +use crate::{ + command::{Command, CommandResult}, + server::Connection, +}; +use async_trait::async_trait; +use std::str::FromStr; +use thrussh::{server::Session, ChannelId}; + +#[derive(Debug, Clone)] +pub struct Exit {} + +#[async_trait] +impl Command for Exit { + async fn new( + _connection: &mut Connection, + params: &[String], + _channel: ChannelId, + _session: &mut Session, + ) -> CommandResult { + let exit_status = params + .get(0) + .map(String::as_str) + .map_or(Ok(0), u32::from_str) + .unwrap_or(2); + + CommandResult::Close(exit_status) + } + + async fn stdin( + self, + _connection: &mut Connection, + _channel: ChannelId, + _data: &[u8], + _session: &mut Session, + ) -> CommandResult { + CommandResult::Exit(0) + } +} diff --git a/pisshoff-server/src/command/ls.rs b/pisshoff-server/src/command/ls.rs new file mode 100644 index 0000000..73aacb0 --- /dev/null +++ b/pisshoff-server/src/command/ls.rs @@ -0,0 +1,58 @@ +use crate::{ + command::{Command, CommandResult}, + server::Connection, +}; +use async_trait::async_trait; +use std::fmt::Write; +use thrussh::{server::Session, ChannelId}; + +#[derive(Debug, Clone)] +pub struct Ls {} + +#[async_trait] +impl Command for Ls { + async fn new( + connection: &mut Connection, + params: &[String], + channel: ChannelId, + session: &mut Session, + ) -> CommandResult { + let resp = if params.is_empty() { + connection.file_system().ls(None).join(" ") + } else if params.len() == 1 { + connection + .file_system() + .ls(Some(params.get(0).unwrap())) + .join(" ") + } else { + let mut out = String::new(); + + for dir in params { + if !out.is_empty() { + out.push_str("\n\n"); + } + + write!(out, "{dir}:").unwrap(); + out.push_str(&connection.file_system().ls(Some(dir)).join(" ")); + } + + out + }; + + if !resp.is_empty() { + session.data(channel, format!("{resp}\n").into()); + } + + CommandResult::Exit(0) + } + + async fn stdin( + self, + _connection: &mut Connection, + _channel: ChannelId, + _data: &[u8], + _session: &mut Session, + ) -> CommandResult { + CommandResult::Exit(0) + } +} diff --git a/pisshoff-server/src/command/pwd.rs b/pisshoff-server/src/command/pwd.rs new file mode 100644 index 0000000..6b68362 --- /dev/null +++ b/pisshoff-server/src/command/pwd.rs @@ -0,0 +1,36 @@ +use crate::{ + command::{Command, CommandResult}, + server::Connection, +}; +use async_trait::async_trait; +use thrussh::{server::Session, ChannelId}; + +#[derive(Debug, Clone)] +pub struct Pwd {} + +#[async_trait] +impl Command for Pwd { + async fn new( + connection: &mut Connection, + _params: &[String], + channel: ChannelId, + session: &mut Session, + ) -> CommandResult { + session.data( + channel, + format!("{}\n", connection.file_system().pwd().display()).into(), + ); + + CommandResult::Exit(0) + } + + async fn stdin( + self, + _connection: &mut Connection, + _channel: ChannelId, + _data: &[u8], + _session: &mut Session, + ) -> CommandResult { + CommandResult::Exit(0) + } +} diff --git a/pisshoff-server/src/command/scp.rs b/pisshoff-server/src/command/scp.rs index 9e47f33..c465c7d 100644 --- a/pisshoff-server/src/command/scp.rs +++ b/pisshoff-server/src/command/scp.rs @@ -1,5 +1,5 @@ use crate::{ - command::{Arg, LongRunningCommand}, + command::{Arg, Command, CommandResult}, server::Connection, }; use async_trait::async_trait; @@ -32,12 +32,13 @@ pub struct Scp { } #[async_trait] -impl LongRunningCommand for Scp { - fn new( +impl Command for Scp { + async fn new( + _connection: &mut Connection, params: &[String], channel: ChannelId, session: &mut Session, - ) -> Result { + ) -> CommandResult { let mut path = None; let mut transfer = false; @@ -53,36 +54,39 @@ impl LongRunningCommand for Scp { path = Some(p); } _ => { - return Err(HELP); + session.data(channel, HELP.to_string().into()); + return CommandResult::Exit(1); } } } let Some(path) = path else { - return Err(AMBIGUOUS_TARGET); + session.data(channel, AMBIGUOUS_TARGET.to_string().into()); + return CommandResult::Exit(1); }; if !transfer { - return Err(HELP); + session.data(channel, HELP.to_string().into()); + return CommandResult::Exit(1); } // signal to the client we've started listening session.data(channel, SUCCESS.to_string().into()); - Ok(Self { + CommandResult::ReadStdin(Self { path: PathBuf::new().join(path), pending_data: BytesMut::new(), state: State::Waiting, }) } - async fn data( + async fn stdin( mut self, connection: &mut Connection, channel: ChannelId, data: &[u8], session: &mut Session, - ) -> Option { + ) -> CommandResult { self.pending_data.extend_from_slice(data); let mut exit = false; @@ -119,7 +123,7 @@ impl LongRunningCommand for Scp { } Err(error) => { warn!(%error, "Rejecting scp modes payload"); - return None; + return CommandResult::Exit(1); } } } @@ -158,7 +162,7 @@ impl LongRunningCommand for Scp { self.state = next_state; } - Some(self) + CommandResult::ReadStdin(self) } } diff --git a/pisshoff-server/src/command/snapshots/pisshoff_server__command__uname__test__.snap b/pisshoff-server/src/command/snapshots/pisshoff_server__command__uname__test__.snap new file mode 100644 index 0000000..39cd9d7 --- /dev/null +++ b/pisshoff-server/src/command/snapshots/pisshoff_server__command__uname__test__.snap @@ -0,0 +1,6 @@ +--- +source: pisshoff-server/src/command/uname.rs +expression: output +--- +Linux + diff --git a/pisshoff-server/src/command/uname.rs b/pisshoff-server/src/command/uname.rs index 8d2db25..bd27e00 100644 --- a/pisshoff-server/src/command/uname.rs +++ b/pisshoff-server/src/command/uname.rs @@ -1,5 +1,10 @@ -use crate::command::Arg; +use crate::{ + command::{Arg, Command, CommandResult}, + server::Connection, +}; +use async_trait::async_trait; use bitflags::bitflags; +use thrussh::{server::Session, ChannelId}; bitflags! { #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] @@ -46,7 +51,35 @@ Full documentation or available locally via: info '(coreutils) uname invocation' "; -pub fn execute(params: &[String]) -> String { +#[derive(Debug, Clone)] +pub struct Uname {} + +#[async_trait] +impl Command for Uname { + async fn new( + _connection: &mut Connection, + params: &[String], + channel: ChannelId, + session: &mut Session, + ) -> CommandResult { + let (out, exit_code) = execute(params); + + session.data(channel, out.into()); + CommandResult::Exit(exit_code) + } + + async fn stdin( + self, + _connection: &mut Connection, + _channel: ChannelId, + _data: &[u8], + _session: &mut Session, + ) -> CommandResult { + CommandResult::Exit(0) + } +} + +pub fn execute(params: &[String]) -> (String, u32) { let mut to_print = ToPrint::empty(); let mut filter_unknown = false; @@ -64,26 +97,39 @@ pub fn execute(params: &[String]) -> String { Arg::Short('p') | Arg::Long("processor") => ToPrint::PROCESSOR, Arg::Short('i') | Arg::Long("hardware-platform") => ToPrint::PLATFORM, Arg::Short('o') | Arg::Long("operating-system") => ToPrint::OPERATING_SYSTEM, - Arg::Long("help") => return HELP_STRING.to_string(), - Arg::Long("version") => return VERSION_STRING.to_string(), + Arg::Long("help") => return (HELP_STRING.to_string(), 0), + Arg::Long("version") => return (VERSION_STRING.to_string(), 0), Arg::Operand(operand) => { - return format!( + return ( + format!( "uname: extra operand '{operand}'\nTry 'uname --help' for more information.\n" + ), + 1, ); } Arg::Short(s) => { - return format!( + return ( + format!( "uname: invalid option -- '{s}'\nTry 'uname --help' for more information.\n" + ), + 1, ); } Arg::Long(s) => { - return format!( - "uname: unrecognized option '--{s}'\nTry 'uname --help' for more information.\n" - ) + return ( + format!( + "uname: unrecognized option '--{s}'\nTry 'uname --help' for more information.\n" + ), + 1, + ); } }; } + if to_print.is_empty() { + to_print |= ToPrint::KERNEL_NAME; + } + let mut out = String::with_capacity(105); macro_rules! write { @@ -130,7 +176,7 @@ pub fn execute(params: &[String]) -> String { out.push('\n'); - out + (out, 0) } #[cfg(test)] @@ -138,17 +184,19 @@ mod test { use crate::command::uname::execute; use test_case::test_case; - #[test_case("-a"; "all")] - #[test_case("-snrvmpio"; "all separate")] - #[test_case("-asnrvmpio"; "all separate with all")] - #[test_case("-sn"; "subset")] - #[test_case("-sn --fake"; "unknown long arg param")] - #[test_case("-sn -z"; "unknown short arg param")] - #[test_case("-sn oper"; "unknown operand")] - fn snapshot(input: &str) { + #[test_case("", 0; "none")] + #[test_case("-a", 0; "all")] + #[test_case("-snrvmpio", 0; "all separate")] + #[test_case("-asnrvmpio", 0; "all separate with all")] + #[test_case("-sn", 0; "subset")] + #[test_case("-sn --fake", 1; "unknown long arg param")] + #[test_case("-sn -z", 1; "unknown short arg param")] + #[test_case("-sn oper", 1; "unknown operand")] + fn snapshot(input: &str, expected_exit_code: u32) { let input_parsed = shlex::split(input).unwrap(); - let output = execute(&input_parsed); + let (output, actual_exit_code) = execute(&input_parsed); insta::assert_display_snapshot!(input, output); + assert_eq!(actual_exit_code, expected_exit_code); } } diff --git a/pisshoff-server/src/command/whoami.rs b/pisshoff-server/src/command/whoami.rs new file mode 100644 index 0000000..6937fa6 --- /dev/null +++ b/pisshoff-server/src/command/whoami.rs @@ -0,0 +1,32 @@ +use crate::{ + command::{Command, CommandResult}, + server::Connection, +}; +use async_trait::async_trait; +use thrussh::{server::Session, ChannelId}; + +#[derive(Debug, Clone)] +pub struct Whoami {} + +#[async_trait] +impl Command for Whoami { + async fn new( + connection: &mut Connection, + _params: &[String], + channel: ChannelId, + session: &mut Session, + ) -> CommandResult { + session.data(channel, format!("{}\n", connection.username()).into()); + CommandResult::Exit(0) + } + + async fn stdin( + self, + _connection: &mut Connection, + _channel: ChannelId, + _data: &[u8], + _session: &mut Session, + ) -> CommandResult { + CommandResult::Exit(0) + } +} diff --git a/pisshoff-server/src/subsystem/shell.rs b/pisshoff-server/src/subsystem/shell.rs index 0776c34..effd21e 100644 --- a/pisshoff-server/src/subsystem/shell.rs +++ b/pisshoff-server/src/subsystem/shell.rs @@ -1,5 +1,5 @@ use crate::{ - command::{run_command, ConcreteLongRunningCommand}, + command::{CommandResult, ConcreteCommand}, server::Connection, subsystem::Subsystem, }; @@ -26,6 +26,19 @@ impl Shell { state: State::Prompt, } } + + fn handle_command_result( + &self, + command_result: CommandResult, + ) -> (State, bool) { + match (command_result, self.interactive) { + (CommandResult::ReadStdin(cmd), _) => (State::Running(cmd), true), + (CommandResult::Exit(exit_status), true) => (State::Exit(exit_status), false), + (CommandResult::Exit(exit_status), false) | (CommandResult::Close(exit_status), _) => { + (State::Quit(exit_status), false) + } + } + } } #[async_trait] @@ -39,38 +52,46 @@ impl Subsystem for Shell { data: &[u8], session: &mut Session, ) { - let next = match std::mem::take(&mut self.state) { - State::Prompt => { - let Some(args) = shlex::split(String::from_utf8_lossy(data).as_ref()) else { - return; - }; + loop { + let (next, terminal) = match std::mem::take(&mut self.state) { + State::Prompt => { + let Some(args) = shlex::split(String::from_utf8_lossy(data).as_ref()) else { + return; + }; - connection - .audit_log() - .push_action(AuditLogAction::ExecCommand(ExecCommandEvent { - args: Box::from(args.clone()), - })); + connection + .audit_log() + .push_action(AuditLogAction::ExecCommand(ExecCommandEvent { + args: Box::from(args.clone()), + })); - run_command(&args, channel, session, connection) - .await - .map_or(State::Prompt, State::Running) - } - State::Running(command) => command - .data(connection, channel, data, session) - .await - .map_or(State::Prompt, State::Running), - }; + self.handle_command_result( + ConcreteCommand::new(connection, &args, channel, session).await, + ) + } + State::Running(command) => self + .handle_command_result(command.stdin(connection, channel, data, session).await), + State::Exit(exit_status) => { + session.exit_status_request(channel, exit_status); + (State::Prompt, true) + } + State::Quit(exit_status) => { + session.exit_status_request(channel, exit_status); + session.close(channel); + break; + } + }; + + self.state = next; - if matches!(next, State::Prompt) { - if self.interactive { - session.data(channel, SHELL_PROMPT.to_string().into()); - } else { - session.exit_status_request(channel, 0); - session.close(channel); + if terminal { + break; } } - self.state = next; + if matches!(self.state, State::Prompt) { + session.data(channel, SHELL_PROMPT.to_string().into()); + } } } @@ -78,5 +99,7 @@ impl Subsystem for Shell { enum State { #[default] Prompt, - Running(ConcreteLongRunningCommand), + Running(ConcreteCommand), + Exit(u32), + Quit(u32), } -- libgit2 1.7.2