sshlogd/src/handler.rs

use std::fs::{OpenOptions, File};
use std::io::Write;
use std::pin::Pin;

use futures::FutureExt;
use log::info;
use russh::server::{Auth, Session};
use russh::*;

fn send_str(session: &mut Session, channel: ChannelId, s: String) {
    let data = CryptoVec::from(s);
    session.data(channel, data);
}

pub struct Handler {
    filename: String,
    lazy_file: Option<File>,
    buffer: Vec<u8>,
    user: String,
}

impl Handler {
    pub fn new<P: ToString>(path: P) -> Self {
        Handler {
            filename: path.to_string(),
            lazy_file: None,
            buffer: vec![],
            user: "root".into(),
        }
    }

    fn file(&mut self) -> &mut File {
        if self.lazy_file.is_none() {
            let file = OpenOptions::new()
                .create(true)
                .append(true)
                .open(&self.filename)
                .unwrap();
            self.lazy_file = Some(file);
        }
        self.lazy_file.as_mut().unwrap()
    }

    fn send_prompt(&self, session: &mut Session, channel: ChannelId) {
        send_str(session, channel, format!("{}@fnordister:~$ ", self.user));
    }

    fn handle_command<F: FnMut(&str)>(&self, command: String, mut respond: F) {
        let program_len = command.find(|c: char| c.is_whitespace())
            .unwrap_or(command.len());
        let program = &command[..program_len];

        match program {
            "whoami" => respond(&self.user),
            "id" => respond(&format!("uid=0({}) gid=0(root)", self.user)),
            "uname" => respond("Linux fnordister 5.10.0-10-amd64 #1 SMP Debian 5.10.84-1 (2021-12-08) x86_64 GNU/Linux"),
            "pwd" => respond("/"),
            "ls" => {
                respond("drwxr-xr-x  18 root root    18 Jan  4  1969 .");
                respond("drwxr-xr-x  18 root root    18 Jan  4  1969 ..");
                respond("drwxr-xr-x  18 root root    18 Jan  4  0000 ...");
            },
            "bash" => {}
            "sh" => {}
            "cd" => {}
            "" => {}
            _ => respond(&format!("{}: command not found", program)),
        }
    }
}

impl server::Handler for Handler {
    type Error = anyhow::Error;
    type FutureAuth =
        Pin<Box<dyn core::future::Future<Output = anyhow::Result<(Self, Auth)>> + Send>>;
    type FutureUnit =
        Pin<Box<dyn core::future::Future<Output = anyhow::Result<(Self, Session)>> + Send>>;
    type FutureBool =
        Pin<Box<dyn core::future::Future<Output = anyhow::Result<(Self, Session, bool)>> + Send>>;

    fn finished_auth(self, auth: Auth) -> Self::FutureAuth {
        async { Ok((self, auth)) }.boxed()
    }

    fn finished_bool(self, b: bool, s: Session) -> Self::FutureBool {
        async move { Ok((self, s, b)) }.boxed()
    }

    fn finished(self, s: Session) -> Self::FutureUnit {
        async { Ok((self, s)) }.boxed()
    }

    fn shell_request(self, channel: ChannelId, mut session: Session) -> Self::FutureUnit {
        info!("shell_request");
        self.send_prompt(&mut session, channel);

        self.finished(session)
    }

    fn exec_request(mut self, channel: ChannelId, data: &[u8], mut session: Session) -> Self::FutureUnit {
        info!("exec_request");
        writeln!(self.file(), "Execute: {}\n", String::from_utf8_lossy(data))
            .unwrap();

        let line = String::from_utf8_lossy(data).into();
        self.handle_command(line, |response| {
            let mut data = Vec::from(response);
            data.extend_from_slice(b"\r\n");
            session.data(channel, data.into());
        });
        session.close(channel);

        self.finished(session)
    }

    fn subsystem_request(
        mut self,
        _channel: ChannelId,
        name: &str,
        session: Session
    ) -> Self::FutureUnit {
        info!("subsystem_request");
        writeln!(self.file(), "Subsystem requested: {}\n", name)
            .unwrap();

        self.finished(session)
    }

    // fn channel_open_session(
    //     self,
    //     channel: ChannelId,
    //     mut session: Session
    // ) -> Self::FutureBool {
    //     info!("channel_open_session");
    //     session.channel_success(channel);
    //     self.finished_bool(true, session)
    // }

    // fn channel_open_confirmation(
    //     self,
    //     id: ChannelId,
    //     max_packet_size: u32,
    //     window_size: u32,
    //     mut session: Session
    // ) -> Self::FutureUnit {
    //     info!("channel_open_confirmation");

    //     self.finished(session)
    // }

    fn auth_password(mut self, user: &str, password: &str) -> Self::FutureAuth {
        writeln!(self.file(), "Authenticated as {} with {}\n", user, password)
            .unwrap();
        self.user = user.into();
        self.finished_auth(server::Auth::Accept)
    }

    fn data(mut self, channel: ChannelId, data: &[u8], mut session: Session) -> Self::FutureUnit {
        self.file().write(data)
            .unwrap();
        // echo input back
        session.data(channel, data.to_vec().into());

        if self.buffer.len() < 1024 {
            self.buffer.extend_from_slice(data);
        }
        if let Some(newline) = self.buffer.iter().position(|b| *b == b'\r') {
            let rest = self.buffer.split_off(newline + 1);
            let line: String = String::from_utf8_lossy(&self.buffer).into();
            self.buffer = rest;

            session.data(channel, b"\n".to_vec().into());
            self.handle_command(line, |response| {
                let mut data = Vec::from(response);
                data.extend_from_slice(b"\r\n");
                session.data(channel, data.into());
            });
            self.send_prompt(&mut session, channel);
        }

        self.finished(session)
    }

    // fn extended_data(
    //     self,
    //     channel: ChannelId,
    //     code: u32,
    //     data: &[u8],
    //     session: Session
    // ) -> Self::FutureUnit {
    //     info!("extended_data {} {} {:?}", channel, code, data);
    //     self.finished(session)
    // }
}