Linux下使用rust进行抓包

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Cap {
    stop: bool,        // ????????
    interface: String, // ??
    duration: u64,     // ??
    #[serde(skip_serializing, skip_deserializing)]
    stop_sender: Option>, // ??sender????
}

pub async fn cmd_capture(
    State(web_env): State,
    Json(mut args): Json,
) -> Result {
    let (stop_tx, stop_rx) = mpsc::channel::<()>();
    args.stop_sender = Some(stop_tx);
    let web_c: ArcWebEnv = web_env.clone();
    let (mac, status) = {
        let env = &mut web_env.env.lock().await;
        let status = env.capture_tasking.clone();
        (env.conf.mac.clone().split(':').collect::>().join(""), status)
    };
    // ???????
    if !args.stop && *status.lock().await {
        return Err(WebError::new(StatusCode::ACCEPTED, anyhow!("capture task is currently in progress. ignore")));
    }
    if args.stop {
        if *status.lock().await {
            if let Some(ref tx) = args.stop_sender {
                tx.send(()).map_err(|e| WebError::new(StatusCode::ACCEPTED, anyhow!("stop capture fail: {e}")))?;
            } else {
                return Err(WebError::new(StatusCode::ACCEPTED, anyhow!("stop sender not available")));
            }
        } else {
            return Err(WebError::new(StatusCode::ACCEPTED, anyhow!("capture task does not exist")));
        }
    }
    let fp = format!("/run/{}.pcap", &args.interface);
    let epx = format!("http://127.0.0.1:{}", LOCAL_EPX_PORT);

    *status.lock().await = true; // ????
    info!("capture upload status change to {}", *status.lock().await);
    spawn(async move {
        if let Err(e) = do_capture(axum::extract::State(web_c), args, fp.clone(), stop_rx, &epx, &mac).await {
            warn!("capture task error: {:?}", e);
        }
    });
    Ok(())
}

async fn do_capture(
    State(web_env): State,
    config: Cap,
    fp: String,
    stop_rec: mpsc::Receiver<()>,
    epx: &str,
    mac: &str,
) -> Result<()> {
    let mut cap = Capture::from_device(config.interface.as_str())
        .context("create device failed")?
        .promisc(true)
        .snaplen(65535)
        .open()
        .context("open capture session failed")?;
    warn!("start capture on: {}", config.interface);
    let duration = std::cmp::min(config.duration, CAPTURE_MAX_TIME as u64); // ?????????3600s
    let mut dump = cap.savefile(&fp).context("fail open capture file")?;
    let start = Instant::now();
    let mut stopped_by_signal = false;

    loop {
        if start.elapsed() >= Duration::from_secs(duration) {
            info!("capture finish, stopping capture");
            break;
        }
        match stop_rec.try_recv() {
            Ok(_) | Err(mpsc::TryRecvError::Disconnected) => {
                info!("received stop signal, stopping capture.");
                stopped_by_signal = true;
                break;
            }
            Err(mpsc::TryRecvError::Empty) => {}
        }
        match cap.next_packet() {
            Ok(packet) => {
                dump.write(&packet);
            }
            Err(e) => {
                warn!("capture fail: {:?}", e);
            }
        }
    }

    if stopped_by_signal {
        if Path::new(&fp).exists() {
            system_async_run("rm", &vec!["-f", &fp]).await?;
        };
        {
            let env = web_env.env.lock().await;
            *env.capture_tasking.lock().await = false;
        }
        info!("capture stopped by signal, file deleted.");
    } else {
        info!("capture finished, start compress file");
        if Path::new(&fp).exists() {
            system_async_run("gzip", &vec![&fp]).await?;
            let gz_path = format!("{}.gz", fp);
            if let Err(e) = upload_capture_file(epx.to_string(), mac.to_string(), gz_path.to_string()).await {
                warn!("upload file fail {}", e);
            };
            sleep(Duration::from_secs(2)).await;
            if let Err(e) = system_async_run("rm", &vec!["-f", &gz_path]).await {
                warn!("rm file fail {}", e);
            }
            {
                let env = web_env.env.lock().await;
                *env.capture_tasking.lock().await = false;
            }
            info!("capture all done.");
        }
    }
    Ok(())
}

async fn upload_capture_file(epx: String, mac: String, fs: String) -> Result<()> {
    let file_size = tokio::fs::metadata(&fs).await?.len();
    let rd = tokio::fs::File::open(&fs).await?;
    let fs = Path::new(&fs).file_name().ok_or(anyhow!("get file failed, {}", &fs))?.to_string_lossy();
    let url = &format!("{epx}/blob/aplog/{mac}/{fs}");
    let client = reqwest::Client::builder().timeout(Duration::from_secs(60)).build()?;
    let rsp = client
        .put(url)
        .timeout(Duration::from_secs(LOG_UPLOAD_TIMEOUT))
        .header("Content-Length", file_size)
        .header("Content-Type", "application/octet-stream")
        .body(rd)
        .send()
        .await?;
    if !rsp.status().is_success() {
        let status = rsp.status();
        let error_text = rsp.text().await.unwrap_or_else(|_| "Unable to read error response".to_string());
        bail!("upload capture file {} failed with status {}: {}", url, status, error_text);
    }

    info!("upload capture file success");
    Ok(())
}

你可能感兴趣的:(rust)