diff --git a/bt_ddos_shield/address.py b/bt_ddos_shield/address.py new file mode 100644 index 0000000..0209b53 --- /dev/null +++ b/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 diff --git a/bt_ddos_shield/address_manager.py b/bt_ddos_shield/address_manager.py new file mode 100644 index 0000000..14e67c4 --- /dev/null +++ b/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 diff --git a/bt_ddos_shield/blockchain_manager.py b/bt_ddos_shield/blockchain_manager.py new file mode 100644 index 0000000..9d926f2 --- /dev/null +++ b/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 diff --git a/bt_ddos_shield/event_processor.py b/bt_ddos_shield/event_processor.py new file mode 100644 index 0000000..b398942 --- /dev/null +++ b/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 diff --git a/bt_ddos_shield/manifest_manager.py b/bt_ddos_shield/manifest_manager.py new file mode 100644 index 0000000..4960535 --- /dev/null +++ b/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 diff --git a/bt_ddos_shield/miner.py b/bt_ddos_shield/miner.py deleted file mode 100644 index b5884bf..0000000 --- a/bt_ddos_shield/miner.py +++ /dev/null @@ -1,6 +0,0 @@ -class Miner: - def __init__(self): - pass - - def encrypt_and_publish(self): - pass diff --git a/bt_ddos_shield/miner_shield.py b/bt_ddos_shield/miner_shield.py new file mode 100644 index 0000000..5c36f7f --- /dev/null +++ b/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 diff --git a/bt_ddos_shield/s3_manager.py b/bt_ddos_shield/s3_manager.py deleted file mode 100644 index 08ebeee..0000000 --- a/bt_ddos_shield/s3_manager.py +++ /dev/null @@ -1,9 +0,0 @@ -class S3Manager: - def __init__(self): - pass - - def upload_file(self): - pass - - def download_file(self): - pass diff --git a/bt_ddos_shield/state_manager.py b/bt_ddos_shield/state_manager.py new file mode 100644 index 0000000..c313946 --- /dev/null +++ b/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 diff --git a/bt_ddos_shield/validators_manager.py b/bt_ddos_shield/validators_manager.py new file mode 100644 index 0000000..d1c2268 --- /dev/null +++ b/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