diff --git a/src/input.rs b/src/input.rs index afa96ab..dccfe2d 100644 --- a/src/input.rs +++ b/src/input.rs @@ -24,7 +24,7 @@ pub enum Command { interactive: bool, /// Block exit signals like CTRL-C - #[clap(short, long)] + #[clap(short, long, conflicts_with = "local-interactive")] block_signals: bool, /// Local interactive diff --git a/src/listener/mod.rs b/src/listener/mod.rs index cae4d64..95188cd 100644 --- a/src/listener/mod.rs +++ b/src/listener/mod.rs @@ -1,5 +1,4 @@ use colored::Colorize; -use rustyline; use rustyline::error::ReadlineError; use std::io::{stdin, stdout, Read, Result, Write}; use std::net::{TcpListener, TcpStream}; @@ -90,11 +89,8 @@ fn listen_tcp_normal(stream: TcpStream, opts: &Opts) -> Result<()> { Ok(()) } -// Listen on given host and port -pub fn listen(opts: &Opts) -> rustyline::Result<()> { - let listener = TcpListener::bind(format!("{}:{}", opts.host, opts.port))?; - - if opts.block_signals { +fn block_signals(should_block: bool) -> Result<()> { + if should_block { #[cfg(unix)] { Signals::new(&[consts::SIGINT])?; @@ -103,24 +99,40 @@ pub fn listen(opts: &Opts) -> rustyline::Result<()> { #[cfg(not(unix))] { print_feature_not_supported(); + + exit(1); + } + } + + Ok(()) +} +// Listen on given host and port +pub fn listen(opts: &Opts) -> rustyline::Result<()> { + let listener = TcpListener::bind(format!("{}:{}", opts.host, opts.port))?; + + #[cfg(not(unix))] + { + if let Mode::Interactive = opts.mode { + print_feature_not_supported(); + + exit(1); } } log::info!("Listening on {}:{}", opts.host.green(), opts.port.cyan()); let (mut stream, _) = listener.accept()?; + match &opts.mode { Mode::Interactive => { + // It exists it if isn't unix above + block_signals(opts.block_signals)?; + #[cfg(unix)] { termios_handler::setup_fd()?; listen_tcp_normal(stream, opts)?; } - - #[cfg(not(unix))] - { - print_feature_not_supported(); - } } Mode::LocalInteractive => { let t = pipe_thread(stream.try_clone()?, stdout()); @@ -132,9 +144,11 @@ pub fn listen(opts: &Opts) -> rustyline::Result<()> { .write_all((command + "\n").as_bytes()) .expect("Failed to send TCP."); })?; + t.join().unwrap(); } Mode::Normal => { + block_signals(opts.block_signals)?; listen_tcp_normal(stream, opts)?; } } diff --git a/src/listener/termios_handler.rs b/src/listener/termios_handler.rs index a5da20c..a8d836e 100644 --- a/src/listener/termios_handler.rs +++ b/src/listener/termios_handler.rs @@ -1,9 +1,10 @@ use std::fs::OpenOptions; +use std::io::Result; use std::os::unix::io::AsRawFd; -use termios::*; +use termios::{Termios, ECHO, ICANON, tcsetattr, TCSADRAIN}; /* https://man7.org/linux/man-pages/man3/tcflow.3.html */ -pub fn setup_fd() -> std::io::Result<()> { +pub fn setup_fd() -> Result<()> { let tty = OpenOptions::new().write(true).read(true).open("/dev/tty")?; let fd = tty.as_raw_fd(); let mut termios = Termios::from_fd(fd)?; diff --git a/src/main.rs b/src/main.rs index 45d9cdd..3feb737 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,9 @@ use clap::Parser; +use fern::Dispatch; use fern::colors::{Color, ColoredLevelConfig}; +use std::io::stdout; +use crate::listener::{Mode, Opts, listen}; +use crate::input::Command; mod input; mod listener; @@ -21,7 +25,7 @@ fn host_from_opts(host: Vec) -> Result<(String, String), String> { fn main() { // Configure logger - if let Err(err) = fern::Dispatch::new() + if let Err(err) = Dispatch::new() .format(|out, message, record| { let colors = ColoredLevelConfig::new() .warn(Color::Yellow) @@ -30,13 +34,14 @@ fn main() { out.finish(format_args!( "{}{} {}", - colors.color(record.level()).to_string().to_lowercase(), // Hardcoded red .red().bold() + colors.color(record.level()).to_string().to_lowercase(), ":", message )) }) - // .level(log::LevelFilter::Error) // For if verbose mode is implemented - .chain(std::io::stdout()) + .level(log::LevelFilter::Warn) + .level(log::LevelFilter::Info) + .chain(stdout()) .apply() { println!("Failed to initialize logger: {}", { err }); @@ -47,7 +52,7 @@ fn main() { let opts = input::Opts::parse(); match opts.command { - input::Command::Listen { + Command::Listen { interactive, block_signals, local_interactive, @@ -63,27 +68,25 @@ fn main() { } }; - let opts = listener::Opts { + let opts = Opts { host, port, exec, block_signals, mode: if interactive { - listener::Mode::Interactive + Mode::Interactive } else if local_interactive { - listener::Mode::LocalInteractive + Mode::LocalInteractive } else { - listener::Mode::Normal + Mode::Normal }, }; - if let Err(err) = listener::listen(&opts) { + if let Err(err) = listen(&opts) { log::error!("{}", err); - - return; }; } - input::Command::Connect { shell, host } => { + Command::Connect { shell, host } => { let (host, port) = match host_from_opts(host) { Ok(value) => value, Err(err) => { @@ -96,15 +99,11 @@ fn main() { #[cfg(unix)] if let Err(err) = unixshell::shell(host, port, shell) { log::error!("{}", err); - - return; } #[cfg(not(unix))] { log::error!("This feature is not supported on your platform"); - - return; } } } @@ -116,6 +115,8 @@ mod tests { #[cfg(unix)] use super::unixshell; + use std::io::ErrorKind; + // Panics if InvalidInput Not returned #[test] #[cfg(unix)] @@ -127,7 +128,7 @@ mod tests { "bash".to_string() ) .map_err(|e| e.kind()), - Err(std::io::ErrorKind::InvalidInput) + Err(ErrorKind::InvalidInput) ) } } diff --git a/src/unixshell.rs b/src/unixshell.rs index ca1d65e..06e12b3 100644 --- a/src/unixshell.rs +++ b/src/unixshell.rs @@ -9,7 +9,7 @@ pub fn shell(host: String, port: String, shell: String) -> Result<()> { let fd = sock.as_raw_fd(); // Open shell - Command::new(format!("{}", shell)) + Command::new(shell) .arg("-i") .stdin(unsafe { Stdio::from_raw_fd(fd) }) .stdout(unsafe { Stdio::from_raw_fd(fd) })