diff --git a/.gitignore b/.gitignore index 0800012..9bf2194 100644 --- a/.gitignore +++ b/.gitignore @@ -12,4 +12,7 @@ run.sh rules.conf* Cargo.lock modsecurity -cache \ No newline at end of file +cache + +/pdfs/*.pdf +/reports \ No newline at end of file diff --git a/Cargo.toml b/Cargo.toml index 892dbe8..c7e8927 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,8 +11,25 @@ path = "src/lib.rs" name = "http_server" path = "src/bin/http_server.rs" +[[bin]] +name = "pdf_gen" +path = "src/bin/pdf_gen.rs" + +[[bin]] +name = "pdf_merge" +path = "src/bin/pdf_merge.rs" + +[[bin]] +name = "waf" +path = "src/bin/waf/main.rs" + [dependencies] tokio = { version = "1.28", features = ["full"] } tokio-rustls = "0.23" modsecurity = "1.0.0" -axum = "0.8.1" \ No newline at end of file +axum = "0.8.1" +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +walkdir = "2.4" +lopdf = "0.36.0" +regex = "1.11.1" diff --git a/pdfs/jf-openhuninn-2.1.ttf b/pdfs/jf-openhuninn-2.1.ttf new file mode 100644 index 0000000..49ce714 Binary files /dev/null and b/pdfs/jf-openhuninn-2.1.ttf differ diff --git a/src/bin/README_Zh-TW.md b/src/bin/README_Zh-TW.md index 61a987a..f3f78a0 100644 --- a/src/bin/README_Zh-TW.md +++ b/src/bin/README_Zh-TW.md @@ -3,9 +3,16 @@ ```plain bin ├── http_server.rs -└── main.rs +├── main.rs +└── pdf_merge.rs +└── pdf_gen.rs ``` +- [Main](#main) (main.rs) +- [HTTP Server](#http-server) (http_server.rs) +- [PDF Generator](#pdf-generator) (pdf_gen.rs) +- [PDF Merge](#pdf-merge) (pdf_merge.rs) + ## 使用方法 ```bash @@ -36,3 +43,12 @@ curl -X GET http://127.0.0.1/ ```bash curl http://localhost:3000 ``` + +## PDF Generator +產生報告文件。運作邏輯是讀取 `api_records.json` ,並產生 `.tex` 文件,再編譯成PDF。 + +### PDF Merge Function +> 參考自:[GitHub-lopdf](https://github.com/J-F-Liu/lopdf) +> A Rust library for PDF document manipulation. + +自動合併相關報告文件。 \ No newline at end of file diff --git a/src/bin/pdf_gen.rs b/src/bin/pdf_gen.rs new file mode 100644 index 0000000..92bd2be --- /dev/null +++ b/src/bin/pdf_gen.rs @@ -0,0 +1,313 @@ +use serde::Deserialize; +use std::process::Command; + +use lopdf::dictionary; +use std::collections::BTreeMap; +use lopdf::content::{Content, Operation}; +use lopdf::{Document, Object, ObjectId, Stream, Bookmark}; + +#[derive(Deserialize, Debug)] +struct SecurityEventData { + id: String, + title: String, + url: String, + date: String, +} + +#[derive(Deserialize, Debug)] +struct APIData { + #[serde(flatten)] + extra_fields: std::collections::HashMap>, +} + + +fn compile_latex(tex_content: &str) -> std::io::Result<()> { + use std::fs; + let tex_path = "pdfs/out.tex"; + if !std::path::Path::new("pdfs").exists() { + std::fs::create_dir("pdfs")?; + } + fs::write(tex_path, tex_content)?; + + // 使用 tectonic 編譯 LaTeX ( sudo pacman -S tectonic ) + let status = Command::new("tectonic") + .arg(tex_path) + .status() + .expect("failed to run tectonic"); + + if status.success() { + println!("PDF 產生成功!"); + } else { + println!("PDF 產生失敗!"); + } + // 刪除臨時檔案 + fs::remove_file(tex_path)?; + Ok(()) +} + +fn main() -> std::io::Result<()> { + // 讀取 JSON + let original_data = std::fs::read_to_string("api_records.json")?; + let api_data: APIData = serde_json::from_str(&original_data).expect("Invalid JSON format"); + + // 產生 LaTeX 字串 + let tex_content = format!(r#" + \documentclass{{article}} + \usepackage{{ctex}} + \usepackage{{xeCJK}} + \usepackage[a4paper, margin=2.5cm]{{geometry}} + \usepackage{{multicol}} + \usepackage{{xcolor}} + \setCJKmainfont{{jf-openhuninn-2.1.ttf}} + \usepackage[colorlinks=true, linkcolor=blue]{{hyperref}} + \title{{近期漏洞報告}} + \author{{WAFFl}} + \begin{{document}} + \renewcommand{{\contentsname}}{{Contents}} + \renewcommand{{\today}}{{\number\year-\number\month-\number\day}} + \maketitle + \tableofcontents + {} + \end{{document}} + "#, + api_data.extra_fields.iter().map(|(key, category)| { + format!( + r#"\section{{{}}} + \begin{{multicols}}{{2}} + \begin{{itemize}} + {} + \end{{itemize}} + \end{{multicols}} + "#, + key, + category.iter().map(|value| { + format!( + r#"\item \textbf{{[{}]-}}\href{{{}}}{{\textcolor{{blue}}{{{}}}}}"#, + value.id, value.url, value.title + ) + }).collect::>().join("\n") + ) + }).collect::>().join("\n\n") + ); + + compile_latex(tex_content.as_str())?; + + merge()?; + + Ok(()) +} + +// Refer to https://github.com/J-F-Liu/lopdf +fn merge() -> std::io::Result<()> { + // Generate a stack of Documents to merge. + let pdf_files = std::fs::read_dir("./pdfs") + .unwrap() + .filter_map(|entry| { + let entry = entry.ok()?; + let path = entry.path(); + if path.is_file() { + if let Some(name) = path.file_name().and_then(|n| n.to_str()) { + if name.starts_with("waf-evaluation-report") && name.ends_with(".pdf") { + return Some(path); + } + } + } + None + }) + .collect::>(); + + let report_files = std::fs::read_dir("./reports") + .unwrap() + .filter_map(|entry| { + let entry = entry.ok()?; + let path = entry.path(); + if path.is_file() { + if let Some(name) = path.file_name().and_then(|n| n.to_str()) { + if name.starts_with("waf-evaluation-report") && name.ends_with(".pdf") { + return Some(path); + } + } + } + None + }) + .collect::>(); + + let mut documents = vec![ + Document::load("./pdfs/out.pdf").unwrap(), + ]; + for pdf_path in pdf_files { documents.push(Document::load(pdf_path).unwrap()); } + for pdf_path in report_files { documents.push(Document::load(pdf_path).unwrap()); } + + // let documents = vec![ + // ]; + + // Define a starting `max_id` (will be used as start index for object_ids). + let mut max_id = 1; + let mut pagenum = 1; + // Collect all Documents Objects grouped by a map + let mut documents_pages = BTreeMap::new(); + let mut documents_objects = BTreeMap::new(); + let mut document = Document::with_version("1.5"); + + for mut doc in documents { + let mut first = false; + doc.renumber_objects_with(max_id); + + max_id = doc.max_id + 1; + + documents_pages.extend( + doc + .get_pages() + .into_iter() + .map(|(_, object_id)| { + if !first { + let bookmark = Bookmark::new(String::from(format!("Page_{}", pagenum)), [0.0, 0.0, 1.0], 0, object_id); + document.add_bookmark(bookmark, None); + first = true; + pagenum += 1; + } + + ( + object_id, + doc.get_object(object_id).unwrap().to_owned(), + ) + }) + .collect::>(), + ); + documents_objects.extend(doc.objects); + } + + // "Catalog" and "Pages" are mandatory. + let mut catalog_object: Option<(ObjectId, Object)> = None; + let mut pages_object: Option<(ObjectId, Object)> = None; + + // Process all objects except "Page" type + for (object_id, object) in documents_objects.iter() { + // We have to ignore "Page" (as are processed later), "Outlines" and "Outline" objects. + // All other objects should be collected and inserted into the main Document. + match object.type_name().unwrap_or(b"") { + b"Catalog" => { + // Collect a first "Catalog" object and use it for the future "Pages". + catalog_object = Some(( + if let Some((id, _)) = catalog_object { + id + } else { + *object_id + }, + object.clone(), + )); + } + b"Pages" => { + // Collect and update a first "Pages" object and use it for the future "Catalog" + // We have also to merge all dictionaries of the old and the new "Pages" object + if let Ok(dictionary) = object.as_dict() { + let mut dictionary = dictionary.clone(); + if let Some((_, ref object)) = pages_object { + if let Ok(old_dictionary) = object.as_dict() { + dictionary.extend(old_dictionary); + } + } + + pages_object = Some(( + if let Some((id, _)) = pages_object { + id + } else { + *object_id + }, + Object::Dictionary(dictionary), + )); + } + } + b"Page" => {} // Ignored, processed later and separately + b"Outlines" => {} // Ignored, not supported yet + b"Outline" => {} // Ignored, not supported yet + _ => { + document.objects.insert(*object_id, object.clone()); + } + } + } + + // If no "Pages" object found, abort. + if pages_object.is_none() { + println!("Pages root not found."); + + return Ok(()); + } + + // Iterate over all "Page" objects and collect into the parent "Pages" created before + for (object_id, object) in documents_pages.iter() { + if let Ok(dictionary) = object.as_dict() { + let mut dictionary = dictionary.clone(); + dictionary.set("Parent", pages_object.as_ref().unwrap().0); + + document + .objects + .insert(*object_id, Object::Dictionary(dictionary)); + } + } + + // If no "Catalog" found, abort. + if catalog_object.is_none() { + println!("Catalog root not found."); + + return Ok(()); + } + + let catalog_object = catalog_object.unwrap(); + let pages_object = pages_object.unwrap(); + + // Build a new "Pages" with updated fields + if let Ok(dictionary) = pages_object.1.as_dict() { + let mut dictionary = dictionary.clone(); + + // Set new pages count + dictionary.set("Count", documents_pages.len() as u32); + + // Set new "Kids" list (collected from documents pages) for "Pages" + dictionary.set( + "Kids", + documents_pages + .into_iter() + .map(|(object_id, _)| Object::Reference(object_id)) + .collect::>(), + ); + + document + .objects + .insert(pages_object.0, Object::Dictionary(dictionary)); + } + + // Build a new "Catalog" with updated fields + if let Ok(dictionary) = catalog_object.1.as_dict() { + let mut dictionary = dictionary.clone(); + dictionary.set("Pages", pages_object.0); + dictionary.remove(b"Outlines"); // Outlines not supported in merged PDFs + + document + .objects + .insert(catalog_object.0, Object::Dictionary(dictionary)); + } + + document.trailer.set("Root", catalog_object.0); + + // Update the max internal ID as wasn't updated before due to direct objects insertion + document.max_id = document.objects.len() as u32; + + // Reorder all new Document objects + document.renumber_objects(); + + // Set any Bookmarks to the First child if they are not set to a page + document.adjust_zero_pages(); + + // Set all bookmarks to the PDF Object tree then set the Outlines to the Bookmark content map. + if let Some(n) = document.build_outline() { + if let Ok(Object::Dictionary(dict)) = document.get_object_mut(catalog_object.0) { + dict.set("Outlines", Object::Reference(n)); + } + } + + document.compress(); + document.save("pdfs/final_report.pdf").unwrap(); + + Ok(()) +} \ No newline at end of file diff --git a/src/bin/main.rs b/src/bin/waf/main.rs similarity index 57% rename from src/bin/main.rs rename to src/bin/waf/main.rs index b3df89f..7e1c94b 100644 --- a/src/bin/main.rs +++ b/src/bin/waf/main.rs @@ -1,4 +1,5 @@ // 搭配 Nginx 使用 + use axum::{ body::{ to_bytes, Body }, extract::State, @@ -10,6 +11,7 @@ use axum::{ use modsecurity::{ ModSecurity, Rules }; use std::{ net::SocketAddr, sync::Arc, usize }; use tokio::{ net::TcpListener, sync::Mutex }; +use regex::Regex; #[derive(Clone)] struct AppState { @@ -89,20 +91,95 @@ async fn handle_request(State(state): State, req: Request) -> im response } +async fn reload_rules(State(state): State) -> impl IntoResponse { + println!("🔄 重新載入 ModSecurity 規則"); + let mut rules = Rules::new(); + match rules.add_file("./rules.conf") { + Ok(_) => { + let mut locked_rules = state.rules.lock().await; + *locked_rules = rules; + println!("✅ WAF 規則集已重新載入"); + StatusCode::OK + } + Err(e) => { + eprintln!("❌ 重新載入規則集失敗: {:?}", e); + StatusCode::INTERNAL_SERVER_ERROR + } + } +} + + +fn checkout_rules() -> Vec> { + use std::io::BufRead; + + let file = std::fs::File::open("rules.conf").expect("Cannot open rules.conf"); + let reader = std::io::BufReader::new(file); + + let mut phase_rules: Vec> = vec![vec![]; 8]; + let mut rule_lines = String::new(); + let mut in_rule = false; + let re_id = Regex::new(r"id:(\d+)").unwrap(); + let re_phase = Regex::new(r"phase:(\d+)").unwrap(); + + for line in reader.lines() { + let line = line.unwrap(); + let trimmed = line.trim(); + if trimmed.is_empty() || trimmed.starts_with('#') { continue; } + if trimmed.starts_with("SecRule") { + rule_lines.clear(); + in_rule = true; + } + if in_rule { + rule_lines.push_str(trimmed); + if !trimmed.ends_with('\\') { + let rule_lines_clone = rule_lines.clone(); + let id = re_id.captures(&rule_lines_clone) + .and_then(|cap| cap.get(1)) + .map(|m| m.as_str().to_string()); + let phase = re_phase.captures(&rule_lines_clone) + .and_then(|cap| cap.get(1)) + .map(|m| m.as_str().to_string()); + if let (Some(id), Some(phase)) = (id, phase) { + // println!("id: {}, phase: {}", id, phase); + phase_rules[phase.parse::().unwrap()].push(id); + } + in_rule = false; + } + } + } + return phase_rules; +} + +async fn dump_rules(State(_state): State) -> impl IntoResponse { + println!("📜 Dump ModSecurity 規則"); + let mut rules_str = String::new(); + + let phase_rules = checkout_rules(); + for (phase, rules) in phase_rules.iter().enumerate() { + println!("Phase {}: {} rules", phase, rules.len()); + rules_str.push_str(&format!("Phase {}: {} rules\n", phase, rules.len())); + for rule in rules { + println!(" Rule ID: {}", rule); + rules_str.push_str(&format!(" Rule ID: {}\n", rule)); + } + } + + (StatusCode::OK, rules_str) +} + #[tokio::main] async fn main() { - // 1. 載入 ModSecurity 規則 let mut rules = Rules::new(); rules.add_file("./rules.conf").expect("Failed to load ModSecurity rules"); let modsec = Arc::new(Mutex::new(ModSecurity::default())); let rules = Arc::new(Mutex::new(rules)); - // 2. 設定 Web 服務 - // .route("/*path", routing::any(handle_request)) let app = Router::new() .without_v07_checks() .route("/", routing::any(handle_request)) + .route("/reload_rules", routing::post(reload_rules)) + .route("/dump_rules", routing::get(dump_rules)) .route("/:path", routing::any(handle_request)) .fallback(handle_request) .with_state(AppState {