22
33from abc import ABC
44from enum import Enum
5- from typing import Annotated , Any , Set
5+ from typing import Annotated , Set
66
77from pydantic import BaseModel , Field , computed_field , model_serializer
88
@@ -45,7 +45,6 @@ class ThreatObjectType(str, Enum):
4545 TLS_HASH = "tls_hash"
4646 URL = "url"
4747
48- #[{"risk_object_field": "dest", "risk_object_type": "system", "risk_score": 64}, {"threat_object_field": "src_ip", "threat_object_type": "ip_address"}]
4948
5049class RiskObject (BaseModel ):
5150 field : str
@@ -62,19 +61,19 @@ def __lt__(self, other: RiskObject) -> bool:
6261 ):
6362 return True
6463 return False
65-
64+
6665 @model_serializer
6766 def serialize_risk_object (self ) -> dict [str , str | int ]:
68- '''
69- We define this explicitly for two reasons, even though the automatic
70- serialization works correctly. First we want to enforce a specific
71- field order for reasons of readability. Second, some of the fields
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
7271 actually have different names than they do in the object.
73- '''
74- return {
72+ """
73+ return {
7574 "risk_object_field" : self .field ,
7675 "risk_object_type" : self .type ,
77- "risk_score" : self .score
76+ "risk_score" : self .score ,
7877 }
7978
8079
@@ -89,16 +88,16 @@ def __lt__(self, other: ThreatObject) -> bool:
8988 if f"{ self .field } { self .type } " < f"{ other .field } { other .type } " :
9089 return True
9190 return False
92-
91+
9392 @model_serializer
9493 def serialize_threat_object (self ) -> dict [str , str ]:
95- '''
96- We define this explicitly for two reasons, even though the automatic
97- serialization works correctly. First we want to enforce a specific
98- field order for reasons of readability. Second, some of the fields
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
9998 actually have different names than they do in the object.
100- '''
101- return {
99+ """
100+ return {
102101 "threat_object_field" : self .field ,
103102 "threat_object_type" : self .type ,
104103 }
0 commit comments