Implement the rest of the pfw command

This commit is contained in:
lemonsh 2023-05-20 11:29:15 +02:00
parent 7d5a93e30d
commit dd6a63187e
4 changed files with 142 additions and 22 deletions

View File

@ -1,3 +1,5 @@
use std::net::Ipv4Addr;
use clap::{Command, Parser, Subcommand}; use clap::{Command, Parser, Subcommand};
use tracing::Level; use tracing::Level;
@ -17,17 +19,43 @@ pub(crate) struct Args {
#[derive(Parser, Debug)] #[derive(Parser, Debug)]
pub(crate) enum ShellCommand { pub(crate) enum ShellCommand {
/// Log out and close the shell
Exit, Exit,
/// Manage port forwards
#[command(name = "pfw")] #[command(name = "pfw")]
PortForwards { PortForwards {
#[command(subcommand)] #[command(subcommand)]
cmd: PortForwardsCommand cmd: PortForwardsCommand,
} },
} }
#[derive(Parser, Debug)] #[derive(Parser, Debug)]
pub(crate) enum PortForwardsCommand { pub(crate) enum PortForwardsCommand {
Show /// List all port forwards
Show,
/// Add a port forward
Add {
/// LAN address of the host to forward the port to
local_ip: Ipv4Addr,
/// External port range
range: String,
/// Internal port range. If unspecified, the same as external
int_range: Option<String>,
/// TCP, UDP or both
#[arg(short, default_value = "both")]
protocol: String,
/// Add this port forward in disabled state
#[arg(short)]
disable: bool,
},
/// Enable, disable or delete a port forward
Edit {
/// ID of the port. You can use `pfw show` to find it
id: u32,
/// Action to perform with the port. Can be either enable, disable or delete.
action: String
}
} }
pub(crate) fn shell_cmd() -> Command { pub(crate) fn shell_cmd() -> Command {

View File

@ -1,8 +1,9 @@
use std::{vec, fmt::Display}; use std::vec;
use ascii_table::{AsciiTable, Align::Right}; use ascii_table::{Align::Right, AsciiTable};
use color_eyre::Result; use color_eyre::Result;
use color_print::cprintln; use color_print::cprintln;
use connectbox::{models::{PortForwardEntry, PortForwardProtocol}, PortForwardAction};
use once_cell::sync::OnceCell; use once_cell::sync::OnceCell;
use crate::{cli::PortForwardsCommand, AppState}; use crate::{cli::PortForwardsCommand, AppState};
@ -13,12 +14,10 @@ fn init_port_forwarding_table() -> AsciiTable {
let mut t = AsciiTable::default(); let mut t = AsciiTable::default();
t.column(0).set_header("ID").set_align(Right); t.column(0).set_header("ID").set_align(Right);
t.column(1).set_header("Local IP"); t.column(1).set_header("Local IP");
t.column(2).set_header("Start port"); t.column(2).set_header("Port range");
t.column(3).set_header("End port"); t.column(3).set_header("Int. port range");
t.column(4).set_header("In. start port"); t.column(4).set_header("Protocol");
t.column(5).set_header("In. end port"); t.column(5).set_header("Enabled");
t.column(6).set_header("Protocol");
t.column(7).set_header("Enabled");
t t
} }
@ -27,13 +26,104 @@ pub(crate) async fn run(cmd: PortForwardsCommand, state: &AppState) -> Result<()
PortForwardsCommand::Show => { PortForwardsCommand::Show => {
cprintln!("<blue!>Retrieving the port forwarding table..."); cprintln!("<blue!>Retrieving the port forwarding table...");
let port_forwards = state.connect_box.port_forwards().await?; let port_forwards = state.connect_box.port_forwards().await?;
let table_entries = port_forwards.entries.iter().map(|e| { let table_entries = port_forwards.entries.into_iter().map(|e| {
let v: Vec<&dyn Display> = vec![&e.id, &e.local_ip, &e.start_port, &e.end_port, &e.start_port_in, &e.end_port_in, &e.protocol, &e.enable]; let port_range = format!("{}-{}", e.start_port, e.end_port);
v let in_port_range = format!("{}-{}", e.start_port_in, e.end_port_in);
vec![
e.id.to_string(),
e.local_ip.to_string(),
port_range,
in_port_range,
e.protocol.to_string(),
e.enable.to_string(),
]
}); });
let rendered_table = PORT_FORWARDING_TABLE.get_or_init(init_port_forwarding_table).format(table_entries); let rendered_table = PORT_FORWARDING_TABLE
cprintln!("<black!>LAN IP: {}\nSubnet mask: {}\n</black!>{rendered_table}", port_forwards.lan_ip, port_forwards.subnet_mask); .get_or_init(init_port_forwarding_table)
.format(table_entries);
cprintln!(
"<black!>LAN IP: {}\nSubnet mask: {}\n</black!>{rendered_table}",
port_forwards.lan_ip,
port_forwards.subnet_mask
);
}
PortForwardsCommand::Add {
local_ip,
range,
int_range,
protocol,
disable,
} => {
let Some(protocol) = PortForwardProtocol::new(&protocol) else {
cprintln!("<red!>Invalid protocol {protocol:?}");
return Ok(())
};
let Some(range) = parse_port_range(&range) else {
cprintln!("<red!>Invalid external range {range:?}");
return Ok(())
};
let int_range = if let Some(r) = int_range {
let Some(r) = parse_port_range(&r) else {
cprintln!("<red!>Invalid internal range {r:?}");
return Ok(())
};
r
} else {
range
};
let port = PortForwardEntry {
id: 0,
local_ip,
start_port: range.0,
end_port: range.1,
start_port_in: int_range.0,
end_port_in: int_range.1,
protocol,
enable: !disable,
};
state.connect_box.add_port_forward(&port).await?;
cprintln!("<green!>Done!");
}
PortForwardsCommand::Edit { id, mut action } => {
action.make_ascii_lowercase();
let action = match action.as_str() {
"enable" => {
cprintln!("<blue!>Enabling port {id}...");
PortForwardAction::Enable
}
"disable" => {
cprintln!("<blue!>Disabling port {id}...");
PortForwardAction::Disable
}
"delete" => {
cprintln!("<blue!>Deleting port {id}...");
PortForwardAction::Delete
}
_ => {
cprintln!("<red!>Invalid action {action:?}");
return Ok(())
}
};
let mut modified = false;
state.connect_box.edit_port_forwards(|p| {
if p.id == id {
modified = true;
action
} else {
PortForwardAction::Keep
}
}).await?;
if !modified {
cprintln!("<red!>No port with id {id} exists");
} else {
cprintln!("<green!>Done!")
}
}, },
} }
Ok(()) Ok(())
} }
fn parse_port_range(s: &str) -> Option<(u16, u16)> {
let (start, end) = s.split_once('-')?;
Some((start.parse().ok()?, end.parse().ok()?))
}

