Change: Update NIST models to use StrEnums

greenbone · Nov 17, 2023 · d98108e · d98108e
1 parent 4bb0657
commit d98108e
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 39 deletions.
diff --git a/pontos/nvd/models/cpe.py b/pontos/nvd/models/cpe.py
@@ -17,10 +17,9 @@
 
 from dataclasses import dataclass, field
 from datetime import datetime
-from enum import Enum
 from typing import List, Optional
 
-from pontos.models import Model
+from pontos.models import Model, StrEnum
 
 __all__ = (
     "DeprecatedBy",
@@ -45,7 +44,7 @@ class Title(Model):
     lang: str
 
 
-class ReferenceType(Enum):
+class ReferenceType(StrEnum):
     """
     A CPE reference type
 

diff --git a/pontos/nvd/models/cve.py b/pontos/nvd/models/cve.py
@@ -17,10 +17,9 @@
 
 from dataclasses import dataclass, field
 from datetime import date, datetime
-from enum import Enum
 from typing import List, Optional
 
-from pontos.models import Model
+from pontos.models import Model, StrEnum
 from pontos.nvd.models.cvss_v2 import CVSSData as CVSSv2Data
 from pontos.nvd.models.cvss_v3 import CVSSData as CVSSv3Data
 
@@ -41,7 +40,7 @@
 )
 
 
-class CVSSType(Enum):
+class CVSSType(StrEnum):
     """
     The CVSS Type: primary or secondary
 
@@ -182,7 +181,7 @@ class VendorComment(Model):
     last_modified: datetime
 
 
-class Operator(Enum):
+class Operator(StrEnum):
     """
     An operator: AND or OR
 

diff --git a/pontos/nvd/models/cvss_v2.py b/pontos/nvd/models/cvss_v2.py
@@ -16,66 +16,65 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 from dataclasses import dataclass
-from enum import Enum
 from typing import Optional
 
-from pontos.models import Model
+from pontos.models import Model, StrEnum
 
 
-class Severity(Enum):
+class Severity(StrEnum):
     LOW = "LOW"
     MEDIUM = "MEDIUM"
     HIGH = "HIGH"
 
 
-class AccessVector(Enum):
+class AccessVector(StrEnum):
     NETWORK = "NETWORK"
     ADJACENT_NETWORK = "ADJACENT_NETWORK"
     LOCAL = "LOCAL"
 
 
-class AccessComplexity(Enum):
+class AccessComplexity(StrEnum):
     HIGH = "HIGH"
     MEDIUM = "MEDIUM"
     LOW = "LOW"
 
 
-class Authentication(Enum):
+class Authentication(StrEnum):
     MULTIPLE = "MULTIPLE"
     SINGLE = "SINGLE"
     NONE = "NONE"
 
 
-class Impact(Enum):
+class Impact(StrEnum):
     NONE = "NONE"
     PARTIAL = "PARTIAL"
     COMPLETE = "COMPLETE"
 
 
-class Exploitability(Enum):
+class Exploitability(StrEnum):
     UNPROVEN = "UNPROVEN"
     PROOF_OF_CONCEPT = "PROOF_OF_CONCEPT"
     FUNCTIONAL = "FUNCTIONAL"
     HIGH = "HIGH"
     NOT_DEFINED = "NOT_DEFINED"
 
 
-class RemediationLevel(Enum):
+class RemediationLevel(StrEnum):
     OFFICIAL_FIX = "OFFICIAL_FIX"
     TEMPORARY_FIX = "TEMPORARY_FIX"
     WORKAROUND = "WORKAROUND"
     UNAVAILABLE = "UNAVAILABLE"
     NOT_DEFINED = "NOT_DEFINED"
 
 
-class ReportConfidence(Enum):
+class ReportConfidence(StrEnum):
     UNCONFIRMED = "UNCONFIRMED"
     UNCORROBORATED = "UNCORROBORATED"
     CONFIRMED = "CONFIRMED"
     NOT_DEFINED = "NOT_DEFINED"
 
 
-class CollateralDamagePotential(Enum):
+class CollateralDamagePotential(StrEnum):
     NONE = "NONE"
     LOW = "LOW"
     LOW_MEDIUM = "LOW_MEDIUM"
@@ -84,15 +83,15 @@ class CollateralDamagePotential(Enum):
     NOT_DEFINED = "NOT_DEFINED"
 
 
-class TargetDistribution(Enum):
+class TargetDistribution(StrEnum):
     NONE = "NONE"
     LOW = "LOW"
     MEDIUM = "MEDIUM"
     HIGH = "HIGH"
     NOT_DEFINED = "NOT_DEFINED"
 
 
-class Requirement(Enum):
+class Requirement(StrEnum):
     LOW = "LOW"
     MEDIUM = "MEDIUM"
     HIGH = "HIGH"

diff --git a/pontos/nvd/models/cvss_v3.py b/pontos/nvd/models/cvss_v3.py
@@ -16,118 +16,117 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 from dataclasses import dataclass
-from enum import Enum
 from typing import Optional
 
-from pontos.models import Model
+from pontos.models import Model, StrEnum
 
 
-class Severity(Enum):
+class Severity(StrEnum):
     NONE = "NONE"
     LOW = "LOW"
     MEDIUM = "MEDIUM"
     HIGH = "HIGH"
     CRITICAL = "CRITICAL"
 
 
-class AttackVector(Enum):
+class AttackVector(StrEnum):
     NETWORK = "NETWORK"
     ADJACENT_NETWORK = "ADJACENT_NETWORK"
     LOCAL = "LOCAL"
     PHYSICAL = "PHYSICAL"
 
 
-class ModifiedAttackVector(Enum):
+class ModifiedAttackVector(StrEnum):
     NETWORK = "NETWORK"
     ADJACENT_NETWORK = "ADJACENT_NETWORK"
     LOCAL = "LOCAL"
     PHYSICAL = "PHYSICAL"
     NOT_DEFINED = "NOT_DEFINED"
 
 
-class AttackComplexity(Enum):
+class AttackComplexity(StrEnum):
     HIGH = "HIGH"
     LOW = "LOW"
 
 
-class ModifiedAttackComplexity(Enum):
+class ModifiedAttackComplexity(StrEnum):
     HIGH = "HIGH"
     LOW = "LOW"
     NOT_DEFINED = "NOT_DEFINED"
 
 
-class PrivilegesRequired(Enum):
+class PrivilegesRequired(StrEnum):
     HIGH = "HIGH"
     LOW = "LOW"
     NONE = "NONE"
 
 
-class ModifiedPrivilegesRequired(Enum):
+class ModifiedPrivilegesRequired(StrEnum):
     HIGH = "HIGH"
     LOW = "LOW"
     NONE = "NONE"
     NOT_DEFINED = "NOT_DEFINED"
 
 
-class UserInteraction(Enum):
+class UserInteraction(StrEnum):
     NONE = "NONE"
     REQUIRED = "REQUIRED"
 
 
-class ModifiedUserInteraction(Enum):
+class ModifiedUserInteraction(StrEnum):
     NONE = "NONE"
     REQUIRED = "REQUIRED"
     NOT_DEFINED = "NOT_DEFINED"
 
 
-class Scope(Enum):
+class Scope(StrEnum):
     UNCHANGED = "UNCHANGED"
     CHANGED = "CHANGED"
 
 
-class ModifiedScope(Enum):
+class ModifiedScope(StrEnum):
     UNCHANGED = "UNCHANGED"
     CHANGED = "CHANGED"
     NOT_DEFINED = "NOT_DEFINED"
 
 
-class Impact(Enum):
+class Impact(StrEnum):
     NONE = "NONE"
     LOW = "LOW"
     HIGH = "HIGH"
 
 
-class ModifiedImpact(Enum):
+class ModifiedImpact(StrEnum):
     NONE = "NONE"
     LOW = "LOW"
     HIGH = "HIGH"
     NOT_DEFINED = "NOT_DEFINED"
 
 
-class ExploitCodeMaturity(Enum):
+class ExploitCodeMaturity(StrEnum):
     UNPROVEN = "UNPROVEN"
     PROOF_OF_CONCEPT = "PROOF_OF_CONCEPT"
     FUNCTIONAL = "FUNCTIONAL"
     HIGH = "HIGH"
     NOT_DEFINED = "NOT_DEFINED"
 
 
-class RemediationLevel(Enum):
+class RemediationLevel(StrEnum):
     OFFICIAL_FIX = "OFFICIAL_FIX"
     TEMPORARY_FIX = "TEMPORARY_FIX"
     WORKAROUND = "WORKAROUND"
     UNAVAILABLE = "UNAVAILABLE"
     NOT_DEFINED = "NOT_DEFINED"
 
 
-class Confidence(Enum):
+class Confidence(StrEnum):
     UNKNOWN = "UNKNOWN"
     REASONABLE = "REASONABLE"
     CONFIRMED = "CONFIRMED"
     NOT_DEFINED = "NOT_DEFINED"
 
 
-class Requirement(Enum):
+class Requirement(StrEnum):
     LOW = "LOW"
     MEDIUM = "MEDIUM"
     HIGH = "HIGH"