增加流量统计

This commit is contained in:
lbl8603
2024-07-06 16:27:08 +08:00
parent 7dc78cc170
commit d81ac6368e
31 changed files with 425 additions and 150 deletions
+21
View File
@@ -74,12 +74,15 @@ pub fn parse_args_config() -> anyhow::Result<Option<(Config, Vec<String>, bool)>
opts.optmulti("", "vnt-mapping", "vnt-mapping", "<mapping>");
opts.optopt("f", "", "配置文件", "<conf>");
opts.optopt("", "compressor", "压缩算法", "<lz4>");
opts.optflag("", "disable-stats", "关闭流量统计");
//"后台运行时,查看其他设备列表"
opts.optflag("", "add", "后台运行时,添加地址");
opts.optflag("", "list", "后台运行时,查看其他设备列表");
opts.optflag("", "all", "后台运行时,查看其他设备完整信息");
opts.optflag("", "info", "后台运行时,查看当前设备信息");
opts.optflag("", "route", "后台运行时,查看数据转发路径");
opts.optflag("", "chart_a", "后台运行时,查看流量统计");
opts.optopt("", "chart_b", "后台运行时,查看流量统计", "<IP>");
opts.optflag("", "stop", "停止后台运行");
opts.optflag("h", "help", "帮助");
let matches = match opts.parse(&args[1..]) {
@@ -110,6 +113,13 @@ pub fn parse_args_config() -> anyhow::Result<Option<(Config, Vec<String>, bool)>
} else if matches.opt_present("all") {
command::command(command::CommandEnum::All);
return Ok(None);
} else if matches.opt_present("chart_a") {
command::command(command::CommandEnum::ChartA);
return Ok(None);
}
if let Some(v) = matches.opt_str("chart_b") {
command::command(command::CommandEnum::ChartB(v));
return Ok(None);
}
let conf = matches.opt_str("f");
let (config, vnt_link_config, cmd) = if conf.is_some() {
@@ -267,6 +277,7 @@ pub fn parse_args_config() -> anyhow::Result<Option<(Config, Vec<String>, bool)>
#[cfg(feature = "port_mapping")]
let port_mapping_list = matches.opt_strs("mapping");
let vnt_mapping_list = matches.opt_strs("vnt-mapping");
let disable_stats = matches.opt_present("disable-stats");
let compressor = if let Some(compressor) = matches.opt_str("compressor").as_ref() {
Compressor::from_str(compressor)
.map_err(|e| anyhow!("{}", e))
@@ -306,6 +317,7 @@ pub fn parse_args_config() -> anyhow::Result<Option<(Config, Vec<String>, bool)>
#[cfg(feature = "port_mapping")]
port_mapping_list,
compressor,
!disable_stats,
) {
Ok(config) => config,
Err(e) => {
@@ -419,6 +431,7 @@ fn print_usage(program: &str, _opts: Options) {
.to_string()
)
);
println!(" --disable-stats 关闭流量统计");
println!();
#[cfg(feature = "command")]
{
@@ -443,6 +456,14 @@ fn print_usage(program: &str, _opts: Options) {
" --route {}",
yellow("后台运行时,查看数据转发路径".to_string())
);
println!(
" --chart_a {}",
yellow("后台运行时,查看所有IP的流量统计".to_string())
);
println!(
" --chart_b <IP> {}",
yellow("后台运行时,查看单个IP的历史流量".to_string())
);
println!(
" --stop {}",
yellow("停止后台运行".to_string())
+12 -1
View File
@@ -4,7 +4,7 @@ use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4, UdpSocket};
use std::str::FromStr;
use std::time::Duration;
use crate::command::entity::{DeviceItem, Info, RouteItem};
use crate::command::entity::{ChartA, ChartB, DeviceItem, Info, RouteItem};
pub struct CommandClient {
buf: Vec<u8>,
@@ -53,6 +53,17 @@ impl CommandClient {
pub fn info(&mut self) -> io::Result<Info> {
self.send_cmd(b"info")
}
pub fn chart_a(&mut self) -> io::Result<ChartA> {
self.send_cmd(b"chart_a")
}
pub fn chart_b(&mut self, input: &str) -> io::Result<ChartB> {
let cmd = if input.is_empty() {
"chart_b".to_string()
} else {
format!("chart_b:{}", input)
};
self.send_cmd(cmd.as_bytes())
}
fn send_cmd<'a, V: Deserialize<'a>>(&'a mut self, cmd: &[u8]) -> io::Result<V> {
self.udp.send(cmd)?;
let len = self.udp.recv(&mut self.buf)?;
+21 -2
View File
@@ -1,5 +1,7 @@
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::net::{Ipv4Addr, SocketAddr};
#[derive(Serialize, Deserialize, Debug)]
pub struct Info {
pub name: String,
@@ -12,8 +14,6 @@ pub struct Info {
pub public_ips: String,
pub local_addr: String,
pub ipv6_addr: String,
pub up: u64,
pub down: u64,
pub port_mapping_list: Vec<(bool, SocketAddr, String)>,
pub in_ips: Vec<(u32, u32, Ipv4Addr)>,
pub out_ips: Vec<(u32, u32)>,
@@ -46,3 +46,22 @@ pub struct DeviceItem {
pub current_client_secret: bool,
pub current_client_secret_hash: Vec<u8>,
}
#[derive(Serialize, Deserialize, Debug, Default)]
pub struct ChartA {
pub disable_stats: bool,
pub up_total: u64,
pub down_total: u64,
pub up_map: HashMap<Ipv4Addr, u64>,
pub down_map: HashMap<Ipv4Addr, u64>,
}
#[derive(Serialize, Deserialize, Debug, Default)]
pub struct ChartB {
pub disable_stats: bool,
pub ip: Option<Ipv4Addr>,
pub up_total: u64,
pub up_list: Vec<usize>,
pub down_total: u64,
pub down_list: Vec<usize>,
}
+98 -6
View File
@@ -1,8 +1,10 @@
use std::collections::HashSet;
use std::io;
use std::net::Ipv4Addr;
use vnt::channel::ConnectProtocol;
use vnt::core::Vnt;
use crate::command::entity::{DeviceItem, Info, RouteItem};
use crate::command::entity::{ChartA, ChartB, DeviceItem, Info, RouteItem};
use crate::console_out;
pub mod client;
@@ -14,6 +16,8 @@ pub enum CommandEnum {
List,
All,
Info,
ChartA,
ChartB(String),
Stop,
}
@@ -21,7 +25,9 @@ pub fn command_str(cmd: &str, vnt: &Vnt) -> bool {
if cmd.is_empty() {
return false;
}
match cmd.to_lowercase().trim() {
let cmd = cmd.to_lowercase();
let cmd = cmd.trim();
match cmd {
"list" => {
let list = command_list(&vnt);
console_out::console_device_list(list);
@@ -38,12 +44,24 @@ pub fn command_str(cmd: &str, vnt: &Vnt) -> bool {
let list = command_list(&vnt);
console_out::console_device_list_all(list);
}
"chart_a" => {
let chart = command_chart_a(&vnt);
console_out::console_chart_a(chart);
}
"stop" => {
let _ = vnt.stop();
return false;
}
_ => {}
}
if let Some(ip) = cmd.strip_prefix("chart_b") {
let chart = if ip.is_empty() {
command_chart_b(&vnt, &vnt.current_device().virtual_gateway.to_string())
} else {
command_chart_b(&vnt, &ip[1..])
};
console_out::console_chart_b(chart);
}
println!();
return true;
}
@@ -73,6 +91,14 @@ fn command_(cmd: CommandEnum) -> io::Result<()> {
let info = command_client.info()?;
console_out::console_info(info);
}
CommandEnum::ChartA => {
let chart = command_client.chart_a()?;
console_out::console_chart_a(chart);
}
CommandEnum::ChartB(input) => {
let chart = command_client.chart_b(&input)?;
console_out::console_chart_b(chart);
}
CommandEnum::Stop => {
command_client.stop()?;
}
@@ -224,8 +250,6 @@ pub fn command_info(vnt: &Vnt) -> Info {
.ipv6()
.map(|v| v.to_string())
.unwrap_or("None".to_string());
let up = vnt.up_stream();
let down = vnt.down_stream();
#[cfg(feature = "port_mapping")]
let port_mapping_list = vnt.config().port_mapping_list.clone();
#[cfg(not(feature = "port_mapping"))]
@@ -249,8 +273,6 @@ pub fn command_info(vnt: &Vnt) -> Info {
public_ips,
local_addr,
ipv6_addr,
up,
down,
port_mapping_list,
in_ips,
out_ips,
@@ -258,3 +280,73 @@ pub fn command_info(vnt: &Vnt) -> Info {
tcp_listen_addr,
}
}
pub fn command_chart_a(vnt: &Vnt) -> ChartA {
let disable_stats = !vnt.config().enable_traffic;
if disable_stats {
let mut chart = ChartA::default();
chart.disable_stats = true;
return chart;
}
let (up_total, up_map) = vnt.up_stream_all().unwrap_or_default();
let (down_total, down_map) = vnt.down_stream_all().unwrap_or_default();
ChartA {
disable_stats,
up_total,
down_total,
up_map,
down_map,
}
}
pub fn command_chart_b(vnt: &Vnt, input_str: &str) -> ChartB {
let disable_stats = !vnt.config().enable_traffic;
if disable_stats {
let mut chart = ChartB::default();
chart.disable_stats = true;
return chart;
}
let (_, up_map) = vnt.up_stream_history().unwrap_or_default();
let (_, down_map) = vnt.down_stream_history().unwrap_or_default();
let up_keys: HashSet<_> = up_map.keys().cloned().collect();
let down_keys: HashSet<_> = down_map.keys().cloned().collect();
let mut keys: Vec<Ipv4Addr> = up_keys.union(&down_keys).cloned().collect();
keys.sort();
if let Some(ip) = find_matching_ipv4_address(input_str, &keys) {
let (up_total, up_list) = up_map.get(&ip).cloned().unwrap_or_default();
let (down_total, down_list) = down_map.get(&ip).cloned().unwrap_or_default();
ChartB {
disable_stats,
ip: Some(ip),
up_total,
up_list,
down_total,
down_list,
}
} else {
ChartB::default()
}
}
fn match_from_end(input_str: &str, ip: &str) -> bool {
let mut input_chars = input_str.chars().rev();
let mut ip_chars = ip.chars().rev();
while let (Some(ic), Some(pc)) = (input_chars.next(), ip_chars.next()) {
if ic != pc {
return false;
}
}
input_chars.next().is_none() // Ensure all input characters matched
}
fn find_matching_ipv4_address(input_str: &str, ip_addresses: &[Ipv4Addr]) -> Option<Ipv4Addr> {
for &ip in ip_addresses {
let ip_str = ip.to_string();
if match_from_end(input_str, &ip_str) {
return Some(ip);
}
}
None
}
+16 -4
View File
@@ -1,3 +1,4 @@
use crate::command::command_chart_b;
use std::io;
use std::io::Write;
use std::net::UdpSocket;
@@ -62,15 +63,26 @@ fn command(cmd: &str, vnt: &Vnt) -> io::Result<String> {
.unwrap_or_else(|e| format!("error {:?}", e)),
"info" => serde_yaml::to_string(&crate::command::command_info(vnt))
.unwrap_or_else(|e| format!("error {:?}", e)),
"chart_a" => serde_yaml::to_string(&crate::command::command_chart_a(vnt))
.unwrap_or_else(|e| format!("error {:?}", e)),
"stop" => {
vnt.stop();
"stopped".to_string()
}
_ => {
format!(
"command '{}' not found. Try to enter: 'route'/'list'/'stop' \n",
cmd
)
if let Some(ip) = cmd.strip_prefix("chart_b") {
let chart = if ip.is_empty() {
command_chart_b(&vnt, &vnt.current_device().virtual_gateway.to_string())
} else {
command_chart_b(&vnt, &ip[1..])
};
serde_yaml::to_string(&chart).unwrap_or_else(|e| format!("error {:?}", e))
} else {
format!(
"command '{}' not found. Try to enter: 'route'/'list'/'stop' \n",
cmd
)
}
}
};
Ok(out_str)
+3
View File
@@ -45,6 +45,7 @@ pub struct FileConfig {
pub mapping: Vec<String>,
pub compressor: Option<String>,
pub vnt_mapping: Vec<String>,
pub disable_stats: bool,
}
impl Default for FileConfig {
@@ -88,6 +89,7 @@ impl Default for FileConfig {
mapping: vec![],
compressor: None,
vnt_mapping: vec![],
disable_stats: false,
}
}
}
@@ -174,6 +176,7 @@ pub fn read_config(file_path: &str) -> anyhow::Result<(Config, Vec<String>, bool
#[cfg(feature = "port_mapping")]
file_conf.mapping,
compressor,
!file_conf.disable_stats,
)?;
Ok((config, file_conf.vnt_mapping, file_conf.cmd))
+117 -3
View File
@@ -1,7 +1,8 @@
use console::{style, Style};
use std::collections::HashSet;
use std::net::Ipv4Addr;
use crate::command::entity::{DeviceItem, Info, RouteItem};
use crate::command::entity::{ChartA, ChartB, DeviceItem, Info, RouteItem};
pub mod table;
@@ -29,8 +30,6 @@ pub fn console_info(status: Info) {
println!("Public ips: {}", style(status.public_ips).green());
println!("Local addr: {}", style(status.local_addr).green());
println!("IPv6: {}", style(status.ipv6_addr).green());
println!("Up: {}", style(convert(status.up)).green());
println!("Down: {}", style(convert(status.down)).green());
if !status.port_mapping_list.is_empty() {
println!("------------------------------------------");
@@ -242,3 +241,118 @@ pub fn console_device_list_all(mut list: Vec<DeviceItem>) {
}
table::println_table(out_list)
}
pub fn console_chart_a(chart_a: ChartA) {
if chart_a.disable_stats {
println!("Traffic statistics not enabled");
return;
}
println!();
println!("-----------------------------------------------------------------");
println!(
"Upload total = {}",
style(convert(chart_a.up_total)).green()
);
println!(
"Download total = {}",
style(convert(chart_a.down_total)).green()
);
println!("-----------------------------------------------------------------");
let up_keys: HashSet<_> = chart_a.up_map.keys().cloned().collect();
let down_keys: HashSet<_> = chart_a.down_map.keys().cloned().collect();
let mut keys: Vec<Ipv4Addr> = up_keys.union(&down_keys).cloned().collect();
// 排序
keys.sort();
// 找到最大的值,用于缩放条形图长度
let up_max_value = *chart_a.up_map.values().max().unwrap_or(&0);
let down_max_value = *chart_a.down_map.values().max().unwrap_or(&0);
let max_value = up_max_value.max(down_max_value);
let max_height = 50;
// 打印条形图
for key in &keys {
if let Some(&value) = chart_a.up_map.get(key) {
let bar = "".repeat(((value as f64 / max_value as f64) * max_height as f64) as usize);
println!(
"{:<10} | {} upload {}",
key,
bar,
style(convert(value)).green()
);
}
if let Some(&value) = chart_a.down_map.get(key) {
let bar = "".repeat(((value as f64 / max_value as f64) * max_height as f64) as usize);
println!(
"{:<10} | {} download {}",
key,
bar,
style(convert(value)).green()
);
}
println!("-");
}
}
pub fn console_chart_b(chart_b: ChartB) {
if chart_b.disable_stats {
println!("Traffic statistics not enabled");
return;
}
let ip = if let Some(ip) = chart_b.ip {
ip
} else {
println!("Ip: None");
return;
};
println!("---------------------------- upload ----------------------------");
println!("IP: {}", ip);
println!("Upload total: {}", style(convert(chart_b.up_total)).green());
println!(
"Max: {}",
style(convert(
chart_b
.up_list
.iter()
.max()
.cloned()
.map_or(0, |v| v as u64)
))
.green()
);
console_chart_b_list(chart_b.up_list);
println!("---------------------------- download ----------------------------");
println!("IP: {}", ip);
println!(
"Download total: {}",
style(convert(chart_b.down_total)).green()
);
println!(
"Max: {}",
style(convert(
chart_b
.down_list
.iter()
.max()
.cloned()
.map_or(0, |v| v as u64)
))
.green()
);
console_chart_b_list(chart_b.down_list);
}
fn console_chart_b_list(list: Vec<usize>) {
let max_value = *list.iter().max().unwrap_or(&0);
let max_height = max_value.min(20);
// 遍历从最大高度到0
for i in (0..=max_height).rev() {
for &value in &list {
let scaled_value = (value as f64 / max_value as f64 * max_height as f64) as usize;
if scaled_value >= i {
print!("");
} else {
print!(" ");
}
}
println!();
}
}