split into multiple rust modules

This commit is contained in:
Astro 2022-09-25 20:06:38 +02:00
parent aaf3c19530
commit 419d1e582d
3 changed files with 226 additions and 182 deletions

194
src/handler.rs Normal file
View File

@ -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<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)
// }
}

View File

@ -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<std::net::SocketAddr>) -> 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<File>,
buffer: Vec<u8>,
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<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 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)
}
}

29
src/server.rs Normal file
View File

@ -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<std::net::SocketAddr>) -> 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())
}
}