Small refactor

This commit is contained in:
Dreaded_X 2025-04-10 17:32:12 +02:00
parent 8a628d51a6
commit c637bd7427
Signed by: Dreaded_X
GPG Key ID: 5A0CBFE3C3377FAA
8 changed files with 123 additions and 102 deletions

23
src/cli.rs Normal file
View File

@ -0,0 +1,23 @@
use clap::Parser;
/// Quickly create http tunnels for development
#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
pub struct Args {
/// Make all tunnels public by default instead of private
#[arg(long, group = "access")]
public: bool,
#[arg(long, group = "access")]
protected: bool,
}
impl Args {
pub fn make_public(&self) -> bool {
self.public
}
pub fn make_protected(&self) -> bool {
self.protected
}
}

View File

@ -1,19 +1,18 @@
use std::{io::Write, iter::once, net::SocketAddr, sync::Arc, time::Duration};
use std::{io::Write, iter::once};
use clap::Parser;
use clap::Parser as _;
use indexmap::IndexMap;
use ratatui::{Terminal, TerminalOptions, Viewport, layout::Rect, prelude::CrosstermBackend};
use russh::{
ChannelId,
keys::PrivateKey,
server::{Auth, Msg, Server as _, Session},
server::{Auth, Msg, Session},
};
use tokio::net::ToSocketAddrs;
use tracing::{debug, trace, warn};
use crate::{
keys::Input,
terminal::TerminalHandle,
cli,
input::Input,
io::TerminalHandle,
tui::Renderer,
tunnel::{Tunnel, TunnelAccess, Tunnels},
};
@ -31,7 +30,7 @@ pub struct Handler {
}
impl Handler {
fn new(all_tunnels: Tunnels) -> Self {
pub fn new(all_tunnels: Tunnels) -> Self {
Self {
all_tunnels,
tunnels: IndexMap::new(),
@ -52,20 +51,19 @@ impl Handler {
}
async fn resize(&mut self, width: u32, height: u32) -> std::io::Result<()> {
let rect = Rect {
x: 0,
y: 0,
width: width as u16,
height: height as u16,
};
if let Some(terminal) = &mut self.terminal {
terminal.resize(rect)?;
} else {
todo!()
}
let rect = Rect {
x: 0,
y: 0,
width: width as u16,
height: height as u16,
};
self.redraw().await?;
terminal.resize(rect)?;
self.redraw().await?;
} else {
warn!("Resize called without valid terminal");
}
Ok(())
}
@ -81,13 +79,12 @@ impl Handler {
async fn redraw(&mut self) -> std::io::Result<()> {
if let Some(terminal) = &mut self.terminal {
trace!("redraw");
self.renderer.update_table(&self.tunnels).await;
self.renderer.select(self.selected);
self.renderer.update(&self.tunnels, self.selected).await;
terminal.draw(|frame| {
self.renderer.render(frame);
})?;
} else {
todo!()
warn!("Redraw called without valid terminal");
}
Ok(())
@ -162,18 +159,6 @@ impl Handler {
}
}
/// Quickly create http tunnels for development
#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
/// Make all tunnels public by default instead of private
#[arg(long, group = "access")]
public: bool,
#[arg(long, group = "access")]
protected: bool,
}
impl russh::server::Handler for Handler {
type Error = russh::Error;
@ -239,15 +224,14 @@ impl russh::server::Handler for Handler {
trace!(?cmd, "exec_request");
let cmd = once("<ssh command> --").chain(cmd.split_whitespace());
match Args::try_parse_from(cmd) {
match cli::Args::try_parse_from(cmd) {
Ok(args) => {
debug!("{args:?}");
if args.public {
if args.make_public() {
trace!("Making tunnels public");
self.set_access_all(TunnelAccess::Public).await;
self.redraw().await?;
}
if args.protected {
} else if args.make_protected() {
trace!("Making tunnels protected");
self.set_access_all(TunnelAccess::Protected).await;
self.redraw().await?;
@ -346,52 +330,3 @@ impl Drop for Handler {
});
}
}
pub struct Server {
tunnels: Tunnels,
}
impl Server {
pub fn new(tunnels: Tunnels) -> Self {
Server { tunnels }
}
pub fn tunnels(&self) -> Tunnels {
self.tunnels.clone()
}
pub fn run(
&mut self,
key: PrivateKey,
addr: impl ToSocketAddrs + Send + std::fmt::Debug,
) -> impl Future<Output = Result<(), std::io::Error>> + Send {
let config = russh::server::Config {
inactivity_timeout: Some(Duration::from_secs(3600)),
auth_rejection_time: Duration::from_secs(3),
auth_rejection_time_initial: Some(Duration::from_secs(0)),
keys: vec![key],
preferred: russh::Preferred {
..Default::default()
},
nodelay: true,
..Default::default()
};
let config = Arc::new(config);
debug!(?addr, "Running ssh");
async move { self.run_on_address(config, addr).await }
}
}
impl russh::server::Server for Server {
type Handler = Handler;
fn new_client(&mut self, _peer_addr: Option<SocketAddr>) -> Self::Handler {
Handler::new(self.tunnels.clone())
}
fn handle_session_error(&mut self, error: <Self::Handler as russh::server::Handler>::Error) {
warn!("Session error: {error:#?}");
}
}

View File

@ -1,10 +1,15 @@
#![feature(impl_trait_in_fn_trait_return)]
#![feature(let_chains)]
pub mod animals;
pub mod auth;
pub mod helper;
pub mod keys;
pub mod ssh;
pub mod terminal;
pub mod tui;
pub mod tunnel;
mod animals;
mod auth;
mod cli;
mod handler;
mod helper;
mod input;
mod io;
mod server;
mod tui;
mod tunnel;
pub use server::Server;
pub use tunnel::{Tunnel, Tunnels};

View File

@ -8,7 +8,7 @@ use rand::rngs::OsRng;
use tokio::net::TcpListener;
use tracing::{error, info};
use tracing_subscriber::{EnvFilter, Registry, layer::SubscriberExt, util::SubscriberInitExt};
use tunnel_rs::{ssh::Server, tunnel::Tunnels};
use tunnel_rs::{Server, Tunnels};
#[tokio::main]
async fn main() -> color_eyre::Result<()> {

56
src/server.rs Normal file
View File

@ -0,0 +1,56 @@
use std::{net::SocketAddr, sync::Arc, time::Duration};
use russh::{keys::PrivateKey, server::Server as _};
use tokio::net::ToSocketAddrs;
use tracing::{debug, warn};
use crate::{handler::Handler, tunnel::Tunnels};
pub struct Server {
tunnels: Tunnels,
}
impl Server {
pub fn new(tunnels: Tunnels) -> Self {
Server { tunnels }
}
pub fn tunnels(&self) -> Tunnels {
self.tunnels.clone()
}
pub fn run(
&mut self,
key: PrivateKey,
addr: impl ToSocketAddrs + Send + std::fmt::Debug,
) -> impl Future<Output = Result<(), std::io::Error>> + Send {
let config = russh::server::Config {
inactivity_timeout: Some(Duration::from_secs(3600)),
auth_rejection_time: Duration::from_secs(3),
auth_rejection_time_initial: Some(Duration::from_secs(0)),
keys: vec![key],
preferred: russh::Preferred {
..Default::default()
},
nodelay: true,
..Default::default()
};
let config = Arc::new(config);
debug!(?addr, "Running ssh");
async move { self.run_on_address(config, addr).await }
}
}
impl russh::server::Server for Server {
type Handler = Handler;
fn new_client(&mut self, _peer_addr: Option<SocketAddr>) -> Self::Handler {
Handler::new(self.tunnels.clone())
}
fn handle_session_error(&mut self, error: <Self::Handler as russh::server::Handler>::Error) {
warn!("Session error: {error:#?}");
}
}

View File

@ -36,13 +36,19 @@ impl Default for Renderer {
impl Renderer {
// NOTE: This needs to be a separate function as the render functions can not be async
pub async fn update_table(&mut self, tunnels: &IndexMap<String, Option<Tunnel>>) {
pub async fn update(
&mut self,
tunnels: &IndexMap<String, Option<Tunnel>>,
index: Option<usize>,
) {
self.table_rows = futures::stream::iter(tunnels.iter())
.then(tunnel::tui::to_row)
.collect::<Vec<_>>()
.await;
self.update_widths();
self.table_state.select(index);
}
pub fn render(&mut self, frame: &mut Frame) {
@ -117,8 +123,4 @@ impl Renderer {
frame.render_stateful_widget(t, rect, &mut self.table_state);
}
pub fn select(&mut self, index: Option<usize>) {
self.table_state.select(index);
}
}