Merge pull request #6 from ggleszczynski/features/structure

Project classes structure

bt_ddos_shield/address.py

@@ -0,0 +1,54 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+
+
+class AddressType(Enum):
+    """
+    Possible types of address.
+    """
+    IP = "ip"         # IPv4 address
+    IPV6 = "ipv6"     # IPv6 address
+    DOMAIN = "domain" # domain name
+
+
+class Address(ABC):
+    """
+    Class describing address, which redirects to original miner's server.
+    """
+
+    def __init__(self, address_id, address_type: AddressType, address: str, port: int):
+        """
+        Args:
+            address_id: Identifier (used by AddressManager) of the address. Type depends on the implementation.
+            address_type: Type of the address.
+            address: Address.
+            port: Port of the address.
+        """
+        self.address_id = address_id
+        self.address_type = address_type
+        self.address = address
+        self.port = port
+
+    @abstractmethod
+    def encrypt(self) -> bytes:
+        """
+        Encrypts address data.
+
+        Returns:
+            bytes: Encrypted address data.
+        """
+        pass
+
+    @classmethod
+    @abstractmethod
+    def decrypt(cls, encrypted_data: bytes) -> Address:
+        """
+        Create address from encrypted address data.
+
+        Args:
+            encrypted_data: Encrypted address data.
+
+        Returns:
+            Address: Created address.
+        """
+        pass

bt_ddos_shield/address_manager.py

@@ -0,0 +1,57 @@
+from abc import ABC, abstractmethod
+
+from bt_ddos_shield.address import Address
+
+
+class AbstractAddressManager(ABC):
+    """
+    Abstract base class for manager handling public IP/domain addresses assigned to validators.
+    """
+
+    def __init__(self, address_id):
+        """
+        Args:
+            address_id: Identifier of the address of original miner's server. All created addresses for validators
+                        should redirect to this address.
+        """
+        self.address_id = address_id
+
+    @abstractmethod
+    def create_address(self) -> Address:
+        """
+        Create a new address.
+
+        Returns:
+            Address: New address to be used by validator.
+        """
+        pass
+
+    @abstractmethod
+    def remove_address(self, address_id):
+        """
+        Remove address.
+
+        Args:
+            address_id: Identifier of the address to remove.
+        """
+        pass
+
+    @abstractmethod
+    def address_exists(self, address_id) -> bool:
+        """
+        Check if address exists.
+
+        Args:
+            address_id: Identifier of the address to check.
+
+            Returns:
+                bool: If address exists.
+        """
+        pass
+
+    def hide_original_server(self):
+        """
+        If method is implemented, it should hide the original server IP address from public access.
+        See auto_hide_original_server in MinerShield options.
+        """
+        pass

bt_ddos_shield/blockchain_manager.py

@@ -0,0 +1,17 @@
+from abc import ABC, abstractmethod
+
+
+class AbstractBlockchainManager(ABC):
+    """
+    Abstract base class for manager handling publishing blocks to blockchain.
+    """
+
+    @abstractmethod
+    def publish(self, data: bytes):
+        """
+        Puts data to blockchain.
+
+        Args:
+            data: Data.
+        """
+        pass

bt_ddos_shield/event_processor.py

@@ -0,0 +1,28 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+
+@dataclass
+class Event:
+    """
+    Class describing event, which happened in the shield.
+    """
+
+    event_description: str       # Description of the event.
+    exception: Exception = None  # Exception which caused the event.
+
+
+class AbstractEventProcessor(ABC):
+    """
+    Abstract base class for processor handling events generated by shield.
+    """
+
+    @abstractmethod
+    def add_event(self, event: Event):
+        """
+        Add new event to be handled by processor.
+
+        Args:
+            event: Event to add.
+        """
+        pass

bt_ddos_shield/manifest_manager.py

@@ -0,0 +1,36 @@
+from abc import ABC, abstractmethod
+
+from bt_ddos_shield.address import Address
+from bt_ddos_shield.miner_shield import Hotkey
+
+
+class AbstractManifestManager(ABC):
+    """
+    Abstract base class for manager handling publishing manifest file containing encrypted addresses for validators.
+    """
+
+    def add_mapping_file(self, address_mapping: dict[Hotkey, Address]) -> Address:
+        """
+        Adds a mapping as file with encrypted addresses to the storage.
+
+        Args:
+            address_mapping: A dictionary containing the address mapping (validator HotKey -> Address).
+
+        Returns:
+            Address: Address where file is put.
+        """
+        # TODO - add implementation (encrypt with EncryptionManager and call put_file)
+        pass
+
+    @abstractmethod
+    def _put_file(self, data: bytes) -> Address:
+        """
+        Puts a file into the storage.
+
+        Args:
+            data: File contents.
+
+        Returns:
+            Address: Address where file is put.
+        """
+        pass

bt_ddos_shield/miner_shield.py

