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 tracing::Level;
@ -17,17 +19,43 @@ pub(crate) struct Args {
#[derive(Parser, Debug)]
pub(crate) enum ShellCommand {
/// Log out and close the shell
Exit,
/// Manage port forwards
#[command(name = "pfw")]
PortForwards {
#[command(subcommand)]
cmd: PortForwardsCommand
}
cmd: PortForwardsCommand,
},
}
#[derive(Parser, Debug)]
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 {

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_print::cprintln;
use connectbox::{models::{PortForwardEntry, PortForwardProtocol}, PortForwardAction};
use once_cell::sync::OnceCell;
use crate::{cli::PortForwardsCommand, AppState};
@ -13,12 +14,10 @@ fn init_port_forwarding_table() -> AsciiTable {
let mut t = AsciiTable::default();
t.column(0).set_header("ID").set_align(Right);
t.column(1).set_header("Local IP");
t.column(2).set_header("Start port");
t.column(3).set_header("End port");
t.column(4).set_header("In. start port");
t.column(5).set_header("In. end port");
t.column(6).set_header("Protocol");
t.column(7).set_header("Enabled");
t.column(2).set_header("Port range");
t.column(3).set_header("Int. port range");
t.column(4).set_header("Protocol");
t.column(5).set_header("Enabled");
t
}
@ -27,13 +26,104 @@ pub(crate) async fn run(cmd: PortForwardsCommand, state: &AppState) -> Result<()
PortForwardsCommand::Show => {
cprintln!("<blue!>Retrieving the port forwarding table...");
let port_forwards = state.connect_box.port_forwards().await?;
let table_entries = port_forwards.entries.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];
v
let table_entries = port_forwards.entries.into_iter().map(|e| {
let port_range = format!("{}-{}", e.start_port, e.end_port);
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);
cprintln!("<black!>LAN IP: {}\nSubnet mask: {}\n</black!>{rendered_table}", port_forwards.lan_ip, port_forwards.subnet_mask);
let rendered_table = PORT_FORWARDING_TABLE
.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(())
}
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 cli::Args;
@ -9,11 +9,11 @@ use rustyline::{error::ReadlineError, DefaultEditor};
use crate::{cli::ShellCommand, utils::QuotableArgs};
mod cli;
mod utils;
mod commands;
mod utils;
pub(crate) struct AppState {
connect_box: ConnectBox
connect_box: ConnectBox,
}
#[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)) {
Ok(mut matches) => ShellCommand::from_arg_matches_mut(&mut matches)?,
Err(e) => {
rl.add_history_entry(line)?;
e.print()?;
continue;
}
};
rl.add_history_entry(line)?;
match cmd {
ShellCommand::Exit => break,
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?;
rl.save_history(&history_path)?;