View File

@ -1,4 +1,4 @@
use color_print::{cstr, cprintln}; use color_print::{cprintln, cstr};
use clap::{FromArgMatches, Parser}; use clap::{FromArgMatches, Parser};
use cli::Args; use cli::Args;
@ -9,11 +9,11 @@ use rustyline::{error::ReadlineError, DefaultEditor};
use crate::{cli::ShellCommand, utils::QuotableArgs}; use crate::{cli::ShellCommand, utils::QuotableArgs};
mod cli; mod cli;
mod utils;
mod commands; mod commands;
mod utils;
pub(crate) struct AppState { pub(crate) struct AppState {
connect_box: ConnectBox connect_box: ConnectBox,
} }
#[tokio::main(flavor = "current_thread")] #[tokio::main(flavor = "current_thread")]
@ -51,10 +51,12 @@ async fn main() -> Result<()> {
let cmd = match shell_cmd.try_get_matches_from_mut(QuotableArgs::new(&line)) { let cmd = match shell_cmd.try_get_matches_from_mut(QuotableArgs::new(&line)) {
Ok(mut matches) => ShellCommand::from_arg_matches_mut(&mut matches)?, Ok(mut matches) => ShellCommand::from_arg_matches_mut(&mut matches)?,
Err(e) => { Err(e) => {
rl.add_history_entry(line)?;
e.print()?; e.print()?;
continue; continue;
} }
}; };
rl.add_history_entry(line)?;
match cmd { match cmd {
ShellCommand::Exit => break, ShellCommand::Exit => break,
ShellCommand::PortForwards { cmd } => commands::pfw::run(cmd, &state).await?, ShellCommand::PortForwards { cmd } => commands::pfw::run(cmd, &state).await?,
@ -67,7 +69,7 @@ async fn main() -> Result<()> {
} }
} }
} }
println!("Logging out..."); cprintln!("<blue!>Logging out...");
state.connect_box.logout().await?; state.connect_box.logout().await?;
rl.save_history(&history_path)?; rl.save_history(&history_path)?;