Merge branch '3411-percent-type' into develop

Issue #3411
PR #3558

Author: mssalvatore

monkey/agent_plugins/payloads/cryptojacker/src/cryptojacker_options.py

@@ -1,6 +1,7 @@
-from pydantic import Field, conint
+from pydantic import Field
 
 from common.base_models import InfectionMonkeyBaseModel
+from common.types import PercentLimited
 
 
 class CryptojackerOptions(InfectionMonkeyBaseModel):
@@ -11,12 +12,12 @@ class CryptojackerOptions(InfectionMonkeyBaseModel):
         default=300,  # 5 minutes
         ge=0,
     )
-    cpu_utilization: conint(ge=0, le=100) = Field(  # type: ignore[valid-type]
+    cpu_utilization: PercentLimited = Field(  # type: ignore[valid-type]
         title="CPU utilization",
         description="The percentage of CPU to use on a machine",
         default=80,
     )
-    memory_utilization: conint(ge=0, le=100) = Field(  # type: ignore[valid-type]
+    memory_utilization: PercentLimited = Field(  # type: ignore[valid-type]
         title="Memory utilization",
         description="The percentage of memory to use on a machine",
         default=20,

monkey/common/types/__init__.py

@@ -5,3 +5,4 @@
 from .networking import NetworkService, NetworkPort, PortStatus, SocketAddress, NetworkProtocol
 from .secrets import OTP, Token
 from .file_extension import FileExtension
+from .percent import Percent, PercentLimited

monkey/common/types/percent.py

@@ -0,0 +1,68 @@
+from typing import Any, Self
+
+from pydantic import NonNegativeFloat
+
+
+class Percent(NonNegativeFloat):
+    """
+    A type representing a percentage
+
+    Note that percentages can be greater than 100. For example, I may have consumed 120% of my quota
+    (if quotas aren't strictly enforced).
+    """
+
+    # This __init__() is required so that instances of Percent can be created. If you try to create
+    # an instance of NonNegativeFloat, no validation is performed.
+    def __init__(self, v: Any):
+        Percent._validate_range(v)
+
+    @classmethod
+    def __get_validators__(cls):
+        for v in super().__get_validators__():
+            yield v
+
+        yield cls.validate
+
+    @classmethod
+    def validate(cls, v: Any) -> Self:
+        cls._validate_range(v)
+
+        # This is required so that floats passed into pydantic models are converted to instances of
+        # Percent objects.
+        return cls(v)
+
+    @staticmethod
+    def _validate_range(v: Any):
+        if v < 0:
+            raise ValueError("value must be non-negative")
+
+    def as_decimal_fraction(self) -> NonNegativeFloat:
+        """
+        Return the percentage as a decimal fraction
+
+        Example: 50% -> 0.5
+
+        return: The percentage as a decimal fraction
+        """
+        return self / 100.0
+
+
+class PercentLimited(Percent):
+    """
+    A type representing a percentage limited to 100
+    """
+
+    le = 100
+
+    def __init__(self, v: Any):
+        PercentLimited._validate_range(v)
+
+    @classmethod
+    def __get_validators__(cls):
+        for v in super().__get_validators__():
+            yield v
+
+    @staticmethod
+    def _validate_range(v: Any):
+        if not (0.0 <= v <= 100.0):
+            raise ValueError("value must be between 0 and 100")

monkey/tests/unit_tests/agent_plugins/payloads/ransomware/cryptojacker/test_cryptojacker_options.py monkey/tests/unit_tests/agent_plugins/payloads/cryptojacker/test_cryptojacker_options.py

@@ -0,0 +1,53 @@
+import pytest
+from pydantic import BaseModel
+
+from common.types import Percent, PercentLimited
+
+
+class Model(BaseModel):
+    p: Percent
+
+
+def test_non_negative_percent():
+    with pytest.raises(ValueError):
+        Percent(-1.0)
+
+
+def test_incorrect_type():
+    with pytest.raises(ValueError):
+        Percent("stuff")
+
+
+@pytest.mark.parametrize("input_value", [0.0, 1.0, 99.9, 100.0, 120.5, 50.123, 25])
+def test_valid_percent(input_value: float):
+    assert Percent(input_value) == input_value  # type: ignore [arg-type]
+
+
+@pytest.mark.parametrize(
+    "input_value,expected_decimal_fraction",
+    [(0.0, 0.0), (99, 0.99), (51.234, 0.51234), (121.2, 1.212)],
+)
+def test_as_decimal_fraction(input_value: float, expected_decimal_fraction: float):
+    assert Model(p=input_value).p.as_decimal_fraction() == expected_decimal_fraction  # type: ignore [arg-type]  # noqa: E501
+
+
+class ModelLimited(BaseModel):
+    p: PercentLimited
+
+
+@pytest.mark.parametrize("input_value", (-1, -0.01, 100.01))
+def test_percent_limited_out_of_range(input_value: float):
+    with pytest.raises(ValueError):
+        PercentLimited(input_value)
+
+
+@pytest.mark.parametrize("input_value", [0.0, 1.0, 99.9, 100.0, 50.123, 25])
+def test_valid_percent_limited(input_value: float):
+    assert PercentLimited(input_value) == input_value  # type: ignore [arg-type]
+
+
+@pytest.mark.parametrize(
+    "input_value,expected_decimal_fraction", [(0.0, 0.0), (99, 0.99), (51.234, 0.51234)]
+)
+def test_as_decimal_fraction_limited(input_value: float, expected_decimal_fraction: float):
+    assert ModelLimited(p=input_value).p.as_decimal_fraction() == expected_decimal_fraction  # type: ignore [arg-type]  # noqa: E501

monkey/tests/unit_tests/common/types/test_percent.py

@@ -31,6 +31,7 @@
 from infection_monkey.network.firewall import FirewallApp, WinAdvFirewall, WinFirewall
 from infection_monkey.utils import commands
 from infection_monkey.utils.decorators import request_cache
+from monkey.common.types import Percent
 from monkey_island.cc.deployment import Deployment
 from monkey_island.cc.models import Machine
 from monkey_island.cc.repositories import IAgentEventRepository, MongoAgentEventRepository
@@ -188,3 +189,6 @@
 AgentPluginRepositoryIndex._infection_monkey_version_parser
 AgentPluginRepositoryIndex._sort_plugins_by_version
 AgentPluginRepositoryIndex.use_enum_values
+
+# Remove when cryptojacker is completed
+Percent.as_decimal_fraction