Progress sync

.gitignore

@@ -12,4 +12,7 @@ run.sh
 rules.conf*
 Cargo.lock
 modsecurity
-cache
+cache
+
+/pdfs/*.pdf
+/reports

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"
+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"

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.  
+
+自動合併相關報告文件。  

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<String, Vec<SecurityEventData>>,
+}
+
+
+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::<Vec<_>>().join("\n")
+            )
+        }).collect::<Vec<_>>().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::<Vec<_>>();
+
+    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::<Vec<_>>();
+
+    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::<BTreeMap<ObjectId, Object>>(),
+        );
+        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::<Vec<_>>(),
+        );
+
+        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(())
+}