From 419d1e582de9ad7be4f3b7ec0e1318ae6c2d5e8c Mon Sep 17 00:00:00 2001 From: Astro Date: Sun, 25 Sep 2022 20:06:38 +0200 Subject: [PATCH] split into multiple rust modules --- src/handler.rs | 194 +++++++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 185 +--------------------------------------------- src/server.rs | 29 ++++++++ 3 files changed, 226 insertions(+), 182 deletions(-) create mode 100644 src/handler.rs create mode 100644 src/server.rs diff --git a/src/handler.rs b/src/handler.rs new file mode 100644 index 0000000..ccbc323 --- /dev/null +++ b/src/handler.rs @@ -0,0 +1,194 @@ +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, + buffer: Vec, + user: String, +} + +impl Handler { + pub fn new(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(&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> + Send>>; + type FutureUnit = + Pin> + Send>>; + type FutureBool = + Pin> + 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) + // } +} diff --git a/src/main.rs b/src/main.rs index 88c21b8..c035ff6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,15 +1,8 @@ -use std::fs::{OpenOptions, File, create_dir_all}; -use std::io::Write; -use std::path::Path; -use std::pin::Pin; use std::str::FromStr; use std::sync::Arc; -use chrono::Utc; -use futures::FutureExt; -use log::info; -use russh::server::{Auth, Session}; -use russh::*; +mod server; +mod handler; #[tokio::main] async fn main() { @@ -31,7 +24,7 @@ async fn main() { config.auth_rejection_time = std::time::Duration::from_secs(1); config.keys.push(russh_keys::key::KeyPair::generate_ed25519().unwrap()); let config = Arc::new(config); - let sh = Server; + let sh = server::Server; russh::server::run( config, &std::net::SocketAddr::from_str(&listen_addr).unwrap(), @@ -40,175 +33,3 @@ async fn main() { .await .unwrap(); } - -struct Server; - -impl server::Server for Server { - type Handler = Handler; - fn new_client(&mut self, addr: Option) -> Handler { - let addr = if let Some(addr) = addr { - format!("{}", addr) - } else { - format!("unknown") - }; - info!("Connection from {}", addr); - - let path = &format!( - "{}-{}.txt", - Utc::now().format("%Y-%m-%d/%H:%M:%S"), - addr - ); - let path = Path::new(path); - create_dir_all(path.parent().unwrap()).unwrap(); - Handler { - filename: path.display().to_string(), - lazy_file: None, - buffer: vec![], - user: "root".into(), - } - } -} - -fn send_str(session: &mut Session, channel: ChannelId, s: String) { - let data = CryptoVec::from(s); - session.data(channel, data); -} - -struct Handler { - filename: String, - lazy_file: Option, - buffer: Vec, - user: String, -} - -impl Handler { - 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(&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> + Send>>; - type FutureUnit = - Pin> + Send>>; - type FutureBool = - Pin> + 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 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) - } -} diff --git a/src/server.rs b/src/server.rs new file mode 100644 index 0000000..494c994 --- /dev/null +++ b/src/server.rs @@ -0,0 +1,29 @@ +use std::fs::create_dir_all; +use std::path::Path; + +use chrono::Utc; +use log::info; +use russh::*; + +pub struct Server; + +impl server::Server for Server { + type Handler = crate::handler::Handler; + fn new_client(&mut self, addr: Option) -> Self::Handler { + let addr = if let Some(addr) = addr { + format!("{}", addr) + } else { + format!("unknown") + }; + info!("Connection from {}", addr); + + let path = &format!( + "{}-{}.txt", + Utc::now().format("%Y-%m-%d/%H:%M:%S"), + addr + ); + let path = Path::new(path); + create_dir_all(path.parent().unwrap()).unwrap(); + crate::handler::Handler::new(path.display()) + } +}