Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implementing logic #7

Merged
Changes from 1 commit
Commits
Show all changes
32 commits
Select commit Hold shift + click to select a range
0dbf21f
feat: implement enabling shield
Oct 30, 2024
9648ecb
fix: fixes after code review
Nov 5, 2024
6f321c7
feat: shield logic finished
Nov 7, 2024
23fccc4
fix: code review fixes
Nov 8, 2024
290fc3d
feat: add serialization logic
Nov 9, 2024
1c69b99
feat: logic finished
Nov 11, 2024
6b4abe1
feat: encryption manager refactored
Nov 11, 2024
34a9cf7
feat: add memory implementation for some managers
Nov 11, 2024
168de86
feat: finish memory implementations for all managers along with full …
Nov 12, 2024
8f975a7
fix: finish tests for shield with needed fixes (there is still proble…
Nov 12, 2024
7f1873c
feat: back to EncryptionManager using ECIES library
Nov 13, 2024
4d7fba3
fix: added hash to manifest file to verify if content is the same
Nov 13, 2024
9802e2d
feat: implemented Route53AddressManager
Nov 14, 2024
16603b6
feat: implemented SQLAlchemyMinerShieldStateManager
Nov 14, 2024
6b24602
feat: implemented S3ManifestManager
Nov 15, 2024
82e330c
feat: added integration test
Nov 15, 2024
968f045
feat: filling AwsAddressManager, work in progress
Nov 22, 2024
935c2e6
feat: added possibility of specyfing IP in AwsAddressManager
Nov 22, 2024
6ba1675
feat: filling AwsAddressManager - creating ELB is working
Nov 22, 2024
2ac5dd7
feat: filling AwsAddressManager - creating HostedZone during ELB crea…
Nov 24, 2024
097b2de
feat: filling AwsAddressManager - creating alias DNS entry in Route53
Nov 24, 2024
dc62a30
feat: finished AwsAddressManager - added creating firewall
Nov 25, 2024
c9d3b75
feat: back to passing hosted_zone_id to AwsAddressManager
Dec 5, 2024
845ec8a
fix: tiny code review fixes
Dec 6, 2024
44ebe33
Add AWSClientFactory
grzegorz-leszczynski-reef Dec 25, 2024
840941e
Use pytest fixture in TestAddressManager
grzegorz-leszczynski-reef Dec 25, 2024
54de75f
Simplify code after CR
grzegorz-leszczynski-reef Dec 25, 2024
d2bc4a4
Upgrade handling AWS API errors
grzegorz-leszczynski-reef Dec 26, 2024
030cba1
Add project dependencies
grzegorz-leszczynski-reef Dec 26, 2024
5e27248
Apply ruff linter fixes
grzegorz-leszczynski-reef Dec 26, 2024
78ad8f2
Add MinerShieldFactory
grzegorz-leszczynski-reef Dec 27, 2024
4af7d75
Implement Validator
grzegorz-leszczynski-reef Dec 31, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Add AWSClientFactory
  • Loading branch information
grzegorz-leszczynski-reef committed Dec 25, 2024
commit 44ebe33e5ff47933c3c105803178f6e860d7fffd
30 changes: 12 additions & 18 deletions bt_ddos_shield/address_manager.py
Original file line number Diff line number Diff line change
@@ -9,8 +9,7 @@
from types import MappingProxyType
from typing import Any, Union, Optional, Callable

import boto3
import route53
from botocore.client import BaseClient
from botocore.exceptions import ClientError
from route53.connection import Route53Connection
from route53.hosted_zone import HostedZone
@@ -19,7 +18,7 @@
from bt_ddos_shield.address import Address, AddressType
from bt_ddos_shield.event_processor import AbstractMinerShieldEventProcessor
from bt_ddos_shield.state_manager import AbstractMinerShieldStateManager, MinerShieldState
from bt_ddos_shield.utils import Hotkey
from bt_ddos_shield.utils import Hotkey, AWSClientFactory


class AddressManagerException(Exception):
@@ -121,48 +120,43 @@ class AwsAddressManager(AbstractAddressManager):
miner_instance: AwsEC2InstanceData
miner_instance_port: int
""" Port where miner server is working. """
waf_client: Any
waf_client: BaseClient
waf_arn: Optional[str]
elb_client: Any
elb_client: BaseClient
elb_data: Optional[AwsELBData]
ec2_client: Any
ec2_client: BaseClient
hosted_zone_id: str
""" ID of hosted zone in Route53 where addresses are located. """
hosted_zone: HostedZone
route53_client: Route53Connection
route53_boto_client: Any
route53_boto_client: BaseClient
event_processor: AbstractMinerShieldEventProcessor
state_manager: AbstractMinerShieldStateManager

HOSTED_ZONE_ID_KEY: str = 'aws_hosted_zone_id'
INSTANCE_PORT_KEY: str = 'ec2_instance_port'
INSTANCE_ID_KEY: str = 'ec2_instance_id'

def __init__(self, aws_access_key_id: str, aws_secret_access_key: str, miner_region_name: str,
def __init__(self, aws_client_factory: AWSClientFactory,
miner_address: Address, hosted_zone_id: str, event_processor: AbstractMinerShieldEventProcessor,
state_manager: AbstractMinerShieldStateManager):
"""
Initialize AWS address manager. miner_address can be passed as EC2 type (then AWS instance_id should be set in
address field) or as IP/IPV6 (then we try to find EC2 instance with this IP, which should be private IP
address of EC2 instance) - where port means destination port and address_id is ignored.
"""
self.miner_region_name = miner_region_name
self.event_processor = event_processor
self.state_manager = state_manager

self.waf_client = boto3.client('wafv2', aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key, region_name=miner_region_name)
self.waf_client = aws_client_factory.boto3_client('wafv2')
self.waf_arn = None
self.elb_client = boto3.client('elbv2', aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key, region_name=miner_region_name)
self.elb_client = aws_client_factory.boto3_client('elbv2')
self.elb_data = None
self.route53_client = route53.connect(aws_access_key_id, aws_secret_access_key)
self.route53_client = aws_client_factory.route53_client()
self.hosted_zone_id = hosted_zone_id
self.hosted_zone = self.route53_client.get_hosted_zone_by_id(hosted_zone_id)
self.route53_boto_client = boto3.client('route53', aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key)
self.ec2_client = boto3.client('ec2', aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key, region_name=miner_region_name)
self.route53_boto_client = aws_client_factory.boto3_client('route53')
self.ec2_client = aws_client_factory.boto3_client('ec2')
self._initialize_miner_instance(miner_address)

def _initialize_miner_instance(self, miner_address: Address):
39 changes: 15 additions & 24 deletions bt_ddos_shield/manifest_manager.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
import base64
import hashlib
import json
import boto3
from botocore.client import BaseClient
from botocore.exceptions import ClientError
from abc import ABC, abstractmethod
@@ -11,7 +10,7 @@

from bt_ddos_shield.address import Address, AbstractAddressSerializer, AddressType, AddressDeserializationException
from bt_ddos_shield.encryption_manager import AbstractEncryptionManager
from bt_ddos_shield.utils import Hotkey, PublicKey
from bt_ddos_shield.utils import Hotkey, PublicKey, AWSClientFactory


class ManifestManagerException(Exception):
@@ -189,43 +188,35 @@ class S3ManifestManager(AbstractManifestManager):
MANIFEST_FILE_NAME: str = "miner_manifest.json"

bucket_name: Optional[str]
region_name: Optional[str]
aws_client_factory: AWSClientFactory
s3_client: Optional[BaseClient]

def __init__(self, address_serializer: AbstractAddressSerializer, manifest_serializer: AbstractManifestSerializer,
encryption_manager: AbstractEncryptionManager,
aws_access_key_id: Optional[str] = None, aws_secret_access_key: Optional[str] = None,
region_name: Optional[str] = None, bucket_name: Optional[str] = None):
aws_client_factory: AWSClientFactory, bucket_name: Optional[str] = None):
"""
Creates S3ManifestManager instance. Credentials, region_name or bucket_name can be None when used on Validator
side and in such case S3 client is created after manifest file address is deserialized - use
create_client_from_address method for it.

Args:
aws_access_key_id: AWS access key ID.
aws_secret_access_key: AWS secret access key.
region_name: AWS region name.
bucket_name: Name of the bucket where manifest file is stored.
Creates S3ManifestManager instance. bucket_name and aws_region_name inside aws_client_factory can be None when
used on Validator side and in such case these fields are initialized when manifest file address is
deserialized - use create_client_from_address method for it.
"""
super().__init__(address_serializer, manifest_serializer, encryption_manager)
self.region_name = region_name
self.aws_client_factory = aws_client_factory
self.bucket_name = bucket_name
if region_name is not None:
self.s3_client = boto3.client("s3", aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key, region_name=region_name)
if self.aws_client_factory.aws_region_name is not None:
self.s3_client = self.aws_client_factory.boto3_client("s3")

def create_client_from_address(self, address: Address, aws_access_key_id: str, aws_secret_access_key: str):
def create_client_from_address(self, address: Address):
region_name, bucket_name, _ = self._deserialize_manifest_address(address)
self.s3_client = boto3.client("s3", aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key, region_name=region_name)
self.region_name = region_name
self.aws_client_factory.set_aws_region_name(region_name)
self.bucket_name = bucket_name
self.s3_client = self.aws_client_factory.boto3_client("s3")

def _put_manifest_file(self, data: bytes) -> Address:
file_key: str = self.MANIFEST_FILE_NAME
self.s3_client.put_object(Bucket=self.bucket_name, Key=file_key, Body=data)
return Address(address_id=file_key, address_type=AddressType.S3,
address=self._serialize_manifest_address(self.region_name, self.bucket_name, file_key), port=0)
serialized_address: str = self._serialize_manifest_address(self.aws_client_factory.aws_region_name,
self.bucket_name, file_key)
return Address(address_id=file_key, address_type=AddressType.S3, address=serialized_address, port=0)

def _get_manifest_file(self, address: Address) -> bytes:
try:
35 changes: 34 additions & 1 deletion bt_ddos_shield/utils.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,38 @@
from typing import TypeAlias
from typing import TypeAlias, Optional

import boto3
import route53
from botocore.client import BaseClient
from route53.connection import Route53Connection


Hotkey: TypeAlias = str
PublicKey: TypeAlias = str
PrivateKey: TypeAlias = str


class AWSClientFactory:
aws_access_key_id: str
aws_secret_access_key: str
aws_region_name: Optional[str]

def __init__(self, aws_access_key_id: str, aws_secret_access_key: str, aws_region_name: Optional[str]=None):
"""
Args:
aws_access_key_id: AWS access key ID.
aws_secret_access_key: AWS secret access key.
aws_region_name: AWS region name. If not known, it can be set later using set_aws_region_name method.
"""
self.aws_access_key_id = aws_access_key_id
self.aws_secret_access_key = aws_secret_access_key
self.aws_region_name = aws_region_name

def set_aws_region_name(self, aws_region_name: str):
self.aws_region_name = aws_region_name

def boto3_client(self, service_name: str) -> BaseClient:
return boto3.client(service_name, aws_access_key_id=self.aws_access_key_id,
aws_secret_access_key=self.aws_secret_access_key, region_name=self.aws_region_name)

def route53_client(self) -> Route53Connection:
return route53.connect(self.aws_access_key_id, self.aws_secret_access_key)
10 changes: 6 additions & 4 deletions tests/test_address_manager.py
Original file line number Diff line number Diff line change
@@ -4,7 +4,7 @@
from bt_ddos_shield.address_manager import AbstractAddressManager, AwsAddressManager, AwsObjectTypes
from bt_ddos_shield.event_processor import PrintingMinerShieldEventProcessor
from bt_ddos_shield.state_manager import MinerShieldState
from bt_ddos_shield.utils import Hotkey
from bt_ddos_shield.utils import Hotkey, AWSClientFactory
from tests.test_credentials import aws_access_key_id, aws_secret_access_key, miner_instance_id, \
miner_instance_ip, miner_instance_port, miner_region_name, aws_route53_hosted_zone_id, \
aws_route53_other_hosted_zone_id
@@ -65,9 +65,11 @@ def create_aws_address_manager(self, port: int = miner_instance_port,
miner_address: Address = get_miner_address_from_credentials(AddressType.IP, port)
if create_state_manager:
self.state_manager: MemoryMinerShieldStateManager = MemoryMinerShieldStateManager()
return AwsAddressManager(aws_access_key_id, aws_secret_access_key, miner_region_name=miner_region_name,
miner_address=miner_address, hosted_zone_id=hosted_zone_id,
event_processor=PrintingMinerShieldEventProcessor(), state_manager=self.state_manager)
aws_client_factory: AWSClientFactory = AWSClientFactory(aws_access_key_id, aws_secret_access_key,
miner_region_name)
return AwsAddressManager(aws_client_factory=aws_client_factory, miner_address=miner_address,
hosted_zone_id=hosted_zone_id, event_processor=PrintingMinerShieldEventProcessor(),
state_manager=self.state_manager)

def test_create_elb(self):
""" Test creating ELB by AwsAddressManager class. """
15 changes: 8 additions & 7 deletions tests/test_manifest_manager.py
Original file line number Diff line number Diff line change
@@ -6,7 +6,7 @@
from bt_ddos_shield.encryption_manager import ECIESEncryptionManager
from bt_ddos_shield.manifest_manager import JsonManifestSerializer, AbstractManifestManager, \
ManifestNotFoundException, Manifest, S3ManifestManager
from bt_ddos_shield.utils import Hotkey
from bt_ddos_shield.utils import Hotkey, AWSClientFactory
from tests.test_credentials import aws_access_key_id, aws_secret_access_key, aws_s3_region_name, aws_s3_bucket_name


@@ -50,9 +50,9 @@ def test_json_serializer(self):

def test_s3_put_get(self):
""" Test S3ManifestManager class. Put manifest file, get it and check if it was stored correctly. """
manifest_manager = S3ManifestManager(aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
region_name=aws_s3_region_name, bucket_name=aws_s3_bucket_name,
aws_client_factory: AWSClientFactory = AWSClientFactory(aws_access_key_id, aws_secret_access_key,
aws_s3_region_name)
manifest_manager = S3ManifestManager(aws_client_factory=aws_client_factory, bucket_name=aws_s3_bucket_name,
address_serializer=DefaultAddressSerializer(),
manifest_serializer=JsonManifestSerializer(),
encryption_manager=ECIESEncryptionManager())
@@ -67,11 +67,12 @@ def test_s3_put_get(self):
retrieved_data: bytes = manifest_manager._get_manifest_file(address)
assert retrieved_data == other_data

validator_aws_client_factory: AWSClientFactory = AWSClientFactory(aws_access_key_id, aws_secret_access_key)
validator_manifest_manager = S3ManifestManager(address_serializer=DefaultAddressSerializer(),
manifest_serializer=JsonManifestSerializer(),
encryption_manager=ECIESEncryptionManager())
validator_manifest_manager.create_client_from_address(address, aws_access_key_id,
aws_secret_access_key)
encryption_manager=ECIESEncryptionManager(),
aws_client_factory=validator_aws_client_factory)
validator_manifest_manager.create_client_from_address(address)
retrieved_data: bytes = validator_manifest_manager._get_manifest_file(address)
assert retrieved_data == other_data

11 changes: 5 additions & 6 deletions tests/test_miner_shield.py
Original file line number Diff line number Diff line change
@@ -8,7 +8,7 @@
from bt_ddos_shield.manifest_manager import Manifest, S3ManifestManager, JsonManifestSerializer
from bt_ddos_shield.miner_shield import MinerShield, MinerShieldOptions
from bt_ddos_shield.state_manager import MinerShieldState, SQLAlchemyMinerShieldStateManager
from bt_ddos_shield.utils import Hotkey, PublicKey
from bt_ddos_shield.utils import Hotkey, PublicKey, AWSClientFactory
from bt_ddos_shield.validators_manager import MemoryValidatorsManager
from tests.test_address_manager import MemoryAddressManager, get_miner_address_from_credentials
from tests.test_blockchain_manager import MemoryBlockchainManager
@@ -109,15 +109,14 @@ def test_integration(self):
state_manager: SQLAlchemyMinerShieldStateManager = SQLAlchemyMinerShieldStateManager(sql_alchemy_db_url)
state_manager.clear_tables()
miner_address: Address = get_miner_address_from_credentials(AddressType.IP)
aws_client_factory: AWSClientFactory = AWSClientFactory(aws_access_key_id, aws_secret_access_key,
miner_region_name)
address_manager: AwsAddressManager = \
AwsAddressManager(aws_access_key_id, aws_secret_access_key,
miner_region_name=miner_region_name, miner_address=miner_address,
AwsAddressManager(aws_client_factory=aws_client_factory, miner_address=miner_address,
hosted_zone_id=aws_route53_hosted_zone_id,
event_processor=PrintingMinerShieldEventProcessor(), state_manager=state_manager)
manifest_manager: S3ManifestManager = \
S3ManifestManager(aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
region_name=aws_s3_region_name, bucket_name=aws_s3_bucket_name,
S3ManifestManager(aws_client_factory=aws_client_factory, bucket_name=aws_s3_bucket_name,
address_serializer=DefaultAddressSerializer(),
manifest_serializer=JsonManifestSerializer(),
encryption_manager=ECIESEncryptionManager())