diff --git a/nebula/addons/attacks/communications/communicationattack.py b/nebula/addons/attacks/communications/communicationattack.py new file mode 100644 index 0000000..6830a9b --- /dev/null +++ b/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() + \ No newline at end of file diff --git a/nebula/addons/attacks/communications/delayerattack.py b/nebula/addons/attacks/communications/delayerattack.py index 61c25b8..8860f0b 100644 --- a/nebula/addons/attacks/communications/delayerattack.py +++ b/nebula/addons/attacks/communications/delayerattack.py @@ -1,18 +1,15 @@ 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. @@ -20,60 +17,40 @@ def __init__(self, engine, attack_params): 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() \ No newline at end of file + return decorator \ No newline at end of file