From 475524d36d05019e30f3be4152d5654ce5d03f88 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BE=90=E6=B6=9B?= <midnite@outlook.com>
Date: Sun, 7 Apr 2024 16:32:09 +0800
Subject: [PATCH] =?UTF-8?q?feat(tools):=E5=A2=9E=E5=8A=A0=E6=94=AF?=
 =?UTF-8?q?=E6=8C=81=E7=9B=B4=E6=8E=A5=E7=94=9F=E6=88=90=E6=96=B0=E7=9A=84?=
 =?UTF-8?q?netfilter=E5=8E=8B=E7=BC=A9=E5=8C=85=E6=96=87=E4=BB=B6=E3=80=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cert_tools/src/power_equal.rs | 62 +++++++++++++++++++++++++++++++++--
 1 file changed, 60 insertions(+), 2 deletions(-)

diff --git a/cert_tools/src/power_equal.rs b/cert_tools/src/power_equal.rs
index 3c17f6f..7983073 100644
--- a/cert_tools/src/power_equal.rs
+++ b/cert_tools/src/power_equal.rs
@@ -1,10 +1,11 @@
 use std::{
     fs::File,
-    io::{BufWriter, Write},
+    io::{BufWriter, Read, Write},
     path::PathBuf,
 };
 
 use clap::Args;
+use zip::{write::FileOptions, ZipWriter};
 
 #[derive(Debug, Args)]
 pub struct PowerEqualResultOption {
@@ -12,13 +13,70 @@ pub struct PowerEqualResultOption {
     key_file: PathBuf,
     #[arg(short, long, default_value = "false", help = "Export to power.conf")]
     export: bool,
+    #[arg(
+        short,
+        long,
+        default_value = "false",
+        help = "Replace power.conf in given zip file, will override export option"
+    )]
+    zip_file: Option<PathBuf>,
 }
 
 pub fn calculate_equal_result(options: PowerEqualResultOption) {
     let cert = cert_lib::load_certificate(options.key_file).expect("load certificate failed");
     let result =
         cert_lib::calculate_power_euqal_result(cert).expect("calculate equal result failed");
-    if options.export {
+    if let Some(zip_file) = options.zip_file {
+        let mut buffer: Vec<u8> = Vec::new();
+        {
+            let mut buf_writer = BufWriter::new(&mut buffer);
+            buf_writer
+                .write("[Args]\n\n[Result]\n".as_bytes())
+                .expect("write power.conf failed");
+            buf_writer
+                .write(result.as_bytes())
+                .expect("write power.conf failed");
+            buf_writer.flush().expect("write power.conf failed");
+        }
+        // 将zip_file的文件名增加一个`_replaced`后缀并形成一个新的位于相同目录的PathBuf实例。
+        let mut zip_file_path = zip_file.clone();
+        let file_name = zip_file_path.file_stem().unwrap().to_str().unwrap();
+        let mut zip_file_name = file_name.to_string();
+        zip_file_name.push_str("_replaced.");
+        zip_file_name.push_str(&zip_file_path.extension().unwrap().to_string_lossy());
+        zip_file_path.set_file_name(zip_file_name);
+        // 创建一个新的zip文件并将zip_file的内容复制到新的zip文件中。
+        let zip_file = File::open(zip_file).expect("open zip file failed");
+        let mut zip_file = zip::ZipArchive::new(zip_file).expect("open zip file failed");
+        let zip_file_path = File::create(zip_file_path).expect("create zip file failed");
+        let mut zip_file_writer = ZipWriter::new(zip_file_path);
+        for i in 0..zip_file.len() {
+            let mut file = zip_file.by_index(i).expect("get file from zip failed");
+            let options = FileOptions::default()
+                .compression_method(file.compression())
+                .unix_permissions(file.unix_mode().unwrap());
+
+            let file_name = file.name().to_owned();
+            let file_name = file_name.replace("power.conf", "power.conf");
+
+            let content = if file_name.ends_with("power.conf") {
+                buffer.clone()
+            } else {
+                let mut content = Vec::new();
+                file.read_to_end(&mut content).expect("read file failed");
+                content
+            };
+
+            zip_file_writer
+                .start_file(file_name, options)
+                .expect("start output file failed");
+            zip_file_writer
+                .write_all(&content)
+                .expect("write output file failed");
+        }
+
+        zip_file_writer.finish().expect("write zip file failed");
+    } else if options.export {
         let mut power_conf_path = PathBuf::new();
         power_conf_path.push(".");
         power_conf_path.push("power.conf");