🏡 index : ~doyle/pisshoff.git

author Jordan Doyle <jordan@doyle.la> 2024-02-17 14:17:07.0 +00:00:00
committer Jordan Doyle <jordan@doyle.la> 2024-02-17 14:17:07.0 +00:00:00
commit
1e652077744d298c6692b425f75c980c177cebb3 [patch]
tree
88b1d67257684f9c86a984d3c78983a8ac2699c8
parent
51d3c21a11c3999e5d3348a0613d3c6d008353ec
download
1e652077744d298c6692b425f75c980c177cebb3.tar.gz

Implement cat command



Diff

 pisshoff-server/src/command.rs     |  14 +-
 pisshoff-server/src/command/cat.rs | 242 ++++++++++++++++++++++++++++++++++++++-
 pisshoff-server/src/command/ls.rs  |  57 ++++++---
 pisshoff-server/src/file_system.rs | 156 +++++++++++++++++++++++-
 4 files changed, 445 insertions(+), 24 deletions(-)

diff --git a/pisshoff-server/src/command.rs b/pisshoff-server/src/command.rs
index 2ef082b..8d33613 100644
--- a/pisshoff-server/src/command.rs
+++ b/pisshoff-server/src/command.rs
@@ -1,3 +1,4 @@
mod cat;
mod echo;
mod exit;
mod ls;
@@ -6,13 +7,14 @@ mod scp;
mod uname;
mod whoami;

use crate::server::{ConnectionState, ThrusshSession};
use std::{borrow::Cow, fmt::Debug};

use async_trait::async_trait;
use itertools::Either;
use std::borrow::Cow;
use std::fmt::Debug;
use thrussh::ChannelId;

use crate::server::{ConnectionState, ThrusshSession};

