refactor: added communicationattacks class

Author: FerTV

nebula/addons/attacks/communications/communicationattack.py

@@ -0,0 +1,51 @@
+from abc import abstractmethod
+import logging
+import types
+from nebula.addons.attacks.attacks import Attack
+
+
+class CommunicationAttack(Attack):
+    def __init__(self, engine, target_class, target_method, round_start_attack, round_stop_attack, decorator_args=None):
+        super().__init__()
+        self.engine = engine
+        self.target_class = target_class
+        self.target_method = target_method
+        self.decorator_args = decorator_args
+        self.round_start_attack = round_start_attack
+        self.round_stop_attack = round_stop_attack
+        self.original_method = getattr(target_class, target_method, None)
+
+        if not self.original_method:
+            raise AttributeError(f"Method {target_method} not found in class {target_class}")
+
+    @abstractmethod
+    def decorator(self, *args):
+        """Decorator that adds malicious behavior to the execution of the original method."""
+        pass
+
+    async def _inject_malicious_behaviour(self):
+        """Inject malicious behavior into the target method."""
+        logging.info("Injecting malicious behavior")
+        
+        decorated_method = self.decorator(self.decorator_args)(self.original_method)
+        
+        setattr(
+            self.target_class,
+            self.target_method,
+            types.MethodType(decorated_method, self.target_class),
+        )
+
+    async def _restore_original_behaviour(self):
+        """Restore the original behavior of the target method."""
+        logging.info(f"Restoring original behavior of {self.target_class}.{self.target_method}")
+        setattr(self.target_class, self.target_method, self.original_method)
+
+    async def attack(self):
+        """Perform the attack logic based on the current round."""
+        if self.engine.round == self.round_stop_attack:
+            logging.info(f"[{self.__class__.__name__}] Restoring original behavior")
+            await self._restore_original_behaviour()
+        elif self.engine.round == self.round_start_attack:
+            logging.info(f"[{self.__class__.__name__}] Injecting malicious behavior")
+            await self._inject_malicious_behaviour()
+        

nebula/addons/attacks/communications/delayerattack.py

@@ -1,79 +1,56 @@
 import asyncio
 from functools import wraps
 import logging
-import types
-from nebula.addons.attacks.attacks import Attack
+from nebula.addons.attacks.communications.communicationattack import CommunicationAttack
 
 
-class DelayerAttack(Attack):
+class DelayerAttack(CommunicationAttack):
     """
     Implements an attack that delays the execution of a target method by a specified amount of time.
-
-    This attack dynamically modifies the `propagate` method of the propagator to
-    introduce a delay during its execution.
     """
-    def __init__(self, engine, attack_params):
+
+    def __init__(self, engine, attack_params: dict):
         """
         Initializes the DelayerAttack with the engine and attack parameters.
 
         Args:
             engine: The engine managing the attack context.
             attack_params (dict): Parameters for the attack, including the delay duration.
         """
-        super().__init__()
-        self.engine = engine
-        self.propagator = self.engine._cm._propagator
-        self.original_propagate = self.propagator.propagate
-        self.delay = int(attack_params["delay"])
-        self.round_start_attack = int(attack_params["round_start_attack"])
-        self.round_stop_attack = int(attack_params["round_stop_attack"])
-
-    def delay_decorator(self, delay):
+        try:
+            self.delay = int(attack_params["delay"])
+            round_start = int(attack_params["round_start_attack"])
+            round_stop = int(attack_params["round_stop_attack"])
+        except KeyError as e:
+            raise ValueError(f"Missing required attack parameter: {e}")
+        except ValueError:
+            raise ValueError("Invalid value in attack_params. Ensure all values are integers.")
+
+        super().__init__(
+            engine,
+            engine._cm._propagator,
+            "propagate",
+            round_start,
+            round_stop,
+            self.delay,
+        )
+
+    def decorator(self, delay: int):
         """
         Decorator that adds a delay to the execution of the original method.
 
         Args:
-            delay (int or float): The time in seconds to delay the method execution.
+            delay (int): The time in seconds to delay the method execution.
 
         Returns:
             function: A decorator function that wraps the target method with the delay logic.
         """
-        # The actual decorator function that will be applied to the target method
         def decorator(func):
-            @wraps(func)  # Preserves the metadata of the original function
-            async def wrapper(*args):
-                logging.info(f"[DelayerAttack] Adding delay of {delay} seconds")
-
+            @wraps(func)
+            async def wrapper(*args, **kwargs):
+                logging.info(f"[DelayerAttack] Adding delay of {delay} seconds to {func.__name__}")
                 await asyncio.sleep(delay)
                 _, *new_args = args  # Exclude self argument
                 return await func(*new_args)
             return wrapper
-        return decorator
-
-    async def _inject_malicious_behaviour(self):
-        """
-        Modifies the `propagate` method of the propagator to include a delay.
-        """
-        decorated_propagate = self.delay_decorator(self.delay)(self.propagator.propagate)
-
-        self.propagator.propagate = types.MethodType(decorated_propagate, self.propagator)
-
-    async def _restore_original_behaviour(self):
-        """
-        Restores the original behaviour of the `propagate` method.
-        """
-        self.propagator.propagate = self.original_propagate
-
-    async def attack(self):
-        """
-        Starts the attack by injecting the malicious behaviour.
-
-        If the current round matches the attack start round, the malicious behavior
-        is injected. If it matches the stop round, the original behavior is restored.
-        """
-        if self.engine.round == self.round_stop_attack:
-            logging.info(f"[DelayerAttack] Stopping Delayer attack")
-            await self._restore_original_behaviour()
-        elif self.engine.round == self.round_start_attack:
-            logging.info("[DelayerAttack] Injecting malicious behaviour")
-            await self._inject_malicious_behaviour()
+        return decorator