sshlogd/src/main.rs
2022-09-19 23:53:20 +02:00

134 lines
3.9 KiB
Rust

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::*;
#[tokio::main]
async fn main() {
env_logger::builder()
.filter_level(log::LevelFilter::Info)
.init();
let listen_addr = std::env::args().skip(1).next()
.expect("Expecting <addr:port>");
let mut config = russh::server::Config::default();
config.connection_timeout = Some(std::time::Duration::from_secs(10000));
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;
russh::server::run(
config,
&std::net::SocketAddr::from_str(&listen_addr).unwrap(),
sh,
)
.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();
let file = OpenOptions::new()
.create(true)
.append(true)
.open(path)
.unwrap();
Handler {
file,
user: "root".into(),
}
}
}
fn send_str(session: &mut Session, channel: ChannelId, s: String) {
let data = CryptoVec::from(s);
session.data(channel, data);
}
struct Handler {
file: File,
user: String,
}
impl Handler {
fn send_prompt(&self, session: &mut Session, channel: ChannelId) {
send_str(session, channel, format!("{}@fnordister:~$ ", self.user));
}
}
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 {
send_str(&mut session, channel, format!("#\r\n# DEMONSTRATION SYSTEM\r\n#\r\n# All input will be published\r\n#\r\n\r\n"));
self.send_prompt(&mut session, channel);
self.finished(session)
}
fn auth_password(mut self, user: &str, password: &str) -> Self::FutureAuth {
writeln!(self.file, "Authenticaed 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());
// pressed return?
if data.contains(&b'\r') {
writeln!(self.file).unwrap();
session.data(channel, b"\n".to_vec().into());
self.send_prompt(&mut session, channel);
}
self.finished(session)
}
}