Unify command implementations behind common interface
Diff
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(-)
@@ -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<ConcreteLongRunningCommand> {
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<T> {
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 => {
session.data(
channel,
format!("bash: {other}: command not found\n").into(),
);
impl<T> CommandResult<T> {
fn map<N>(self, f: fn(T) -> N) -> CommandResult<N> {
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<Self, &'static str>;
) -> CommandResult<Self>;
async fn data(
async fn stdin(
self,
connection: &mut Connection,
channel: ChannelId,
data: &[u8],
session: &mut Session,
) -> Option<Self>;
) -> CommandResult<Self>;
}
#[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<Self> {
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<Self> {
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 => {
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<Self> {
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),
@@ -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<Self> {
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<Self> {
CommandResult::Exit(0)
}
}
@@ -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<Self> {
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<Self> {
CommandResult::Exit(0)
}
}
@@ -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<Self> {
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<Self> {
CommandResult::Exit(0)
}
}
@@ -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<Self> {
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<Self> {
CommandResult::Exit(0)
}
}
@@ -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<Self, &'static str> {
) -> CommandResult<Self> {
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);
}
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<Self> {
) -> CommandResult<Self> {
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)
}
}
@@ -0,0 +1,6 @@
---
source: pisshoff-server/src/command/uname.rs
expression: output
---
Linux
@@ -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 <https://www.gnu.org/software/coreutils/uname>
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<Self> {
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<Self> {
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);
}
}
@@ -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<Self> {
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<Self> {
CommandResult::Exit(0)
}
}
@@ -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<ConcreteCommand>,
) -> (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),
}