catalystwan 0.41.2__py3-none-any.whl → 0.41.2.dev0__py3-none-any.whl

catalystwan/models/policy/definition/dns_security.py

@@ -0,0 +1,130 @@
+# Copyright 2024 Cisco Systems, Inc. and its affiliates
+
+from typing import List, Literal, Optional
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic.networks import IPvAnyAddress
+
+from catalystwan.models.common import IntStr, VpnId
+from catalystwan.models.policy.policy_definition import (
+    PolicyDefinitionBase,
+    PolicyDefinitionGetResponse,
+    PolicyDefinitionId,
+    Reference,
+)
+
+
+class TargetVpn(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+    dns_server_ip: Optional[IPvAnyAddress] = Field(
+        default=None, validation_alias="dnsServerIP", serialization_alias="dnsServerIP"
+    )
+
+    local_domain_bypass_enabled: bool = Field(
+        default=True, validation_alias="localDomainBypassEnabled", serialization_alias="localDomainBypassEnabled"
+    )
+
+    uid: Optional[int] = Field(default=None)
+    umbrella_default: bool = Field(validation_alias="umbrellaDefault", serialization_alias="umbrellaDefault")
+    vpns: List[VpnId]
+
+
+class DnsSecurityDefinition(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+    local_domain_bypass_enabled: Optional[bool] = Field(
+        default=None, validation_alias="localDomainBypassEnabled", serialization_alias="localDomainBypassEnabled"
+    )
+
+    local_domain_bypass_list: Optional[Reference] = Field(
+        default=None, validation_alias="localDomainBypassList", serialization_alias="localDomainBypassList"
+    )
+
+    dns_crypt: bool = Field(default=True, validation_alias="dnsCrypt", serialization_alias="dnsCrypt")
+    match_all_vpn: bool = Field(default=True, validation_alias="matchAllVpn", serialization_alias="matchAllVpn")
+
+    umbrella_data: Reference = Field(validation_alias="umbrellaData", serialization_alias="umbrellaData")
+    child_org_id: Optional[IntStr] = Field(
+        default=None, validation_alias="childOrgId", serialization_alias="childOrgId"
+    )
+
+    umbrella_default: Optional[bool] = Field(
+        default=None, validation_alias="umbrellaDefault", serialization_alias="umbrellaDefault"
+    )
+    dns_server_ip: Optional[str] = Field(  # using st, because server accepts non ip strings also
+        default=None, validation_alias="dnsServerIP", serialization_alias="dnsServerIP"
+    )
+
+    target_vpns: Optional[List[TargetVpn]] = Field(
+        default=None, validation_alias="targetVpns", serialization_alias="targetVpns"
+    )
+
+    @classmethod
+    def create_match_all_vpns_config(
+        cls,
+        umbrella_default: bool,
+        umbrella_data: Reference,
+        dns_crypt: bool = True,
+        dns_server_ip: Optional[str] = None,
+        local_domain_bypass_list: Optional[Reference] = None,
+        child_org_id: Optional[IntStr] = None,
+    ):
+        if not umbrella_default and not dns_server_ip:
+            raise ValueError("When umbrella default is False, then dns_server_ip must be provided.")
+
+        local_domain_bypass_enabled = True if local_domain_bypass_list else False
+
+        return cls(
+            match_all_vpn=True,
+            dns_crypt=dns_crypt,
+            dns_server_ip=dns_server_ip,
+            child_org_id=child_org_id,
+            local_domain_bypass_enabled=local_domain_bypass_enabled,
+            local_domain_bypass_list=local_domain_bypass_list,
+            umbrella_data=umbrella_data,
+            umbrella_default=umbrella_default,
+        )
+
+    @classmethod
+    def create_match_custom_vpns_config(
+        cls,
+        dns_crypt: bool,
+        umbrella_data: Reference,
+        target_vpns=List[TargetVpn],
+        local_domain_bypass_list: Optional[Reference] = None,
+        child_org_id: Optional[IntStr] = None,
+    ):
+        return cls(
+            match_all_vpn=False,
+            dns_crypt=dns_crypt,
+            child_org_id=child_org_id,
+            local_domain_bypass_list=local_domain_bypass_list,
+            umbrella_data=umbrella_data,
+            target_vpns=target_vpns,
+        )
+
+    @field_validator("local_domain_bypass_list", mode="before")
+    @classmethod
+    def convert_empty_dict_to_none(cls, value):
+        if not value:
+            return None
+        return value
+
+    @field_validator("child_org_id", mode="before")
+    @classmethod
+    def convert_empty_int_str_to_none(cls, value):
+        if value == "":
+            return None
+        return value
+
+
+class DnsSecurityPolicy(PolicyDefinitionBase):
+    type: Literal["dnsSecurity"] = "dnsSecurity"
+    definition: DnsSecurityDefinition
+
+
+class DnsSecurityPolicyEditPayload(DnsSecurityPolicy, PolicyDefinitionId):
+    pass
+
+
+class DnsSecurityPolicyGetResponse(DnsSecurityPolicy, PolicyDefinitionGetResponse):
+    pass

catalystwan/models/policy/definition/fxo_port.py

@@ -0,0 +1,63 @@
+# Copyright 2024 Cisco Systems, Inc. and its affiliates
+
+from typing import List, Literal, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from catalystwan.models.policy.policy_definition import (
+    CommonStation,
+    FxoTuningParams,
+    LineParams,
+    PolicyDefinitionBase,
+    PolicyDefinitionGetResponse,
+    PolicyDefinitionId,
+    SupervisoryDisconnectEntry,
+    TranslationProfileEntry,
+    TrunkGroupPreference,
+    VoicePortType,
+)
+
+
+class FxoPortPolicyDefinition(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+    common_station: Optional[CommonStation] = Field(
+        default=None, serialization_alias="commonStation", validation_alias="commonStation"
+    )
+    line_params: Optional[LineParams] = Field(
+        default=None, serialization_alias="lineParams", validation_alias="lineParams"
+    )
+    tuning_params: Optional[FxoTuningParams] = Field(
+        default=None, serialization_alias="tuningParams", validation_alias="tuningParams"
+    )
+    trunk_group: Optional[List[TrunkGroupPreference]] = Field(
+        default=None, serialization_alias="trunkGroup", validation_alias="trunkGroup"
+    )
+    incoming_translation_profile: Optional[TranslationProfileEntry] = Field(
+        default=None, serialization_alias="incomingTranslationProfile", validation_alias="incomingTranslationProfile"
+    )
+    outgoing_translation_profile: Optional[TranslationProfileEntry] = Field(
+        default=None, serialization_alias="outgoingTranslationProfile", validation_alias="outgoingTranslationProfile"
+    )
+    supervisory_disc_custom: Optional[SupervisoryDisconnectEntry] = Field(
+        default=None, serialization_alias="supervisoryDiscCustom", validation_alias="supervisoryDiscCustom"
+    )
+    supervisory_disc_dualtone: Optional[SupervisoryDisconnectEntry] = Field(
+        default=None, serialization_alias="supervisoryDiscDualtone", validation_alias="supervisoryDiscDualtone"
+    )
+
+
+class FxoPortPolicy(PolicyDefinitionBase):
+    model_config = ConfigDict(populate_by_name=True)
+    type: Literal["fxoPort", "fxoport"] = "fxoPort"
+    port_type: Optional[VoicePortType] = Field(
+        default="voicePort", serialization_alias="portType", validation_alias="portType"
+    )
+    definition: FxoPortPolicyDefinition = FxoPortPolicyDefinition()
+
+
+class FxoPortPolicyEditPayload(FxoPortPolicy, PolicyDefinitionId):
+    pass
+
+
+class FxoPortPolicyGetResponse(FxoPortPolicy, PolicyDefinitionGetResponse):
+    pass

catalystwan/models/policy/definition/fxs_did_port.py

@@ -0,0 +1,58 @@
+# Copyright 2024 Cisco Systems, Inc. and its affiliates
+
+from typing import List, Literal, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from catalystwan.models.policy.policy_definition import (
+    CommonStation,
+    DidTimers,
+    FxsTuningParams,
+    LineParams,
+    PolicyDefinitionBase,
+    PolicyDefinitionGetResponse,
+    PolicyDefinitionId,
+    TranslationProfileEntry,
+    TrunkGroupPreference,
+    VoicePortType,
+)
+
+
+class FxsDidPortPolicyDefinition(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+    common_station: Optional[CommonStation] = Field(
+        default=None, serialization_alias="commonStation", validation_alias="commonStation"
+    )
+    line_params: Optional[LineParams] = Field(
+        default=None, serialization_alias="lineParams", validation_alias="lineParams"
+    )
+    tuning_params: Optional[FxsTuningParams] = Field(
+        default=None, serialization_alias="tuningParams", validation_alias="tuningParams"
+    )
+    trunk_group: Optional[List[TrunkGroupPreference]] = Field(
+        default=None, serialization_alias="trunkGroup", validation_alias="trunkGroup"
+    )
+    incoming_translation_profile: Optional[TranslationProfileEntry] = Field(
+        default=None, serialization_alias="incomingTranslationProfile", validation_alias="incomingTranslationProfile"
+    )
+    outgoing_translation_profile: Optional[TranslationProfileEntry] = Field(
+        default=None, serialization_alias="outgoingTranslationProfile", validation_alias="outgoingTranslationProfile"
+    )
+    did_timers: Optional[DidTimers] = Field(default=None, serialization_alias="didTimers", validation_alias="didTimers")
+
+
+class FxsDidPortPolicy(PolicyDefinitionBase):
+    model_config = ConfigDict(populate_by_name=True)
+    type: Literal["fxsDidPort", "fxsdidport"] = "fxsDidPort"
+    port_type: Optional[VoicePortType] = Field(
+        default="voicePort", serialization_alias="portType", validation_alias="portType"
+    )
+    definition: FxsDidPortPolicyDefinition = FxsDidPortPolicyDefinition()
+
+
+class FxsDidPortPolicyEditPayload(FxsDidPortPolicy, PolicyDefinitionId):
+    pass
+
+
+class FxsDidPortPolicyGetResponse(FxsDidPortPolicy, PolicyDefinitionGetResponse):
+    pass

catalystwan/models/policy/definition/fxs_port.py

@@ -0,0 +1,56 @@
+# Copyright 2024 Cisco Systems, Inc. and its affiliates
+
+from typing import List, Literal, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from catalystwan.models.policy.policy_definition import (
+    CommonStation,
+    FxsTuningParams,
+    LineParams,
+    PolicyDefinitionBase,
+    PolicyDefinitionGetResponse,
+    PolicyDefinitionId,
+    TranslationProfileEntry,
+    TrunkGroupPreference,
+    VoicePortType,
+)
+
+
+class FxsPortPolicyDefinition(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+    common_station: Optional[CommonStation] = Field(
+        default=None, serialization_alias="commonStation", validation_alias="commonStation"
+    )
+    line_params: Optional[LineParams] = Field(
+        default=None, serialization_alias="lineParams", validation_alias="lineParams"
+    )
+    tuning_params: Optional[FxsTuningParams] = Field(
+        default=None, serialization_alias="tuningParams", validation_alias="tuningParams"
+    )
+    trunk_group: Optional[List[TrunkGroupPreference]] = Field(
+        default=None, serialization_alias="trunkGroup", validation_alias="trunkGroup"
+    )
+    incoming_translation_profile: Optional[TranslationProfileEntry] = Field(
+        default=None, serialization_alias="incomingTranslationProfile", validation_alias="incomingTranslationProfile"
+    )
+    outgoing_translation_profile: Optional[TranslationProfileEntry] = Field(
+        default=None, serialization_alias="outgoingTranslationProfile", validation_alias="outgoingTranslationProfile"
+    )
+
+
+class FxsPortPolicy(PolicyDefinitionBase):
+    model_config = ConfigDict(populate_by_name=True)
+    type: Literal["fxsPort", "fxsport"] = "fxsPort"
+    port_type: Optional[VoicePortType] = Field(
+        default="voicePort", serialization_alias="portType", validation_alias="portType"
+    )
+    definition: FxsPortPolicyDefinition = FxsPortPolicyDefinition()
+
+
+class FxsPortPolicyEditPayload(FxsPortPolicy, PolicyDefinitionId):
+    pass
+
+
+class FxsPortPolicyGetResponse(FxsPortPolicy, PolicyDefinitionGetResponse):
+    pass

catalystwan/models/policy/{definitions → definition}/hub_and_spoke.py

@@ -5,16 +5,22 @@ from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict, Field
 
-from catalystwan.models.policy.policy_definition import PolicyDefinitionBase
+from catalystwan.models.policy.policy_definition import (
+    PolicyDefinitionBase,
+    PolicyDefinitionGetResponse,
+    PolicyDefinitionId,
+)
 
 
 class Hub(BaseModel):
     model_config = ConfigDict(populate_by_name=True)
     site_list: UUID = Field(validation_alias="siteList", serialization_alias="siteList")
     preference: Optional[str] = None
-    prefix_lists: List[UUID] = Field(default=[], validation_alias="prefixLists", serialization_alias="prefixLists")
+    prefix_lists: List[UUID] = Field(
+        default_factory=list, validation_alias="prefixLists", serialization_alias="prefixLists"
+    )
     ipv6_prefix_lists: List[UUID] = Field(
-        default=[], validation_alias="ipv6PrefixLists", serialization_alias="ipv6PrefixLists"
+        default_factory=list, validation_alias="ipv6PrefixLists", serialization_alias="ipv6PrefixLists"
     )
 
 
@@ -39,7 +45,7 @@ class HubAndSpokePolicyDefinition(BaseModel):
     model_config = ConfigDict(populate_by_name=True)
     vpn_list: UUID = Field(validation_alias="vpnList", serialization_alias="vpnList")
     sub_definitions: List[HubAndSpokePolicySubDefinition] = Field(
-        default=[], validation_alias="subDefinitions", serialization_alias="subDefinitions"
+        default_factory=list, validation_alias="subDefinitions", serialization_alias="subDefinitions"
     )
 
 
@@ -71,3 +77,11 @@ class HubAndSpokePolicy(PolicyDefinitionBase):
         )
         self.definition.sub_definitions.append(sub_definition)
         return sub_definition
+
+
+class HubAndSpokePolicyEditPayload(HubAndSpokePolicy, PolicyDefinitionId):
+    pass
+
+
+class HubAndSpokePolicyGetResponse(HubAndSpokePolicy, PolicyDefinitionGetResponse):
+    pass

catalystwan/models/policy/definition/intrusion_prevention.py

@@ -0,0 +1,55 @@
+# Copyright 2024 Cisco Systems, Inc. and its affiliates
+
+from typing import List, Literal, Optional
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+from catalystwan.models.common import PolicyModeType, VpnId
+from catalystwan.models.policy.policy_definition import (
+    PolicyDefinitionBase,
+    PolicyDefinitionGetResponse,
+    PolicyDefinitionId,
+    Reference,
+)
+
+SignatureSetType = Literal["balanced", "connectivity", "security"]
+InspectionModeType = Literal["protection", "detection"]
+LogLevel = Literal["emergency", "alert", "critical", "error", "warning", "notice", "info", "debug"]
+
+
+class IntrusionPreventionDefinition(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+    signature_set: SignatureSetType = Field(validation_alias="signatureSet", serialization_alias="signatureSet")
+    inspection_mode: InspectionModeType = Field(validation_alias="inspectionMode", serialization_alias="inspectionMode")
+    signature_white_list: Optional[Reference] = Field(
+        default=None, validation_alias="signatureWhiteList", serialization_alias="signatureWhiteList"
+    )
+    log_level: Optional[LogLevel] = Field(default="error", validation_alias="logLevel", serialization_alias="logLevel")
+    logging: List[str] = Field(default_factory=list)
+    target_vpns: List[VpnId] = Field(
+        default_factory=list, validation_alias="targetVpns", serialization_alias="targetVpns"
+    )
+    custom_signature: bool = Field(
+        default=False, validation_alias="customSignature", serialization_alias="customSignature"
+    )
+
+    @field_validator("signature_white_list", mode="before")
+    @classmethod
+    def convert_empty_dict_to_none(cls, value):
+        if not value:
+            return None
+        return value
+
+
+class IntrusionPreventionPolicy(PolicyDefinitionBase):
+    type: Literal["intrusionPrevention"] = "intrusionPrevention"
+    mode: PolicyModeType = "security"
+    definition: IntrusionPreventionDefinition
+
+
+class IntrusionPreventionPolicyEditPayload(IntrusionPreventionPolicy, PolicyDefinitionId):
+    pass
+
+
+class IntrusionPreventionPolicyGetResponse(IntrusionPreventionPolicy, PolicyDefinitionGetResponse):
+    pass

catalystwan/models/policy/{definitions → definition}/mesh.py

@@ -5,7 +5,11 @@ from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict, Field
 
-from catalystwan.models.policy.policy_definition import PolicyDefinitionBase
+from catalystwan.models.policy.policy_definition import (
+    PolicyDefinitionBase,
+    PolicyDefinitionGetResponse,
+    PolicyDefinitionId,
+)
 
 
 class Region(BaseModel):
@@ -33,3 +37,11 @@ class MeshPolicy(PolicyDefinitionBase):
         region = Region(name=name, site_lists=site_lists)
         self.definition.regions.append(region)
         return region
+
+
+class MeshPolicyEditPayload(MeshPolicy, PolicyDefinitionId):
+    pass
+
+
+class MeshPolicyGetResponse(MeshPolicy, PolicyDefinitionGetResponse):
+    pass

catalystwan/models/policy/definition/pri_isdn_port.py

@@ -0,0 +1,48 @@
+# Copyright 2024 Cisco Systems, Inc. and its affiliates
+
+from typing import List, Literal, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from catalystwan.models.policy.policy_definition import (
+    LineParams,
+    PolicyDefinitionBase,
+    PolicyDefinitionGetResponse,
+    PolicyDefinitionId,
+    TranslationProfileEntry,
+    TrunkGroupPreference,
+    VoicePortType,
+)
+
+
+class PriIsdnPortPolicyDefinition(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+    line_params: Optional[LineParams] = Field(
+        default=None, serialization_alias="lineParams", validation_alias="lineParams"
+    )
+    trunk_group: Optional[List[TrunkGroupPreference]] = Field(
+        default=None, serialization_alias="trunkGroup", validation_alias="trunkGroup"
+    )
+    incoming_translation_profile: Optional[TranslationProfileEntry] = Field(
+        default=None, serialization_alias="incomingTranslationProfile", validation_alias="incomingTranslationProfile"
+    )
+    outgoing_translation_profile: Optional[TranslationProfileEntry] = Field(
+        default=None, serialization_alias="outgoingTranslationProfile", validation_alias="outgoingTranslationProfile"
+    )
+
+
+class PriIsdnPortPolicy(PolicyDefinitionBase):
+    model_config = ConfigDict(populate_by_name=True)
+    type: Literal["priIsdnPort", "priisdnport"] = "priIsdnPort"
+    port_type: Optional[VoicePortType] = Field(
+        default="voicePort", serialization_alias="portType", validation_alias="portType"
+    )
+    definition: PriIsdnPortPolicyDefinition = PriIsdnPortPolicyDefinition()
+
+
+class PriIsdnPortPolicyEditPayload(PriIsdnPortPolicy, PolicyDefinitionId):
+    pass
+
+
+class PriIsdnPortPolicyGetResponse(PriIsdnPortPolicy, PolicyDefinitionGetResponse):
+    pass

catalystwan/models/policy/{definitions → definition}/qos_map.py

@@ -5,7 +5,12 @@ from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
-from catalystwan.models.policy.policy_definition import PolicyDefinitionBase
+from catalystwan.models.common import IntStr
+from catalystwan.models.policy.policy_definition import (
+    PolicyDefinitionBase,
+    PolicyDefinitionGetResponse,
+    PolicyDefinitionId,
+)
 
 QoSScheduling = Literal[
     "llq",
@@ -19,11 +24,17 @@ QoSDropType = Literal[
 
 
 class QoSScheduler(BaseModel):
-    queue: str
-    class_map_ref: Union[UUID, Literal[""]] = Field(serialization_alias="classMapRef", validation_alias="classMapRef")
-    bandwidth_percent: str = Field("1", serialization_alias="bandwidthPercent", validation_alias="bandwidthPercent")
-    buffer_percent: str = Field("1", serialization_alias="bufferPercent", validation_alias="bufferPercent")
-    burst: Optional[str] = None
+    queue: IntStr = Field(ge=0, le=7)
+    class_map_ref: Optional[UUID] = Field(
+        default=None, serialization_alias="classMapRef", validation_alias="classMapRef"
+    )
+    bandwidth_percent: IntStr = Field(
+        default=1, ge=1, le=100, serialization_alias="bandwidthPercent", validation_alias="bandwidthPercent"
+    )
+    buffer_percent: IntStr = Field(
+        default=1, ge=1, le=100, serialization_alias="bufferPercent", validation_alias="bufferPercent"
+    )
+    burst: Optional[IntStr] = Field(default=None, ge=5000, le=10_000_000)
     scheduling: QoSScheduling = "wrr"
     drops: QoSDropType = "tail-drop"
     temp_key_values: Optional[str] = Field(
@@ -33,35 +44,23 @@ class QoSScheduler(BaseModel):
     @staticmethod
     def get_default_control_scheduler() -> "QoSScheduler":
         return QoSScheduler(
-            queue="0",
-            class_map_ref="",
-            bandwidth_percent="100",
-            buffer_percent="100",
-            burst="15000",
+            queue=0,
+            bandwidth_percent=100,
+            buffer_percent=100,
+            burst=15000,
             scheduling="llq",
             drops="tail-drop",
         )
 
     model_config = ConfigDict(populate_by_name=True)
 
-    @field_validator("queue")
+    @field_validator("class_map_ref", mode="before")
     @classmethod
-    def check_queue(cls, queue_str: str):
-        assert 0 <= int(queue_str) <= 7
-        return queue_str
-
-    @field_validator("bandwidth_percent", "buffer_percent")
-    @classmethod
-    def check_bandwidth_and_buffer_percent(cls, percent_str: str):
-        assert 1 <= int(percent_str) <= 100
-        return percent_str
-
-    @field_validator("burst")
-    @classmethod
-    def check_burst(cls, burst_val: Union[str, None]):
-        if burst_val is not None:
-            assert 5000 <= int(burst_val) <= 10_000_000
-        return burst_val
+    def check_optional_class_map_ref(cls, class_map_ref: Union[str, None]):
+        # None and "" indicates missing value, both can be found in server responses
+        if not class_map_ref:
+            return None
+        return class_map_ref
 
 
 class QoSMapDefinition(BaseModel):
@@ -86,11 +85,11 @@ class QoSMapPolicy(PolicyDefinitionBase):
     ) -> None:
         self.definition.qos_schedulers.append(
             QoSScheduler(
-                queue=str(queue),
+                queue=queue,
                 class_map_ref=class_map_ref,
-                bandwidth_percent=str(bandwidth),
-                buffer_percent=str(buffer),
-                burst=str(burst) if burst is not None else None,
+                bandwidth_percent=bandwidth,
+                buffer_percent=buffer,
+                burst=burst,
                 scheduling=scheduling,
                 drops=drops,
             )
@@ -102,3 +101,11 @@ class QoSMapPolicy(PolicyDefinitionBase):
             # Only when creating (not when value obtained from remote is present)
             self.definition = QoSMapDefinition(qos_schedulers=[QoSScheduler.get_default_control_scheduler()])
         return self
+
+
+class QoSMapPolicyEditPayload(QoSMapPolicy, PolicyDefinitionId):
+    pass
+
+
+class QoSMapPolicyGetResponse(QoSMapPolicy, PolicyDefinitionGetResponse):
+    pass

catalystwan/models/policy/{definitions → definition}/rewrite.py

@@ -9,6 +9,8 @@ from catalystwan.models.policy.policy_definition import (
     DefinitionWithSequencesCommonBase,
     PLPEntryType,
     PolicyDefinitionBase,
+    PolicyDefinitionGetResponse,
+    PolicyDefinitionId,
 )
 
 
@@ -20,7 +22,7 @@ class RewritePolicyRule(BaseModel):
     class_: UUID = Field(serialization_alias="class", validation_alias="class")
     plp: str
     dscp: str
-    l2cos: str = Field(serialization_alias="layer2Cos", validation_alias="layer2Cos")
+    l2cos: int = Field(serialization_alias="layer2Cos", validation_alias="layer2Cos")
     model_config = ConfigDict(populate_by_name=True)
 
 
@@ -32,6 +34,14 @@ class RewritePolicy(RewritePolicyHeader, DefinitionWithSequencesCommonBase):
     definition: RewritePolicyDefinition = RewritePolicyDefinition()
 
     def add_rule(self, class_map_ref: UUID, dscp: int, l2cos: int, plp: PLPEntryType) -> None:
-        self.definition.rules.append(RewritePolicyRule(class_=class_map_ref, plp=plp, dscp=str(dscp), l2cos=str(l2cos)))
+        self.definition.rules.append(RewritePolicyRule(class_=class_map_ref, plp=plp, dscp=str(dscp), l2cos=l2cos))
 
     model_config = ConfigDict(populate_by_name=True)
+
+
+class RewritePolicyEditPayload(RewritePolicy, PolicyDefinitionId):
+    pass
+
+
+class RewritePolicyGetResponse(RewritePolicy, PolicyDefinitionGetResponse):
+    pass