added proper validation to ensure security

Author: Andrei Neagu

packages/models-library/src/models_library/api_schemas_dynamic_sidecar/telemetry.py

@@ -1,8 +1,17 @@
 from abc import abstractmethod
-from typing import Protocol
+from typing import Any, Final, Protocol
 
 from models_library.projects_nodes_io import NodeID
-from pydantic import BaseModel, ByteSize, Field
+from pydantic import (
+    BaseModel,
+    ByteSize,
+    Field,
+    NonNegativeFloat,
+    root_validator,
+    validator,
+)
+
+_EPSILON: Final[NonNegativeFloat] = 1e-16
 
 
 class SDiskUsageProtocol(Protocol):
@@ -27,28 +36,70 @@ def percent(self) -> float:
         ...
 
 
+def _get_percent(used: float, total: float) -> float:
+    return round(used * 100 / (total + _EPSILON), 2)
+
+
 class DiskUsage(BaseModel):
     used: ByteSize = Field(description="used space")
     free: ByteSize = Field(description="remaining space")
 
     total: ByteSize = Field(description="total space = free + used")
-    used_percent: float = Field(
+    used_percent: NonNegativeFloat = Field(
         gte=0.00,
         lte=100.00,
         description="Percent of used space relative to the total space",
     )
 
+    @validator("free")
+    @classmethod
+    def _free_positive(cls, v: float) -> float:
+        if v < 0:
+            msg = f"free={v} cannot be a negative value"
+            raise ValueError(msg)
+        return v
+
+    @validator("used")
+    @classmethod
+    def _used_positive(cls, v: float) -> float:
+        if v < 0:
+            msg = f"used={v} cannot be a negative value"
+            raise ValueError(msg)
+        return v
+
+    @root_validator(pre=True)
+    @classmethod
+    def _check_total(cls, values: dict[str, Any]) -> dict[str, Any]:
+        total = values["total"]
+        free = values["free"]
+        used = values["used"]
+        if total != free + used:
+            msg = f"{total=} is different than the sum of {free=}+{used=} => sum={free+used}"
+            raise ValueError(msg)
+        return values
+
     @classmethod
     def from_ps_util_disk_usage(
         cls, ps_util_disk_usage: SDiskUsageProtocol
     ) -> "DiskUsage":
         total = ps_util_disk_usage.free + ps_util_disk_usage.used
-        used_percent = round(ps_util_disk_usage.used * 100 / total, 2)
         return cls(
             used=ByteSize(ps_util_disk_usage.used),
             free=ByteSize(ps_util_disk_usage.free),
             total=ByteSize(total),
-            used_percent=used_percent,
+            used_percent=_get_percent(ps_util_disk_usage.used, total),
+        )
+
+    @classmethod
+    def from_efs_guardian(
+        cls, used: NonNegativeFloat, total: NonNegativeFloat
+    ) -> "DiskUsage":
+        free = total - used
+        return cls(
+            used=ByteSize(used),
+            free=ByteSize(free),
+            total=ByteSize(total),
+            used_percent=_get_percent(used, total),
         )
 
     def __hash__(self):

packages/models-library/tests/test_api_schemas_dynamic_sidecar_telemetry.py

@@ -2,6 +2,7 @@
 import pytest
 from models_library.api_schemas_dynamic_sidecar.telemetry import DiskUsage
 from psutil._common import sdiskusage
+from pydantic import ByteSize, ValidationError
 
 
 def _assert_same_value(ps_util_disk_usage: sdiskusage) -> None:
@@ -27,3 +28,35 @@ def test_disk_usage_regression_cases(ps_util_disk_usage: sdiskusage):
 def test_disk_usage():
     ps_util_disk_usage = psutil.disk_usage("/")
     _assert_same_value(ps_util_disk_usage)
+
+
+def test_from_efs_guardian_constructor():
+    result = DiskUsage.from_efs_guardian(10, 100)
+    assert result.used == ByteSize(10)
+    assert result.free == ByteSize(90)
+    assert result.total == ByteSize(100)
+    assert result.used_percent == 10
+
+
+def test_failing_validation():
+    with pytest.raises(ValidationError) as exc:
+        assert DiskUsage.from_efs_guardian(100, 10)
+
+    assert "free=" in f"{exc.value}"
+    assert "negative value" in f"{exc.value}"
+
+    with pytest.raises(ValidationError) as exc:
+        assert DiskUsage(
+            used=-10,  # type: ignore
+            free=ByteSize(10),
+            total=ByteSize(0),
+            used_percent=-10,
+        )
+    assert "used=" in f"{exc.value}"
+    assert "negative value" in f"{exc.value}"
+
+    with pytest.raises(ValidationError) as exc:
+        DiskUsage(
+            used=ByteSize(10), free=ByteSize(10), total=ByteSize(21), used_percent=0
+        )
+    assert "is different than the sum of" in f"{exc.value}"