From 6d96fec1731073be21af4b4c1c7dd9c9d57bcc18 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Grzegorz=20Leszczy=C5=84ski?= Date: Tue, 29 Oct 2024 13:10:51 +0100 Subject: [PATCH 1/2] feat: add project structure --- bt_ddos_shield/blockchain_manager.py | 19 +++++++ bt_ddos_shield/dns_manager.py | 52 +++++++++++++++++ bt_ddos_shield/domain.py | 50 +++++++++++++++++ bt_ddos_shield/event_manager.py | 36 ++++++++++++ bt_ddos_shield/manifest_manager.py | 38 +++++++++++++ bt_ddos_shield/miner.py | 6 -- bt_ddos_shield/miner_shield.py | 73 ++++++++++++++++++++++++ bt_ddos_shield/s3_manager.py | 9 --- bt_ddos_shield/state_manager.py | 83 ++++++++++++++++++++++++++++ bt_ddos_shield/validators_manager.py | 26 +++++++++ 10 files changed, 377 insertions(+), 15 deletions(-) create mode 100644 bt_ddos_shield/blockchain_manager.py create mode 100644 bt_ddos_shield/dns_manager.py create mode 100644 bt_ddos_shield/domain.py create mode 100644 bt_ddos_shield/event_manager.py create mode 100644 bt_ddos_shield/manifest_manager.py delete mode 100644 bt_ddos_shield/miner.py create mode 100644 bt_ddos_shield/miner_shield.py delete mode 100644 bt_ddos_shield/s3_manager.py create mode 100644 bt_ddos_shield/state_manager.py create mode 100644 bt_ddos_shield/validators_manager.py diff --git a/bt_ddos_shield/blockchain_manager.py b/bt_ddos_shield/blockchain_manager.py new file mode 100644 index 0000000..5ec895f --- /dev/null +++ b/bt_ddos_shield/blockchain_manager.py @@ -0,0 +1,19 @@ +from abc import ABC, abstractmethod + +class BlockchainManager(ABC): + """ + Abstract base class for manager handling publishing blocks to blockchain. + """ + + def __init__(self): + pass + + @abstractmethod + def publish(self, data: bytes): + """ + Puts data to blockchain. + + Args: + data: Data. + """ + pass diff --git a/bt_ddos_shield/dns_manager.py b/bt_ddos_shield/dns_manager.py new file mode 100644 index 0000000..e488e88 --- /dev/null +++ b/bt_ddos_shield/dns_manager.py @@ -0,0 +1,52 @@ +from abc import ABC, abstractmethod + +from bt_ddos_shield.domain import Domain + + +class DNSManager(ABC): + """ + Abstract base class for manager handling public IP/domain addresses assigned to validators. + """ + + def __init__(self): + pass + + @abstractmethod + def create_domain(self) -> Domain: + """ + Create a new domain. + + Returns: + Domain: New domain. + """ + pass + + @abstractmethod + def remove_domain(self, domain_id): + """ + Remove domain. + + Args: + domain_id: Identifier of the domain to remove. + """ + pass + + @abstractmethod + def domain_exists(self, domain_id): + """ + Check if domain exists. + + Args: + domain_id: Identifier of the domain to check. + + Returns: + bool: If domain 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/domain.py b/bt_ddos_shield/domain.py new file mode 100644 index 0000000..57eda8a --- /dev/null +++ b/bt_ddos_shield/domain.py @@ -0,0 +1,50 @@ +from abc import ABC, abstractmethod +from enum import Enum + + +class Domain(ABC): + """ + Class describing domain. + """ + + class DomainType(Enum): + """ + Possible types of domain. + """ + IP = "ip" # domain is an IPv4 address + IPV6 = "ipv6" # domain is an IPv6 address + DOMAIN = "domain" # domain is a domain name + + def __init__(self, domain_id, domain_type: DomainType, address: str, port: int): + """ + Args: + domain_id: Identifier (used by DNSManager) of the domain. Type depends on the implementation. + domain_type: Type of the domain. + address: Address of the domain. + port: Port of the domain. + """ + self.domain_id = domain_id + self.domain_type = domain_type + self.address = address + self.port = port + + @abstractmethod + def encrypt(self) -> str: + """ + Encrypts domain data. + + Returns: + str: Encrypted domain data. + """ + pass + + @classmethod + @abstractmethod + def decrypt(cls, encrypted_data: str) -> Domain: + """ + Create domain from encrypted domain data. + + Returns: + Domain: Created domain. + """ + pass diff --git a/bt_ddos_shield/event_manager.py b/bt_ddos_shield/event_manager.py new file mode 100644 index 0000000..8ce8bc7 --- /dev/null +++ b/bt_ddos_shield/event_manager.py @@ -0,0 +1,36 @@ +from abc import ABC, abstractmethod + +class Event: + """ + Class describing event, which happened in the shield. + """ + + event_description: str + exception: Exception + + def __init__(self, event_description: str, exception: Exception = None): + """ + Args: + event_description: Description of the event. + exception: Exception which caused the event. + """ + self.event_description = event_description + self.exception = exception + +class EventManager(ABC): + """ + Abstract base class for manager handling events generated by shield. + """ + + def __init__(self): + pass + + @abstractmethod + def add_event(self, event: Event): + """ + Add new event to manager. + + 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..7562330 --- /dev/null +++ b/bt_ddos_shield/manifest_manager.py @@ -0,0 +1,38 @@ +from abc import ABC, abstractmethod + +from bt_ddos_shield.dns_manager import Domain + + +class ManifestManager(ABC): + """ + Abstract base class for manager handling publishing manifest file containing encrypted domains for validators. + """ + + def __init__(self): + pass + + def add_mapping_file(self, domain_mapping: dict[str, Domain]) -> Domain: + """ + Adds a mapping as file with encrypted domains to the storage. + + Args: + domain_mapping: A dictionary containing the domain mapping (validator HotKey -> Domain). + + Returns: + Domain: Domain where file is put. + """ + # TODO - add implementation (encrypt with EncryptionManager and call put_file) + pass + + @abstractmethod + def put_file(self, data: str) -> Domain: + """ + Puts a file into the storage. + + Args: + data: File contents. + + Returns: + Domain: Domain 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..8fbac6f --- /dev/null +++ b/bt_ddos_shield/miner_shield.py @@ -0,0 +1,73 @@ +from bt_ddos_shield.blockchain_manager import BlockchainManager +from bt_ddos_shield.event_manager import EventManager +from bt_ddos_shield.dns_manager import DNSManager +from bt_ddos_shield.validators_manager import ValidatorsManager +from bt_ddos_shield.manifest_manager import ManifestManager +from bt_ddos_shield.state_manager import StateManager + + +class MinerShieldOptions: + """ + A class to represent the configuration options for the MinerShield. + """ + + def __init__(self, auto_hide_original_server: bool = False, auto_hide_delay_sec: int = 600): + """ + Initialize the MinerShieldOptions class. + + Args: + auto_hide_original_server: If True, the original server will be hidden after some time after shield gets + enabled. Method hide_original_server in DNSManager will be called. + auto_hide_delay_sec: Time in seconds after which the original server will be hidden if + auto_hide_original_server is set to True. + """ + self.auto_hide_original_server = auto_hide_original_server + self.auto_hide_delay_sec = auto_hide_delay_sec + +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: ValidatorsManager, dns_manager: DNSManager, manifest_manager: ManifestManager, + blockchain_manager: BlockchainManager, state_manager: StateManager, event_manager: EventManager, + options: MinerShieldOptions): + """ + Initialize the MinerShield class. + + Args: + validators_manager: Instance of ValidatorsManager to manage validators and their keys. + dns_manager: Instance of DNSManager to manage public IP/domain addresses assigned to validators. + manifest_manager: Instance of ManifestManager to manage publishing manifest file. + blockchain_manager: Instance of BlockchainManager to manage blockchain operations. + state_manager: Instance of StateManager to manage state of the shield. + event_manager: Handles 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 IPs/domains 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: str): + """ + 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..dd5dfc6 --- /dev/null +++ b/bt_ddos_shield/state_manager.py @@ -0,0 +1,83 @@ +from abc import ABC, abstractmethod +from enum import Enum + +from bt_ddos_shield.domain import Domain + + +class State: + """ + Class representing state of MinerShield. + """ + + class ShieldPhase(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 + + phase: ShieldPhase # current phase of the shield + banned_validators: dict[str, int] # banned validators with ban time (HotKey -> ban timestamp) + active_domains: dict[str, Domain] # active domains (validator HotKey -> Domain created for him) + + def __init__(self): + self.phase = self.ShieldPhase.DISABLED + self.banned_validators = {} + self.active_domains = {} + +class StateManager(ABC): + """ + Abstract base class for manager handling state of MinerShield. + """ + + state: State # current state of MinerShield + + def __init__(self): + pass + + @abstractmethod + def load_state(self): + """ + Load current state. + """ + pass + + @abstractmethod + def save_state(self): + """ + Save current state. + """ + pass + + @abstractmethod + def ban_validator(self, validator_hotkey: str): + """ + Add validator to the list of banned validators. + + Args: + validator_hotkey: The hotkey of the validator. + """ + pass + + @abstractmethod + def remove_validator(self, validator_hotkey: str): + """ + Remove validator from the lists of banned validators or active domains. + + Args: + validator_hotkey: The hotkey of the validator. + """ + pass + + @abstractmethod + def add_domain(self, validator_hotkey: str, domain: Domain): + """ + Add new domain to state. + + Args: + validator_hotkey: The hotkey of the validator. + domain: Domain 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..1ff497a --- /dev/null +++ b/bt_ddos_shield/validators_manager.py @@ -0,0 +1,26 @@ +from abc import ABC, abstractmethod + +class ValidatorsManager(ABC): + """ + Abstract base class for manager of validators and their public keys used for encryption. + """ + + def __init__(self): + pass + + @abstractmethod + def get_validators(self) -> dict[str, 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 From 6ea94cb50912280908a93d4e6bdaa08b4ad2f539 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Grzegorz=20Leszczy=C5=84ski?= Date: Wed, 30 Oct 2024 14:42:51 +0100 Subject: [PATCH 2/2] fix: fixes after code review --- bt_ddos_shield/address.py | 54 ++++++++++++++++++++++++ bt_ddos_shield/address_manager.py | 57 ++++++++++++++++++++++++++ bt_ddos_shield/blockchain_manager.py | 14 +++---- bt_ddos_shield/dns_manager.py | 52 ------------------------ bt_ddos_shield/domain.py | 50 ----------------------- bt_ddos_shield/event_manager.py | 36 ---------------- bt_ddos_shield/event_processor.py | 28 +++++++++++++ bt_ddos_shield/manifest_manager.py | 22 +++++----- bt_ddos_shield/miner_shield.py | 55 +++++++++++++------------ bt_ddos_shield/state_manager.py | 61 +++++++++++++--------------- bt_ddos_shield/validators_manager.py | 10 ++--- 11 files changed, 216 insertions(+), 223 deletions(-) create mode 100644 bt_ddos_shield/address.py create mode 100644 bt_ddos_shield/address_manager.py delete mode 100644 bt_ddos_shield/dns_manager.py delete mode 100644 bt_ddos_shield/domain.py delete mode 100644 bt_ddos_shield/event_manager.py create mode 100644 bt_ddos_shield/event_processor.py 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 index 5ec895f..9d926f2 100644 --- a/bt_ddos_shield/blockchain_manager.py +++ b/bt_ddos_shield/blockchain_manager.py @@ -1,19 +1,17 @@ from abc import ABC, abstractmethod -class BlockchainManager(ABC): + +class AbstractBlockchainManager(ABC): """ Abstract base class for manager handling publishing blocks to blockchain. """ - def __init__(self): - pass - @abstractmethod def publish(self, data: bytes): """ - Puts data to blockchain. + Puts data to blockchain. - Args: - data: Data. - """ + Args: + data: Data. + """ pass diff --git a/bt_ddos_shield/dns_manager.py b/bt_ddos_shield/dns_manager.py deleted file mode 100644 index e488e88..0000000 --- a/bt_ddos_shield/dns_manager.py +++ /dev/null @@ -1,52 +0,0 @@ -from abc import ABC, abstractmethod - -from bt_ddos_shield.domain import Domain - - -class DNSManager(ABC): - """ - Abstract base class for manager handling public IP/domain addresses assigned to validators. - """ - - def __init__(self): - pass - - @abstractmethod - def create_domain(self) -> Domain: - """ - Create a new domain. - - Returns: - Domain: New domain. - """ - pass - - @abstractmethod - def remove_domain(self, domain_id): - """ - Remove domain. - - Args: - domain_id: Identifier of the domain to remove. - """ - pass - - @abstractmethod - def domain_exists(self, domain_id): - """ - Check if domain exists. - - Args: - domain_id: Identifier of the domain to check. - - Returns: - bool: If domain 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/domain.py b/bt_ddos_shield/domain.py deleted file mode 100644 index 57eda8a..0000000 --- a/bt_ddos_shield/domain.py +++ /dev/null @@ -1,50 +0,0 @@ -from abc import ABC, abstractmethod -from enum import Enum - - -class Domain(ABC): - """ - Class describing domain. - """ - - class DomainType(Enum): - """ - Possible types of domain. - """ - IP = "ip" # domain is an IPv4 address - IPV6 = "ipv6" # domain is an IPv6 address - DOMAIN = "domain" # domain is a domain name - - def __init__(self, domain_id, domain_type: DomainType, address: str, port: int): - """ - Args: - domain_id: Identifier (used by DNSManager) of the domain. Type depends on the implementation. - domain_type: Type of the domain. - address: Address of the domain. - port: Port of the domain. - """ - self.domain_id = domain_id - self.domain_type = domain_type - self.address = address - self.port = port - - @abstractmethod - def encrypt(self) -> str: - """ - Encrypts domain data. - - Returns: - str: Encrypted domain data. - """ - pass - - @classmethod - @abstractmethod - def decrypt(cls, encrypted_data: str) -> Domain: - """ - Create domain from encrypted domain data. - - Returns: - Domain: Created domain. - """ - pass diff --git a/bt_ddos_shield/event_manager.py b/bt_ddos_shield/event_manager.py deleted file mode 100644 index 8ce8bc7..0000000 --- a/bt_ddos_shield/event_manager.py +++ /dev/null @@ -1,36 +0,0 @@ -from abc import ABC, abstractmethod - -class Event: - """ - Class describing event, which happened in the shield. - """ - - event_description: str - exception: Exception - - def __init__(self, event_description: str, exception: Exception = None): - """ - Args: - event_description: Description of the event. - exception: Exception which caused the event. - """ - self.event_description = event_description - self.exception = exception - -class EventManager(ABC): - """ - Abstract base class for manager handling events generated by shield. - """ - - def __init__(self): - pass - - @abstractmethod - def add_event(self, event: Event): - """ - Add new event to manager. - - Args: - event: Event to add. - """ - 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 index 7562330..4960535 100644 --- a/bt_ddos_shield/manifest_manager.py +++ b/bt_ddos_shield/manifest_manager.py @@ -1,31 +1,29 @@ from abc import ABC, abstractmethod -from bt_ddos_shield.dns_manager import Domain +from bt_ddos_shield.address import Address +from bt_ddos_shield.miner_shield import Hotkey -class ManifestManager(ABC): +class AbstractManifestManager(ABC): """ - Abstract base class for manager handling publishing manifest file containing encrypted domains for validators. + Abstract base class for manager handling publishing manifest file containing encrypted addresses for validators. """ - def __init__(self): - pass - - def add_mapping_file(self, domain_mapping: dict[str, Domain]) -> Domain: + def add_mapping_file(self, address_mapping: dict[Hotkey, Address]) -> Address: """ - Adds a mapping as file with encrypted domains to the storage. + Adds a mapping as file with encrypted addresses to the storage. Args: - domain_mapping: A dictionary containing the domain mapping (validator HotKey -> Domain). + address_mapping: A dictionary containing the address mapping (validator HotKey -> Address). Returns: - Domain: Domain where file is put. + Address: Address where file is put. """ # TODO - add implementation (encrypt with EncryptionManager and call put_file) pass @abstractmethod - def put_file(self, data: str) -> Domain: + def _put_file(self, data: bytes) -> Address: """ Puts a file into the storage. @@ -33,6 +31,6 @@ def put_file(self, data: str) -> Domain: data: File contents. Returns: - Domain: Domain where file is put. + Address: Address where file is put. """ pass diff --git a/bt_ddos_shield/miner_shield.py b/bt_ddos_shield/miner_shield.py index 8fbac6f..5c36f7f 100644 --- a/bt_ddos_shield/miner_shield.py +++ b/bt_ddos_shield/miner_shield.py @@ -1,47 +1,48 @@ -from bt_ddos_shield.blockchain_manager import BlockchainManager -from bt_ddos_shield.event_manager import EventManager -from bt_ddos_shield.dns_manager import DNSManager -from bt_ddos_shield.validators_manager import ValidatorsManager -from bt_ddos_shield.manifest_manager import ManifestManager -from bt_ddos_shield.state_manager import StateManager +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. """ - def __init__(self, auto_hide_original_server: bool = False, auto_hide_delay_sec: int = 600): - """ - Initialize the MinerShieldOptions class. + 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. - Args: - auto_hide_original_server: If True, the original server will be hidden after some time after shield gets - enabled. Method hide_original_server in DNSManager will be called. - auto_hide_delay_sec: Time in seconds after which the original server will be hidden if - auto_hide_original_server is set to True. - """ - self.auto_hide_original_server = auto_hide_original_server - self.auto_hide_delay_sec = auto_hide_delay_sec 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: ValidatorsManager, dns_manager: DNSManager, manifest_manager: ManifestManager, - blockchain_manager: BlockchainManager, state_manager: StateManager, event_manager: EventManager, + 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 ValidatorsManager to manage validators and their keys. - dns_manager: Instance of DNSManager to manage public IP/domain addresses assigned to validators. - manifest_manager: Instance of ManifestManager to manage publishing manifest file. - blockchain_manager: Instance of BlockchainManager to manage blockchain operations. - state_manager: Instance of StateManager to manage state of the shield. - event_manager: Handles events generated by the shield. + 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 @@ -50,7 +51,7 @@ def enable(self): """ Enable shield. It asynchronously starts the shield, which consists of such steps: 1. Fetch validators keys. - 2. Creates IPs/domains for all validators. + 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. @@ -62,7 +63,7 @@ def enable(self): """ pass - def ban_validator(self, validator_hotkey: str): + 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. diff --git a/bt_ddos_shield/state_manager.py b/bt_ddos_shield/state_manager.py index dd5dfc6..c313946 100644 --- a/bt_ddos_shield/state_manager.py +++ b/bt_ddos_shield/state_manager.py @@ -1,58 +1,53 @@ from abc import ABC, abstractmethod +from datetime import datetime from enum import Enum -from bt_ddos_shield.domain import Domain +from bt_ddos_shield.address import Address +from bt_ddos_shield.miner_shield import Hotkey -class State: +class MinerShieldPhase(Enum): """ - Class representing state of MinerShield. + 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 ShieldPhase(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 - phase: ShieldPhase # current phase of the shield - banned_validators: dict[str, int] # banned validators with ban time (HotKey -> ban timestamp) - active_domains: dict[str, Domain] # active domains (validator HotKey -> Domain created for him) +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 = self.ShieldPhase.DISABLED + self.phase = MinerShieldPhase.DISABLED self.banned_validators = {} - self.active_domains = {} + self.active_addresses = {} + -class StateManager(ABC): +class AbstractMinerShieldStateManager(ABC): """ Abstract base class for manager handling state of MinerShield. """ - state: State # current state of MinerShield - - def __init__(self): - pass + current_miner_shield_state: MinerShieldState @abstractmethod def load_state(self): - """ - Load current state. - """ pass @abstractmethod def save_state(self): - """ - Save current state. - """ pass @abstractmethod - def ban_validator(self, validator_hotkey: str): + def ban_validator(self, validator_hotkey: Hotkey): """ Add validator to the list of banned validators. @@ -62,9 +57,9 @@ def ban_validator(self, validator_hotkey: str): pass @abstractmethod - def remove_validator(self, validator_hotkey: str): + def remove_validator(self, validator_hotkey: Hotkey): """ - Remove validator from the lists of banned validators or active domains. + Remove validator from the lists of banned validators or active addresses. Args: validator_hotkey: The hotkey of the validator. @@ -72,12 +67,12 @@ def remove_validator(self, validator_hotkey: str): pass @abstractmethod - def add_domain(self, validator_hotkey: str, domain: Domain): + def add_address(self, validator_hotkey: Hotkey, address: Address): """ - Add new domain to state. + Add new address to state. Args: validator_hotkey: The hotkey of the validator. - domain: Domain to add. + address: Address to add. """ pass diff --git a/bt_ddos_shield/validators_manager.py b/bt_ddos_shield/validators_manager.py index 1ff497a..d1c2268 100644 --- a/bt_ddos_shield/validators_manager.py +++ b/bt_ddos_shield/validators_manager.py @@ -1,15 +1,15 @@ from abc import ABC, abstractmethod -class ValidatorsManager(ABC): +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. """ - def __init__(self): - pass - @abstractmethod - def get_validators(self) -> dict[str, str]: + def get_validators(self) -> dict[Hotkey, str]: """ Get cached dictionary of validators - maps HotKey of validator to public key. """