network-monitor/src/main.rs

use std::{format, time::Duration};

use clap::Parser;
use futures_util::{pin_mut, StreamExt};
use log::info;
use tokio::{
    io::{stdout, AsyncWriteExt},
    time::sleep,
};
use tokio_tungstenite::connect_async;

mod clash_conn_msg;
mod statistics;
mod udp_server;

#[derive(Parser)]
#[clap(
    version = "0.1.0",
    author = "Ivan Li",
    about = "Watch Clash network traffic"
)]
struct Args {
    #[clap(
        short,
        long,
        default_value = "ws://192.168.31.1:9090/connections?token=123456"
    )]
    clash_url: String,
    #[clap(short, long, default_value = "17890")]
    listen_port: u16,
}

#[tokio::main]
async fn main() {
    env_logger::init();

    println!("Hello, world!");

    let args = Args::parse();

    let connect_addr = args.clash_url;

    loop {
        pipe(connect_addr.clone()).await;
        sleep(Duration::from_secs(1)).await;
        info!("restart!");
    }
}

async fn pipe(connect_addr: String) {
    let url = url::Url::parse(&connect_addr).unwrap();

    let (ws_stream, _) = connect_async(url).await.expect("Failed to connect");
    println!("WebSocket handshake has been successfully completed");

    let (_, read) = ws_stream.split();

    let ws_to_stdout = {
        read.for_each(|message| async {
            let data = message.unwrap().into_data();

            let wrapper = serde_json::from_slice::<clash_conn_msg::ClashConnectionsWrapper>(&data);

            if let Err(err) = wrapper {
                stdout()
                    .write_all(format!("Error: {}\n", err).as_bytes())
                    .await
                    .unwrap();
                return;
            }
            let wrapper = wrapper.unwrap();

            let statistics_manager = statistics::StatisticsManager::global().await;

            statistics_manager.update(&wrapper).await;

            let state = statistics_manager.get_state().await;

            // stdout()
            //     .write_all(
            //         format!(
            //             "len: {},speed_upload: {}, speed_download: {}, update_direct_speed: {}, download_direct_speed: {}, update_proxy_speed: {}, download_proxy_speed: {}\n",
            //             state.connections,
            //             state.speed_upload,
            //             state.speed_download,
            //             state.direct_upload_speed,
            //             state.direct_download_speed,
            //             state.proxy_upload_speed,
            //             state.proxy_download_speed,
            //         )
            //         .as_bytes(),
            //     )
            //     .await
            //     .unwrap();

            let udp_server = udp_server::UdpServer::global().await;
            let mut buf = [0; 32];
            buf[0..8].copy_from_slice(&state.direct_upload_speed.to_le_bytes());
            buf[8..16].copy_from_slice(&state.direct_download_speed.to_le_bytes());
            buf[16..24].copy_from_slice(&state.proxy_upload_speed.to_le_bytes());
            buf[24..32].copy_from_slice(&state.proxy_download_speed.to_le_bytes());
            udp_server.publish(&buf).await;
        })
    };

    pin_mut!(ws_to_stdout);
    ws_to_stdout.await;
}