From dc97d211d8441410aafdddf0adb010298ac3456b Mon Sep 17 00:00:00 2001 From: Hatter Jiang Date: Sat, 2 Aug 2025 13:57:40 +0800 Subject: [PATCH] feat: update ipset m --- Cargo.lock | 81 +++++++++++++++++++++-------- Cargo.toml | 5 +- README.md | 7 ++- src/ipset.rs | 18 ++++--- src/main.rs | 140 ++++++++++++++++++++++++++++++++++++++++++++++++++- 5 files changed, 218 insertions(+), 33 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0f0fe1e..ea8d9a9 100644 --- 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" diff --git a/Cargo.toml b/Cargo.toml index b166c64..d70fec5 100644 --- 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"] } diff --git a/README.md b/README.md index 0389a91..6143846 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,8 @@ # ipset-management -Linux ipset management \ No newline at end of file +Linux ipset management + + +```shell +xh GET 0:2688/ipset x-ssrf-token:token +``` diff --git a/src/ipset.rs b/src/ipset.rs index 56bffc3..7d1c00a 100644 --- a/src/ipset.rs +++ b/src/ipset.rs @@ -26,13 +26,7 @@ impl IpSet { impl IpSet { pub fn list(&self) -> XResult> { - 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> { Ok(parse_ipset_list(&ipset_list_names)) } -fn list_ipset(ipset_name: &str) {} +fn list_ipset_ips(ipset_name: &str) -> XResult> { + 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]) { diff --git a/src/main.rs b/src/main.rs index 7e4e22a..7586f1b 100644 --- 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, +} + +#[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) -> 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, + Path(params): Path, +) -> impl IntoResponse { + check_header_ssrf_token!(header_map); + let ipset = get_value_or_bad_request!(IpSet::new(¶ms.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) -> impl IntoResponse { + let ipset = get_value_or_bad_request!(IpSet::new(¶ms.ipset)); + get_value_or_bad_request!(ipset.del(¶ms.ip)); + (StatusCode::OK, Json(json!({}))) +} + +// POST /ipset/{}/ips with parameter: ip=? +async fn post_ipset_ips( + Path(params): Path, + Form(request): Form, +) -> impl IntoResponse { + let ipset = get_value_or_bad_request!(IpSet::new(¶ms.ipset)); + get_value_or_bad_request!(ipset.add(&request.ip)); + (StatusCode::OK, Json(json!({}))) }