From 809d87897e8d9a876683ac94a745f9bf72867bcf Mon Sep 17 00:00:00 2001
From: Jia Chao <jiac13@chinaunicom.cn>
Date: Tue, 11 Jun 2024 16:40:32 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BD=BF=E7=94=A8=20enum=20=E5=A8=81=E8=83=81?=
 =?UTF-8?q?=E7=AD=89=E7=BA=A7=20Severity?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Jia Chao <jiac13@chinaunicom.cn>
---
 src/lib.rs  | 56 ++++++++++++++++++++++++++++++++++++++++++++++++++---
 src/test.rs |  2 ++
 2 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index f9fabb1..f2418bd 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -3,8 +3,10 @@
     allow(dead_code, unused_imports, unused_variables, unused_mut)
 )]
 use std::collections::HashMap;
+use std::fmt;
 use std::fs::File;
 use std::io::{self, BufReader};
+use std::str::FromStr;
 
 use serde::{Deserialize, Serialize};
 use tracing::{debug, error, instrument, trace};
@@ -901,14 +903,14 @@ pub struct Threat {
     pub r#type: String,
 
     // As threat level
-    pub description: String,
+    pub description: Severity,
 }
 
 impl Threat {
     pub fn new() -> Self {
         Threat {
             r#type: String::new(),
-            description: String::new(),
+            description: Severity::new(),
         }
     }
 
@@ -920,7 +922,7 @@ impl Threat {
                     if xmlreader.depth == 4 {
                         self.r#type = attributes[0].value.clone();
                     } else {
-                        self.description = xmlreader.next_characters();
+                        self.description = xmlreader.next_characters().parse::<Severity>().unwrap();
                     }
                 }
                 Ok(XmlEvent::EndElement { .. }) => {
@@ -938,6 +940,54 @@ impl Threat {
     }
 }
 
+#[derive(Debug, Clone, PartialOrd, Ord, PartialEq, Eq, Serialize, Deserialize)]
+pub enum Severity {
+    Null,
+    Low,
+    Moderate,
+    Important,
+    Critical,
+}
+
+impl Severity {
+    pub fn new() -> Self {
+        Severity::Null
+    }
+}
+
+// 为枚举 Severity 实现 FromStr trait
+impl FromStr for Severity {
+    type Err = ParseSeverityError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.to_lowercase().as_str() {
+            "low" => Ok(Severity::Low),
+            "moderate" | "medium" => Ok(Severity::Moderate),
+            "important" | "high" => Ok(Severity::Important),
+            "critical" => Ok(Severity::Critical),
+            _ => Err(ParseSeverityError::InvalidSeverity),
+        }
+    }
+}
+
+// 定义 ParseSeverityError 枚举类型来表示解析错误
+#[derive(Debug, Clone)]
+pub enum ParseSeverityError {
+    InvalidSeverity,
+}
+
+// 为 ParseSeverityError 实现 Display trait，以便更好地显示错误信息
+impl fmt::Display for ParseSeverityError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            ParseSeverityError::InvalidSeverity => write!(f, "Invalid severity level"),
+        }
+    }
+}
+
+// 为 ParseSeverityError 实现 std::error::Error trait
+impl std::error::Error for ParseSeverityError {}
+
 // depth = 4
 // <ScoreSet>
 //   <BaseScore>7.5</BaseScore>
diff --git a/src/test.rs b/src/test.rs
index 6e7edc0..fa9b7c8 100644
--- a/src/test.rs
+++ b/src/test.rs
@@ -90,6 +90,7 @@ fn cvrf_works() {
     let cvrf_vulner_cve = "CVE-2023-45288";
     let cvrf_vulner_productstatues_status = "Fixed";
     let cvrf_vulner_productstatues_product = "openEuler-22.03-LTS";
+    let cvrf_vulner_threat = Severity::Important;
     let cvrf_vulner_basescore = "7.5";
     let cvrf_vulner_vector = "AV:N/AC:L/PR:N/UI:N/S:U/C:N/I:N/A:H";
     let cvrf_vulner_remedition_type = "Vendor Fix";
@@ -108,6 +109,7 @@ fn cvrf_works() {
         cvrf.vulnerabilities[0].productstatuses[0].products[2],
         cvrf_vulner_productstatues_product
     );
+    assert_eq!(cvrf.vulnerabilities[0].threats[0].description, cvrf_vulner_threat);
     assert_eq!(
         cvrf.vulnerabilities[0].cvssscoresets[0].basescore,
         cvrf_vulner_basescore