use super::{client_listen, config::ServerConfig}; use crate::message::PortDesc; use anyhow::Result; use crossterm::{ cursor::{MoveTo, RestorePosition, SavePosition}, event::{Event, EventStream, KeyCode, KeyEvent, KeyModifiers}, execute, queue, style::{Color, PrintStyledContent, Stylize}, terminal::{ disable_raw_mode, enable_raw_mode, size, Clear, ClearType, DisableLineWrap, EnableLineWrap, EnterAlternateScreen, LeaveAlternateScreen, }, }; use log::{error, info, Level, Metadata, Record}; use open; use std::collections::vec_deque::VecDeque; use std::collections::{HashMap, HashSet}; use std::io::{stdout, Write}; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio_stream::StreamExt; pub enum UIEvent { Connected(u16), Disconnected, ServerLine(String), LogLine(log::Level, String), Ports(Vec), } #[derive(Debug, Clone)] pub struct Logger { line_sender: mpsc::Sender, } impl Logger { pub fn new(line_sender: mpsc::Sender) -> Box { Box::new(Logger { line_sender }) } } impl log::Log for Logger { fn enabled(&self, metadata: &Metadata) -> bool { metadata.level() <= Level::Info } fn log(&self, record: &Record) { if self.enabled(record.metadata()) { let line = format!("{}", record.args()); _ = self .line_sender .try_send(UIEvent::LogLine(record.level(), line)); } } fn flush(&self) {} } #[derive(Debug)] struct Listener { enabled: bool, stop: Option>, desc: Option, } impl Listener { pub fn from_desc( socks_port: Option, desc: PortDesc, enabled: bool, ) -> Listener { let mut listener = Listener { enabled, stop: None, desc: Some(desc) }; if enabled { listener.start(socks_port); } listener } pub fn enabled(&self) -> bool { self.enabled } pub fn set_enabled(&mut self, socks_port: Option, enabled: bool) { if enabled { self.enabled = true; self.start(socks_port); } else { self.enabled = false; self.stop = None; } } pub fn connect(&mut self, socks_port: Option, desc: PortDesc) { self.desc = Some(desc); self.start(socks_port); } pub fn disconnect(&mut self) { self.desc = None; self.stop = None; } pub fn start(&mut self, socks_port: Option) { if self.enabled { if let (Some(desc), Some(socks_port), None) = (&self.desc, socks_port, &self.stop) { info!("Starting port {port} to {socks_port}", port = desc.port); let (l, stop) = oneshot::channel(); let port = desc.port; tokio::spawn(async move { let result = tokio::select! { r = client_listen(port, socks_port) => r, _ = stop => Ok(()), }; if let Err(e) = result { error!("Error listening on port {port}: {e:?}"); } else { info!("Stopped listening on port {port}"); } }); self.stop = Some(l); } } } } #[derive(Debug)] pub struct UI { events: mpsc::Receiver, ports: HashMap, socks_port: Option, lines: VecDeque, config: ServerConfig, selection: usize, running: bool, show_logs: bool, alternate_screen: bool, raw_mode: bool, } impl UI { pub fn new(events: mpsc::Receiver, config: ServerConfig) -> UI { UI { events, ports: HashMap::new(), socks_port: None, running: true, show_logs: false, selection: 0, lines: VecDeque::with_capacity(1024), config, alternate_screen: false, raw_mode: false, } } pub async fn run(&mut self) -> Result<()> { let mut console_events = EventStream::new(); self.running = true; while self.running { self.handle_events(&mut console_events).await; if self.connected() { self.render_connected()?; } else { self.render_disconnected()?; } } Ok(()) } fn render_disconnected(&mut self) -> Result<()> { self.enter_alternate_screen()?; self.disable_raw_mode()?; let mut stdout = stdout(); let (columns, _) = size()?; let columns: usize = columns.into(); execute!( stdout, SavePosition, MoveTo(0, 0), PrintStyledContent( format!("{:^columns$}\r\n", "Not Connected") .with(Color::Black) .on(Color::Red) ), RestorePosition, )?; Ok(()) } fn render_connected(&mut self) -> Result<()> { self.enter_alternate_screen()?; self.enable_raw_mode()?; let mut stdout = stdout(); let (columns, rows) = size()?; let columns: usize = columns.into(); queue!(stdout, Clear(ClearType::All), MoveTo(0, 0))?; // List of open ports // How wide are all the things? let padding = 1; let enabled_width = 1; // Just 1 character let port_width = 5; // 5 characters for 16-bit number let description_width = columns - (padding + padding + enabled_width + padding + port_width + padding); print!( "{}", format!( " {enabled:>enabled_width$} {port:>port_width$} {description: &port_desc.desc, None => "", }; print!("{caret} {state:>enabled_width$} {port:port_width$} {desc:description_width$}\r\n"); } // Log if self.show_logs { let hr: usize = ((rows / 2) - 2).into(); let start: usize = if self.lines.len() > hr { self.lines.len() - hr } else { 0 }; queue!(stdout, MoveTo(0, rows / 2))?; print!("{}", format!("{:columns$}", " Log").negative()); for line in self.lines.range(start..) { print!("{}\r\n", line); } } queue!(stdout, MoveTo(0, rows - 1))?; print!( "{}", format!( "{:columns$}", " q - quit | l - toggle log | \u{2191}/\u{2193} - select port | e - enable/disable | - browse" ).negative() ); stdout.flush()?; Ok(()) } fn connected(&self) -> bool { match self.socks_port { Some(_) => true, None => false, } } fn get_ui_ports(&self) -> Vec { let mut ports: Vec = self.ports.keys().copied().collect(); ports.sort(); ports } fn get_selected_port(&self) -> Option { if self.selection < self.ports.len() { Some(self.get_ui_ports()[self.selection]) } else { None } } fn enable_disable_port(&mut self, port: u16) { if let Some(listener) = self.ports.get_mut(&port) { listener.set_enabled(self.socks_port, !listener.enabled()); } } fn enable_raw_mode(&mut self) -> Result<()> { if !self.raw_mode { enable_raw_mode()?; self.raw_mode = true; } Ok(()) } fn disable_raw_mode(&mut self) -> Result<()> { if self.raw_mode { disable_raw_mode()?; self.raw_mode = false; } Ok(()) } fn enter_alternate_screen(&mut self) -> Result<()> { if !self.alternate_screen { enable_raw_mode()?; execute!(stdout(), EnterAlternateScreen, DisableLineWrap)?; self.alternate_screen = true; } Ok(()) } fn leave_alternate_screen(&mut self) -> Result<()> { if self.alternate_screen { execute!(stdout(), LeaveAlternateScreen, EnableLineWrap)?; disable_raw_mode()?; self.alternate_screen = false; } Ok(()) } async fn handle_events(&mut self, console_events: &mut EventStream) { tokio::select! { ev = console_events.next() => self.handle_console_event(ev), ev = self.events.recv() => self.handle_internal_event(ev), } } fn handle_console_event( &mut self, ev: Option>, ) { match ev { Some(Ok(Event::Key(ev))) => match ev { KeyEvent { code: KeyCode::Char('c'), .. } => { if ev.modifiers.intersects(KeyModifiers::CONTROL) { self.running = false; } } KeyEvent { code: KeyCode::Esc, .. } | KeyEvent { code: KeyCode::Char('q'), .. } => { self.running = false; } KeyEvent { code: KeyCode::Char('l'), .. } => { self.show_logs = !self.show_logs; } KeyEvent { code: KeyCode::Char('e'), .. } => { if let Some(p) = self.get_selected_port() { self.enable_disable_port(p); } } KeyEvent { code: KeyCode::Up, .. } | KeyEvent { code: KeyCode::Char('j'), .. } => { if self.selection > 0 { self.selection -= 1; } } KeyEvent { code: KeyCode::Down, .. } | KeyEvent { code: KeyCode::Char('k'), .. } => { if self.selection != self.ports.len() - 1 { self.selection += 1; } } KeyEvent { code: KeyCode::Enter, .. } => { if let Some(p) = self.get_selected_port() { _ = open::that(format!("http://127.0.0.1:{}/", p)); } } _ => (), }, Some(Ok(_)) => (), // Don't care about this event... Some(Err(_)) => (), // Hmmmmmm.....? None => (), // ....no events? what? } } fn handle_internal_event(&mut self, event: Option) { match event { Some(UIEvent::Disconnected) => { self.socks_port = None; for port in self.ports.values_mut() { port.disconnect(); } } Some(UIEvent::Connected(sp)) => { info!("Socks port {sp}"); self.socks_port = Some(sp); for port in self.ports.values_mut() { port.start(self.socks_port); } } Some(UIEvent::Ports(p)) => { let mut leftover_ports: HashSet = HashSet::from_iter(self.ports.keys().copied()); for port_desc in p.into_iter() { leftover_ports.remove(&port_desc.port); if let Some(listener) = self.ports.get_mut(&port_desc.port) { listener.connect(self.socks_port, port_desc); } else { let config = self.config.get(port_desc.port); info!("Port config {port_desc:?} -> {config:?}"); self.ports.insert( port_desc.port, Listener::from_desc( self.socks_port, port_desc, config.enabled, ), ); } } for port in leftover_ports { let mut enabled = false; if let Some(listener) = self.ports.get_mut(&port) { enabled = listener.enabled; listener.disconnect(); } if !enabled { // Just... forget it? Or leave it around forever? self.ports.remove(&port); } } if self.selection >= self.ports.len() && self.ports.len() > 0 { self.selection = self.ports.len() - 1; } } Some(UIEvent::ServerLine(line)) => { while self.lines.len() >= 1024 { self.lines.pop_front(); } self.lines.push_back(format!("[SERVER] {line}")); } Some(UIEvent::LogLine(_level, line)) => { while self.lines.len() >= 1024 { self.lines.pop_front(); } self.lines.push_back(format!("[CLIENT] {line}")); } None => { self.running = false; } } } } impl Drop for UI { fn drop(&mut self) { _ = self.disable_raw_mode(); _ = self.leave_alternate_screen(); } }