feat: update ipset m

2025-08-02 13:57:40 +08:00
parent 3c30a313dd
commit dc97d211d8
5 changed files with 218 additions and 33 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -76,17 +76,6 @@ dependencies = [
 "windows-sys",
 ]

-[[package]]
-name = "async-trait"
-version = "0.1.88"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e539d3fca749fcee5236ab05e93a52867dd549cc157c8cb7f99595f3cedffdb5"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn",
-]
-
 [[package]]
 name = "autocfg"
 version = "1.5.0"
@@ -95,13 +84,13 @@ checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"

 [[package]]
 name = "axum"
-version = "0.7.9"
+version = "0.8.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f"
+checksum = "021e862c184ae977658b36c4500f7feac3221ca5da43e3f25bd04ab6c79a29b5"
 dependencies = [
- "async-trait",
 "axum-core",
 "bytes",
+ "form_urlencoded",
 "futures-util",
 "http",
 "http-body",
@@ -129,13 +118,12 @@ dependencies = [

 [[package]]
 name = "axum-core"
-version = "0.4.5"
+version = "0.5.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "09f2bd6146b97ae3359fa0cc6d6b376d9539582c7b4220f041a33ec24c226199"
+checksum = "68464cd0412f486726fb3373129ef5d2993f90c34bc2bc1c1e9943b2f4fc7ca6"
 dependencies = [
- "async-trait",
 "bytes",
- "futures-util",
+ "futures-core",
 "http",
 "http-body",
 "http-body-util",
@@ -181,6 +169,46 @@ version = "1.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9555578bc9e57714c812a1f84e4fc5b4d21fcb063490c624de019f7464c91268"

+[[package]]
+name = "clap"
+version = "4.5.42"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ed87a9d530bb41a67537289bafcac159cb3ee28460e0a4571123d2a778a6a882"
+dependencies = [
+ "clap_builder",
+ "clap_derive",
+]
+
+[[package]]
+name = "clap_builder"
+version = "4.5.42"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "64f4f3f3c77c94aff3c7e9aac9a2ca1974a5adf392a8bb751e827d6d127ab966"
+dependencies = [
+ "anstream",
+ "anstyle",
+ "clap_lex",
+ "strsim",
+]
+
+[[package]]
+name = "clap_derive"
+version = "4.5.41"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ef4f52386a59ca4c860f7393bcf8abd8dfd91ecccc0f774635ff68e92eeef491"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "clap_lex"
+version = "0.7.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b94f61472cee1439c0b966b47e3aca9ae07e45d070759512cd390ea2bebc6675"
+
 [[package]]
 name = "colorchoice"
 version = "1.0.4"
@@ -296,6 +324,12 @@ version = "0.31.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f"

+[[package]]
+name = "heck"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
+
 [[package]]
 name = "http"
 version = "1.3.1"
@@ -393,6 +427,7 @@ name = "ipset-management"
 version = "0.1.0"
 dependencies = [
 "axum",
+ "clap",
 "env_logger",
 "log",
 "rust_util",
@@ -477,9 +512,9 @@ checksum = "13dc2df351e3202783a1fe0d44375f7295ffb4049267b0f3018346dc122a1d94"

 [[package]]
 name = "matchit"
-version = "0.7.3"
+version = "0.8.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94"
+checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3"

 [[package]]
 name = "memchr"
@@ -778,6 +813,12 @@ dependencies = [
 "windows-sys",
 ]

+[[package]]
+name = "strsim"
+version = "0.11.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f"
+
 [[package]]
 name = "syn"
 version = "2.0.104"
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -6,10 +6,11 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

 [dependencies]
-axum = "0.7"
+axum = "0.8"
+clap = { version = "4.5", features = ["derive"] }
 env_logger = "0.11"
 log = "0.4"
-rust_util = "0.6.48"
+rust_util = "0.6"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 tokio = { version = "1.37", features = ["full"] }
--- a/README.md
+++ b/README.md
@@ -1,3 +1,8 @@
 # ipset-management

-Linux ipset management
+Linux ipset management
+
+
+```shell
+xh GET 0:2688/ipset x-ssrf-token:token
+```
--- a/src/ipset.rs
+++ b/src/ipset.rs
@@ -26,13 +26,7 @@ impl IpSet {

 impl IpSet {
    pub fn list(&self) -> XResult<Vec<String>> {
-        let (stdout, _stderr) = execute_ipset(&["list", &self.name])?;
-        let ipset_list_ips = opt_result!(
-            String::from_utf8(stdout.clone()),
-            "Parse output: {} failed: {}",
-            String::from_utf8_lossy(&stdout)
-        );
-        Ok(parse_ipset_list_ips(&ipset_list_ips))
+        list_ipset_ips(&self.name)
    }

    pub fn add(&self, ip: &str) -> XResult<()> {
@@ -54,7 +48,15 @@ fn list_names() -> XResult<Vec<String>> {
    Ok(parse_ipset_list(&ipset_list_names))
 }

-fn list_ipset(ipset_name: &str) {}
+fn list_ipset_ips(ipset_name: &str) -> XResult<Vec<String>> {
+    let (stdout, _stderr) = execute_ipset(&["list", ipset_name])?;
+    let ipset_list_ips = opt_result!(
+        String::from_utf8(stdout.clone()),
+        "Parse output: {} failed: {}",
+        String::from_utf8_lossy(&stdout)
+    );
+    Ok(parse_ipset_list_ips(&ipset_list_ips))
+}

 fn add_ipset(ipset_name: &str, ip: &str) -> XResult<()> {
    if let Err(err) = execute_ipset(&["add", ipset_name, ip]) {
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,5 +1,141 @@
+use crate::ipset::IpSet;
+use axum::extract::Path;
+use axum::http::{HeaderMap, HeaderValue, StatusCode};
+use axum::response::IntoResponse;
+use axum::routing::{get, post};
+use axum::{Form, Json, Router};
+use clap::Parser;
+use rust_util::XResult;
+use serde::Deserialize;
+use serde_json::json;
+use std::{env, process};
+use tokio::net::TcpListener;
+
+const SSRF_TOKEN: &str = "SSRF_TOKEN";
+
 mod ipset;

-fn main() {
-    println!("Hello, world!");
+#[derive(Debug, Parser)]
+#[command(name = "ipset-management")]
+#[command(about = "A ipset management tool", long_about = None)]
+struct Args {
+    /// SSRF token
+    #[arg(long, required = true)]
+    ssrf: String,
+
+    /// Port
+    #[arg(long)]
+    port: Option<u16>,
+}
+
+#[tokio::main]
+async fn main() {
+    if let Err(err) = inner_main().await {
+        eprintln!("Error: {}", err);
+        process::exit(1);
+    }
+}
+async fn inner_main() -> XResult<()> {
+    let args = Args::parse();
+    env::set_var(SSRF_TOKEN, args.ssrf);
+
+    let app = Router::new() // --
+        .route("/", get(root)) // --
+        .route("/ipset", get(get_ipset)) // --
+        .route("/ipset/{ipset}/ips", get(get_ipset_ips)) // --
+        .route("/ipset/{ipset}/ips/{ip}", post(delete_ipset_ips)) // --
+        .route("/ipset/{ipset}/ips", post(post_ipset_ips));
+
+    let listen_addr = format!("127.0.0.1:{}", args.port.unwrap_or(2688));
+    let listener = TcpListener::bind(&listen_addr).await.unwrap();
+    axum::serve(listener, app).await.unwrap();
+    Ok(())
+}
+
+async fn root() -> String {
+    format!("{} v{}", env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION"))
+}
+
+#[derive(Deserialize)]
+struct IpsetParams {
+    ipset: String,
+}
+#[derive(Deserialize)]
+struct IpsetIpParams {
+    ipset: String,
+    ip: String,
+}
+#[derive(Deserialize, Debug)]
+// #[allow(dead_code)]
+struct PostIpsetIpsRequest {
+    ip: String,
+}
+
+macro_rules! get_value_or_bad_request {
+    ($ex: expr) => {
+        match $ex {
+            Ok(o) => o,
+            Err(e) => {
+                return (StatusCode::BAD_REQUEST, Json(json!({"message": format!("error: {}", e)})));
+            }
+        }
+    };
+}
+
+macro_rules! check_header_ssrf_token {
+    ($header_map: expr) => {
+        match $header_map.get("x-ssrf-token") {
+            None => {
+                return (
+                    StatusCode::BAD_REQUEST,
+                    Json(json!({"message": format!("SSRF token required")})),
+                )
+            }
+            Some(ssrf_token) => {
+                if ssrf_token.to_str().unwrap_or("")
+                    != env::var(SSRF_TOKEN).unwrap_or("default-ssrf-token".to_string())
+                {
+                    return (
+                        StatusCode::BAD_REQUEST,
+                        Json(json!({"message": format!("invalid SSRF token: {:?}", ssrf_token)})),
+                    );
+                }
+            }
+        }
+    };
+}
+
+// GET /ipset
+async fn get_ipset(header_map: HeaderMap<HeaderValue>) -> impl IntoResponse {
+    check_header_ssrf_token!(header_map);
+    let names = get_value_or_bad_request!(IpSet::list_names());
+    (StatusCode::OK, Json(json!({"names": names})))
+}
+
+// GET /ipset/{}/ips
+async fn get_ipset_ips(
+    header_map: HeaderMap<HeaderValue>,
+    Path(params): Path<IpsetParams>,
+) -> impl IntoResponse {
+    check_header_ssrf_token!(header_map);
+    let ipset = get_value_or_bad_request!(IpSet::new(&params.ipset));
+    let ips = get_value_or_bad_request!(ipset.list());
+    (StatusCode::OK, Json(json!({"ips": ips})))
+}
+
+// DELETE /ipset/{}/ips/{}
+async fn delete_ipset_ips(Path(params): Path<IpsetIpParams>) -> impl IntoResponse {
+    let ipset = get_value_or_bad_request!(IpSet::new(&params.ipset));
+    get_value_or_bad_request!(ipset.del(&params.ip));
+    (StatusCode::OK, Json(json!({})))
+}
+
+// POST /ipset/{}/ips with parameter: ip=?
+async fn post_ipset_ips(
+    Path(params): Path<IpsetParams>,
+    Form(request): Form<PostIpsetIpsRequest>,
+) -> impl IntoResponse {
+    let ipset = get_value_or_bad_request!(IpSet::new(&params.ipset));
+    get_value_or_bad_request!(ipset.add(&request.ip));
+    (StatusCode::OK, Json(json!({})))
 }