#[derive(Debug)]
pub enum CommandResult<T> {
    /// Wait for stdin
@@ -146,7 +148,8 @@ define_commands! {
    Pwd(pwd::Pwd) = b"pwd",
    Scp(scp::Scp) = b"scp",
    Uname(uname::Uname) = b"uname",
    Whoami(whoami::Whoami) = b"whoami"
    Whoami(whoami::Whoami) = b"whoami",
    Cat(cat::Cat) = b"cat"
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
@@ -170,9 +173,10 @@ fn argparse(args: &[String]) -> impl Iterator<Item = Arg<'_>> {

#[cfg(test)]
mod test {
    use super::Arg;
    use test_case::test_case;

    use super::Arg;

    #[test_case("-a", &[Arg::Short('a')]; "single short parameter")]
    #[test_case("-abc", &[Arg::Short('a'), Arg::Short('b'), Arg::Short('c')]; "multiple short parameter")]
    #[test_case("-a --long operand -b -", &[Arg::Short('a'), Arg::Long("long"), Arg::Operand("operand"), Arg::Short('b'), Arg::Operand("-")]; "full hit")]
diff --git a/pisshoff-server/src/command/cat.rs b/pisshoff-server/src/command/cat.rs
new file mode 100644
index 0000000..efd3d0b
--- /dev/null
+++ b/pisshoff-server/src/command/cat.rs
@@ -0,0 +1,242 @@
use std::{collections::VecDeque, path::Path};

use async_trait::async_trait;
use thrussh::ChannelId;

use crate::{
    command::{Command, CommandResult},
    server::{ConnectionState, ThrusshSession},
};

#[derive(Debug, Clone)]
pub struct Cat {
    remaining_params: VecDeque<String>,
    status: u32,
}

impl Cat {
    fn run<S: ThrusshSession + Send>(
        mut self,
        connection: &mut ConnectionState,
        channel: ChannelId,
        session: &mut S,
    ) -> CommandResult<Self> {
        while let Some(param) = self.remaining_params.pop_front() {
            if param == "-" {
                return CommandResult::ReadStdin(self);
            }

            match connection.file_system().read(Path::new(&param)) {
                Ok(content) => {
                    session.data(channel, content.to_vec().into());
                }
                Err(e) => {
                    self.status = 1;
                    // TODO: stderr
                    eprintln!("{e}");
                    session.data(channel, format!("cat: {param}: {e}").into());
                }
            }
        }

        CommandResult::Exit(self.status)
    }
}

#[async_trait]
impl Command for Cat {
    async fn new<S: ThrusshSession + Send>(
        connection: &mut ConnectionState,
        params: &[String],
        channel: ChannelId,
        session: &mut S,
    ) -> CommandResult<Self> {
        let this = Self {
            remaining_params: params.to_vec().into(),
            status: 0,
        };

        if params.is_empty() {
            CommandResult::ReadStdin(this)
        } else {
            this.run(connection, channel, session)
        }
    }

    async fn stdin<S: ThrusshSession + Send>(
        self,
        connection: &mut ConnectionState,
        channel: ChannelId,
        data: &[u8],
        session: &mut S,
    ) -> CommandResult<Self> {
        session.data(channel, data.to_vec().into());
        self.run(connection, channel, session)
    }
}

#[cfg(test)]
mod test {
    use std::path::Path;

    use mockall::predicate::{self, always};

    use crate::{
        command::{cat::Cat, Command, CommandResult},
        server::{
            test::{fake_channel_id, predicate::eq_string},
            ConnectionState, MockThrusshSession,
        },
    };

    #[tokio::test]
    async fn no_args() {
        let mut session = MockThrusshSession::default();

        let out = Cat::new(
            &mut ConnectionState::mock(),
            [].as_slice(),
            fake_channel_id(),
            &mut session,
        )
        .await;

        assert!(matches!(out, CommandResult::ReadStdin(_)), "{out:?}");
    }

    #[tokio::test]
    async fn file_args_with_missing() {
        let mut session = MockThrusshSession::default();
        let mut state = ConnectionState::mock();

        state.file_system().mkdirall(Path::new("/rootdir")).unwrap();

        state
            .file_system()
            .write(Path::new("a"), "hello".as_bytes().into())
            .unwrap();
        state
            .file_system()
            .write(Path::new("/rootdir/c"), "world".as_bytes().into())
            .unwrap();

        session
            .expect_data()
            .once()
            .with(always(), eq_string("hello"))
            .returning(|_, _| ());

        session
            .expect_data()
            .once()
            .with(always(), eq_string("cat: b: No such file or directory"))
            .returning(|_, _| ());

        session
            .expect_data()
            .once()
            .with(always(), eq_string("world"))
            .returning(|_, _| ());

        let out = Cat::new(
            &mut state,
            ["a".to_string(), "b".to_string(), "/rootdir/c".to_string()].as_slice(),
            fake_channel_id(),
            &mut session,
        )
        .await;

        assert!(matches!(out, CommandResult::Exit(1)), "{out:?}");
    }

    #[tokio::test]
    async fn file_args() {
        let mut session = MockThrusshSession::default();
        let mut state = ConnectionState::mock();

        state
            .file_system()
            .write(Path::new("a"), "hello".as_bytes().into())
            .unwrap();
        state
            .file_system()
            .write(Path::new("b"), "world".as_bytes().into())
            .unwrap();

        session
            .expect_data()
            .once()
            .with(always(), eq_string("hello"))
            .returning(|_, _| ());

        session
            .expect_data()
            .once()
            .with(always(), eq_string("world"))
            .returning(|_, _| ());

        let out = Cat::new(
            &mut state,
            ["a".to_string(), "b".to_string()].as_slice(),
            fake_channel_id(),
            &mut session,
        )
        .await;

        assert!(matches!(out, CommandResult::Exit(0)), "{out:?}");
    }

    #[tokio::test]
    async fn stdin() {
        let mut session = MockThrusshSession::default();
        let mut state = ConnectionState::mock();

        state
            .file_system()
            .write(Path::new("a"), "hello".as_bytes().into())
            .unwrap();

        state
            .file_system()
            .write(Path::new("b"), "world".as_bytes().into())
            .unwrap();

        session
            .expect_data()
            .once()
            .with(always(), eq_string("hello"))
            .returning(|_, _| ());

        session
            .expect_data()
            .once()
            .with(always(), eq_string("the whole"))
            .returning(|_, _| ());

        session
            .expect_data()
            .once()
            .with(always(), eq_string("world"))
            .returning(|_, _| ());

        let out = Cat::new(
            &mut state,
            ["a".to_string(), "-".to_string(), "b".to_string()].as_slice(),
            fake_channel_id(),
            &mut session,
        )
        .await
        .unwrap_stdin();

        let out = out
            .stdin(
                &mut state,
                fake_channel_id(),
                "the whole".as_bytes(),
                &mut session,
            )
            .await;

        assert!(matches!(out, CommandResult::Exit(0)), "{out:?}");
    }
}
diff --git a/pisshoff-server/src/command/ls.rs b/pisshoff-server/src/command/ls.rs
index 81278bd..75e4e87 100644
--- a/pisshoff-server/src/command/ls.rs
+++ b/pisshoff-server/src/command/ls.rs
@@ -1,10 +1,12 @@
use std::{fmt::Write, path::Path};

use async_trait::async_trait;
use thrussh::ChannelId;

use crate::{
    command::{Command, CommandResult},
    server::{ConnectionState, ThrusshSession},
};
use async_trait::async_trait;
use std::fmt::Write;
use thrussh::ChannelId;

#[derive(Debug, Clone)]
pub struct Ls {}
@@ -17,33 +19,55 @@ impl Command for Ls {
        channel: ChannelId,
        session: &mut S,
    ) -> CommandResult<Self> {
        let mut error = false;

        let resp = if params.is_empty() {
            connection.file_system().ls(None).join("  ")
            match connection.file_system().ls(None) {
                Ok(v) => v.join("  "),
                Err(e) => {
                    error = true;
                    format!("ls: {}: {e}", connection.file_system().pwd().display())
                }
            }
        } else if params.len() == 1 {
            connection
            match connection
                .file_system()
                .ls(Some(params.first().unwrap()))
                .join("  ")
                .ls(Some(Path::new(params.first().unwrap())))
            {
                Ok(v) => v.join("  "),
                Err(e) => {
                    error = true;
                    format!("ls: {}: {e}", params.first().unwrap())
                }
            }
        } else {
            let mut out = String::new();

            for dir in params {
                if !out.is_empty() {
                    out.push_str("\n\n");
                    out.push('\n');
                }

                write!(out, "{dir}:").unwrap();
                out.push_str(&connection.file_system().ls(Some(dir)).join("  "));
                match connection.file_system().ls(Some(Path::new(dir))) {
                    Ok(v) => {
                        write!(out, "{dir}:\n{}", v.join("  ")).unwrap();
                    }
                    Err(e) => {
                        error = true;
                        write!(out, "ls: {dir}: {e}").unwrap();
                    }
                }
            }

            out
        };

        if !resp.is_empty() {
            let resp = resp.trim();
            session.data(channel, format!("{resp}\n").into());
        }

        CommandResult::Exit(0)
        CommandResult::Exit(u32::from(error))
    }

    async fn stdin<S: ThrusshSession + Send>(
@@ -59,6 +83,10 @@ impl Command for Ls {

#[cfg(test)]
mod test {
    use std::path::Path;

    use mockall::predicate::always;

    use crate::{
        command::{ls::Ls, Command, CommandResult},
        server::{
@@ -66,7 +94,6 @@ mod test {
            ConnectionState, MockThrusshSession,
        },
    };
    use mockall::predicate::always;

    #[tokio::test]
    async fn empty_pwd() {
@@ -93,8 +120,12 @@ mod test {
            .with(always(), eq_string("a:\n\nb:\n"))
            .returning(|_, _| ());

        let mut state = ConnectionState::mock();
        state.file_system().mkdirall(Path::new("/root/a")).unwrap();
        state.file_system().mkdirall(Path::new("/root/b")).unwrap();

        let out = Ls::new(
            &mut ConnectionState::mock(),
            &mut state,
            ["a".to_string(), "b".to_string()].as_slice(),
            fake_channel_id(),
            &mut session,
diff --git a/pisshoff-server/src/file_system.rs b/pisshoff-server/src/file_system.rs
index 5f96385..3882137 100644
--- a/pisshoff-server/src/file_system.rs
+++ b/pisshoff-server/src/file_system.rs
@@ -1,25 +1,57 @@
#![allow(dead_code)]

use std::path::{Path, PathBuf};
use std::{
    borrow::Cow,
    collections::{btree_map::Entry, BTreeMap},
    fmt::{Display, Formatter},
    path::{Path, PathBuf},
};

/// A fake file system, stored in memory only active for the current session.
pub struct FileSystem {
    pwd: PathBuf,
    home: PathBuf,
    data: Tree,
}

pub enum Tree {
    Directory(BTreeMap<String, Box<Tree>>),
    File(Box<[u8]>),
}

impl FileSystem {
    pub fn new(user: &str) -> Self {
        let pwd = if user == "root" {
            PathBuf::new().join("/root")
            PathBuf::from("/root")
        } else {
            PathBuf::new().join("/home").join(user)
            PathBuf::from("/home").join(user)
        };

        Self {
        let mut this = Self {
            home: pwd.clone(),
            pwd,
            data: Tree::Directory(BTreeMap::new()),
        };

        let _res = this.mkdirall(&this.pwd.clone());
        this
    }

    pub fn mkdirall(&mut self, path: &Path) -> Result<(), LsError> {
        let mut tree = &mut self.data;

        for c in path {
            match tree {
                Tree::Directory(d) => {
                    tree = d
                        .entry(c.to_str().unwrap().to_string())
                        .or_insert_with(|| Box::new(Tree::Directory(BTreeMap::new())));
                }
                Tree::File(_) => return Err(LsError::FileExists),
            }
        }

        Ok(())
    }

    pub fn cd(&mut self, v: Option<&str>) {
@@ -34,8 +66,120 @@ impl FileSystem {
        &self.pwd
    }

    pub fn read(&self, path: &Path) -> Result<&[u8], LsError> {
        let canonical = self.pwd().join(path);
        let mut tree = &self.data;

        for c in &canonical {
            match tree {
                Tree::Directory(d) => {
                    tree = d
                        .get(c.to_str().unwrap())
                        .ok_or(LsError::NoSuchFileOrDirectory)?;
                }
                Tree::File(_) => {
                    return Err(LsError::NotDirectory);
                }
            }
        }

        match tree {
            Tree::Directory(_) => Err(LsError::IsADirectory),
            Tree::File(content) => Ok(content),
        }
    }

    pub fn write(&mut self, path: &Path, content: Box<[u8]>) -> Result<(), LsError> {
        let canonical = self.pwd().join(path);
        let mut tree = &mut self.data;

        if let Some(parents) = canonical.parent() {
            for c in parents {
                match tree {
                    Tree::Directory(d) => {
                        tree = d
                            .get_mut(c.to_str().unwrap())
                            .ok_or(LsError::NoSuchFileOrDirectory)?;
                    }
                    Tree::File(_) => {
                        return Err(LsError::NotDirectory);
                    }
                }
            }
        }

        match tree {
            Tree::Directory(v) => {
                match v.entry(
                    canonical
                        .components()
                        .next_back()
                        .unwrap()
                        .as_os_str()
                        .to_str()
                        .unwrap()
                        .to_string(),
                ) {
                    Entry::Vacant(v) => {
                        v.insert(Box::new(Tree::File(content)));
                        Ok(())
                    }
                    Entry::Occupied(mut o) if matches!(o.get().as_ref(), Tree::File(_)) => {
                        o.insert(Box::new(Tree::File(content)));
                        Ok(())
                    }
                    Entry::Occupied(_) => Err(LsError::IsADirectory),
                }
            }
            Tree::File(_) => Err(LsError::NotDirectory),
        }
    }

    #[allow(clippy::unused_self)]
    pub fn ls(&self, _dir: Option<&str>) -> &[&str] {
        &[]
    pub fn ls<'a>(&'a self, dir: Option<&'a Path>) -> Result<Vec<&'a str>, LsError> {
        let canonical = if let Some(dir) = dir {
            Cow::Owned(self.pwd().join(dir))
        } else {
            Cow::Borrowed(self.pwd())
        };

        let mut tree = &self.data;

        for c in canonical.as_ref() {
            match tree {
                Tree::Directory(d) => {
                    tree = d
                        .get(c.to_str().unwrap())
                        .ok_or(LsError::NoSuchFileOrDirectory)?;
                }
                Tree::File(_) => {
                    return Err(LsError::NotDirectory);
                }
            }
        }

        match tree {
            Tree::Directory(v) => Ok(v.keys().map(String::as_str).collect()),
            Tree::File(_) => Ok(vec![dir.unwrap_or(self.pwd()).to_str().unwrap()]),
        }
    }
}

#[derive(Debug)]
pub enum LsError {
    NotDirectory,
    NoSuchFileOrDirectory,
    IsADirectory,
    FileExists,
}

impl Display for LsError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            LsError::NoSuchFileOrDirectory => "No such file or directory",
            LsError::NotDirectory => "Not a directory",
            LsError::IsADirectory => "Is a directory",
            LsError::FileExists => "File exists",
        })
    }
}