1- from enum import Enum
2- from pydantic import BaseModel , computed_field , Field
1+ from __future__ import annotations
2+
33from abc import ABC
4- from typing import Set , Annotated
5- from contentctl . objects . enums import RiskSeverity
4+ from enum import Enum
5+ from typing import Annotated , Set
66
7+ from pydantic import BaseModel , Field , computed_field , model_serializer
8+
9+ from contentctl .objects .enums import RiskSeverity
710
811RiskScoreValue_Type = Annotated [int , Field (ge = 1 , le = 100 )]
912
@@ -51,6 +54,28 @@ class RiskObject(BaseModel):
5154 def __hash__ (self ):
5255 return hash ((self .field , self .type , self .score ))
5356
57+ def __lt__ (self , other : RiskObject ) -> bool :
58+ if (
59+ f"{ self .field } { self .type } { self .score } "
60+ < f"{ other .field } { other .type } { other .score } "
61+ ):
62+ return True
63+ return False
64+
65+ @model_serializer
66+ def serialize_risk_object (self ) -> dict [str , str | int ]:
67+ """
68+ We define this explicitly for two reasons, even though the automatic
69+ serialization works correctly. First we want to enforce a specific
70+ field order for reasons of readability. Second, some of the fields
71+ actually have different names than they do in the object.
72+ """
73+ return {
74+ "risk_object_field" : self .field ,
75+ "risk_object_type" : self .type ,
76+ "risk_score" : self .score ,
77+ }
78+
5479
5580class ThreatObject (BaseModel ):
5681 field : str
@@ -59,6 +84,24 @@ class ThreatObject(BaseModel):
5984 def __hash__ (self ):
6085 return hash ((self .field , self .type ))
6186
87+ def __lt__ (self , other : ThreatObject ) -> bool :
88+ if f"{ self .field } { self .type } " < f"{ other .field } { other .type } " :
89+ return True
90+ return False
91+
92+ @model_serializer
93+ def serialize_threat_object (self ) -> dict [str , str ]:
94+ """
95+ We define this explicitly for two reasons, even though the automatic
96+ serialization works correctly. First we want to enforce a specific
97+ field order for reasons of readability. Second, some of the fields
98+ actually have different names than they do in the object.
99+ """
100+ return {
101+ "threat_object_field" : self .field ,
102+ "threat_object_type" : self .type ,
103+ }
104+
62105
63106class RBAObject (BaseModel , ABC ):
64107 message : str
@@ -94,3 +137,11 @@ def severity(self) -> RiskSeverity:
94137 raise Exception (
95138 f"Error getting severity - risk_score must be between 0-100, but was actually { self .risk_score } "
96139 )
140+
141+ @model_serializer
142+ def serialize_rba (self ) -> dict [str , str | list [dict [str , str | int ]]]:
143+ return {
144+ "message" : self .message ,
145+ "risk_objects" : [obj .model_dump () for obj in sorted (self .risk_objects )],
146+ "threat_objects" : [obj .model_dump () for obj in sorted (self .threat_objects )],
147+ }
0 commit comments