@@ -0,0 +1,74 @@
+from dataclasses import dataclass
+
+from bt_ddos_shield.blockchain_manager import AbstractBlockchainManager
+from bt_ddos_shield.event_processor import AbstractEventProcessor
+from bt_ddos_shield.address_manager import AbstractAddressManager
+from bt_ddos_shield.validators_manager import AbstractValidatorsManager
+from bt_ddos_shield.manifest_manager import AbstractManifestManager
+from bt_ddos_shield.state_manager import AbstractMinerShieldStateManager
+
+
+Hotkey = str # type of Hotkey
+
+
+@dataclass
+class MinerShieldOptions:
+    """
+        A class to represent the configuration options for the MinerShield.
+    """
+
+    auto_hide_original_server: bool = False # If True, the original server will be hidden after some time after shield
+                                            # gets enabled. Method hide_original_server in AddressManager will be called.
+
+    auto_hide_delay_sec: int = 600          # Time in seconds after which the original server will be hidden if
+                                            # auto_hide_original_server is set to True.
+
+
+class MinerShield:
+    """
+    Main class to be used by Miner to shield himself from DDoS. Call enable() to start the shield.
+    """
+
+    def __init__(self, validators_manager: AbstractValidatorsManager, address_manager: AbstractAddressManager,
+                 manifest_manager: AbstractManifestManager, blockchain_manager: AbstractBlockchainManager,
+                 state_manager: AbstractMinerShieldStateManager, event_processor: AbstractEventProcessor,
+                 options: MinerShieldOptions):
+        """
+        Initialize the MinerShield class.
+
+        Args:
+            validators_manager: Instance of AbstractValidatorsManager to manage validators and their keys.
+            address_manager: Instance of AbstractAddressManager to manage public IP/domain addresses assigned to validators.
+            manifest_manager: Instance of AbstractManifestManager to manage publishing manifest file.
+            blockchain_manager: Instance of AbstractBlockchainManager to manage blockchain operations.
+            state_manager: Instance of AbstractMinerShieldStateManager to manage state of the shield.
+            event_processor: Instance of AbstractEventProcessor to handle events generated by the shield.
+            options: Instance of MinerShieldOptions.
+        """
+        pass
+
+    def enable(self):
+        """
+        Enable shield. It asynchronously starts the shield, which consists of such steps:
+        1. Fetch validators keys.
+        2. Creates addresses for all validators.
+        3. Save manifest file.
+        4. Publish link to manifest file to blockchain.
+        5. Eventually close public access to original IP after some time.
+
+        It puts events to event_manager after each step. Current state is managed by state_manager. If shielding
+        process had been interrupted it is continued from the last step.
+
+        When shield is running, changing validators set triggers shield reconfiguration.
+        """
+        pass
+
+    def ban_validator(self, validator_hotkey: Hotkey):
+        """
+        Ban a validator by its hotkey. Function blocks execution until manifest file is updated and info about file
+        is published to Bittensor.
+
+        Args:
+            validator_hotkey: The hotkey of the validator.
+        """
+        pass

bt_ddos_shield/state_manager.py

@@ -0,0 +1,78 @@
+from abc import ABC, abstractmethod
+from datetime import datetime
+from enum import Enum
+
+from bt_ddos_shield.address import Address
+from bt_ddos_shield.miner_shield import Hotkey
+
+
+class MinerShieldPhase(Enum):
+    """
+    Possible phases of shield.
+    """
+    DISABLED = "disabled"                         # disabled - initial state
+    MANIFEST_PUBLISHED = "manifest_published"     # manifest is saved to storage
+    MANIFEST_BROADCASTED = "manifest_broadcasted" # info about manifest is published to blockchain
+    ENABLED = "enabled"                           # shield is enabled
+
+
+class MinerShieldState:
+    """
+    Class representing state of MinerShield.
+    """
+
+    phase: MinerShieldPhase                   # current phase of the shield
+    banned_validators: dict[Hotkey, datetime] # banned validators with ban time (HotKey -> time of ban)
+    active_addresses: dict[Hotkey, Address]   # active addresses (validator HotKey -> Address created for him)
+
+    def __init__(self):
+        self.phase = MinerShieldPhase.DISABLED
+        self.banned_validators = {}
+        self.active_addresses = {}
+
+
+class AbstractMinerShieldStateManager(ABC):
+    """
+    Abstract base class for manager handling state of MinerShield.
+    """
+
+    current_miner_shield_state: MinerShieldState
+
+    @abstractmethod
+    def load_state(self):
+        pass
+
+    @abstractmethod
+    def save_state(self):
+        pass
+
+    @abstractmethod
+    def ban_validator(self, validator_hotkey: Hotkey):
+        """
+        Add validator to the list of banned validators.
+
+        Args:
+            validator_hotkey: The hotkey of the validator.
+        """
+        pass
+
+    @abstractmethod
+    def remove_validator(self, validator_hotkey: Hotkey):
+        """
+        Remove validator from the lists of banned validators or active addresses.
+
+        Args:
+            validator_hotkey: The hotkey of the validator.
+        """
+        pass
+
+    @abstractmethod
+    def add_address(self, validator_hotkey: Hotkey, address: Address):
+        """
+        Add new address to state.
+
+        Args:
+            validator_hotkey: The hotkey of the validator.
+            address: Address to add.
+        """
+        pass

bt_ddos_shield/validators_manager.py

@@ -0,0 +1,26 @@
+from abc import ABC, abstractmethod
+
+from bt_ddos_shield.miner_shield import Hotkey
+
+
+class AbstractValidatorsManager(ABC):
+    """
+    Abstract base class for manager of validators and their public keys used for encryption.
+    """
+
+    @abstractmethod
+    def get_validators(self) -> dict[Hotkey, str]:
+        """
+        Get cached dictionary of validators - maps HotKey of validator to public key.
+        """
+        pass
+
+    @abstractmethod
+    def refresh_validators(self) -> bool:
+        """
+        Refresh validators dictionary. Blocks code execution until new validators set is fetched.
+
+        Returns:
+            bool: True if validators set is different after refreshing.
+        """
+        pass