AWS Snowball: Data Migration, Edge Computing, and Physical Data Transport at Scale

package io.thecodeforge.aws;

# Snowball capacity and transfer time calculator.
# Use this to determine whether Snowball or direct transfer is faster.

from dataclasses import dataclass
from enum import Enum


class SnowballDevice(Enum):
    """Snowball device tiers with usable capacity in TB."""
    SNOWBALL_STANDARD = 80
    SNOWBALL_EDGE_STORAGE = 210
    SNOWBALL_EDGE_COMPUTE = 100
    SNOWBALL_EDGE_COMPUTE_GPU = 100
    SNOWBALL_EDGE_EBS = 28
    SNOWMOBILE = 100_000


@dataclass
class TransferEstimate:
    device: SnowballDevice
    data_size_tb: float
    network_speed_gbps: float
    snowball_shipping_days: int = 14  # average round-trip shipping time
    aws_loading_days: int = 2  # time for AWS to load data into S3

    def direct_transfer_days(self) -> float:
        """Calculate days to transfer over network."""
        # Convert TB to bits, Gbps to bits/day
        data_bits = self.data_size_tb * 1e12 * 8
        bits_per_day = self.network_speed_gbps * 1e9 * 86400
        return data_bits / bits_per_day

    def snowball_transfer_days(self) -> float:
        """Calculate total Snowball timeline."""
        # Copy to device: assume 100MB/s average (realistic for mixed files)
        copy_speed_mbps = 100
        data_mb = self.data_size_tb * 1e6
        copy_days = data_mb / (copy_speed_mbps * 86400)
        return copy_days + self.snowball_shipping_days + self.aws_loading_days

    def recommendation(self) -> str:
        """Return recommendation: Snowball or direct transfer."""
        direct = self.direct_transfer_days()
        snowball = self.snowball_transfer_days()

        if self.data_size_tb > self.device.value:
            return (
                f"DATA EXCEEDS DEVICE CAPACITY. "
                f"{self.data_size_tb}TB > {self.device.value}TB. "
                f"Use multiple devices or Snowmobile."
            )

        if direct < snowball:
            return (
                f"DIRECT TRANSFER is faster. "
                f"Direct: {direct:.1f} days vs Snowball: {snowball:.1f} days. "
                f"At {self.network_speed_gbps}Gbps, network transfer wins."
            )
        else:
            savings = snowball - direct
            return (
                f"SNOWBALL is faster. "
                f"Snowball: {snowball:.1f} days vs Direct: {direct:.1f} days. "
                f"Saves {savings:.1f} days."
            )


def main():
    scenarios = [
        TransferEstimate(SnowballDevice.SNOWBALL_STANDARD, 50, 1),
        TransferEstimate(SnowballDevice.SNOWBALL_STANDARD, 50, 10),
        TransferEstimate(SnowballDevice.SNOWBALL_EDGE_STORAGE, 150, 1),
        TransferEstimate(SnowballDevice.SNOWBALL_EDGE_STORAGE, 150, 10),
        TransferEstimate(SnowballDevice.SNOWBALL_EDGE_STORAGE, 200, 1),
    ]

    for scenario in scenarios:
        print(f"\n--- {scenario.data_size_tb}TB @ {scenario.network_speed_gbps}Gbps ---")
        print(f"  Device: {scenario.device.name} ({scenario.device.value}TB)")
        print(f"  Direct transfer: {scenario.direct_transfer_days():.1f} days")
        print(f"  Snowball total:  {scenario.snowball_transfer_days():.1f} days")
        print(f"  {scenario.recommendation()}")


if __name__ == '__main__':
    main()

import boto3
import json
import hashlib
import os
from dataclasses import dataclass, field
from typing import Optional, List
from datetime import datetime


@dataclass
class SnowballJobConfig:
    """Configuration for a Snowball migration job."""
    job_id: str
    device_type: str  # STANDARD | EDGE_STORAGE | EDGE_COMPUTE
    s3_bucket: str
    s3_prefix: str
    kms_key_arn: str
    iam_role_arn: str
    source_directory: str
    shipping_address: dict
    notification_sns_topic: str


@dataclass
class ManifestEntry:
    """Single entry in the transfer manifest."""
    file_path: str
    s3_key: str
    size_bytes: int
    sha256_hash: str
    last_modified: str


class SnowballJobManager:
    """Manages Snowball job lifecycle and pre-ship validation."""

    def __init__(self, config: SnowballJobConfig):
        self.config = config
        self.snowball_client = boto3.client('snowball')
        self.kms_client = boto3.client('kms')
        self.iam_client = boto3.client('iam')
        self.s3_client = boto3.client('s3')

    def generate_manifest(self, source_dir: str) -> List[ManifestEntry]:
        """Generate a manifest of all files to be transferred."""
        manifest = []
        total_size = 0
        skipped_files = []

        for dirpath, dirnames, filenames in os.walk(source_dir):
            for filename in filenames:
                filepath = os.path.join(dirpath, filename)

                # Check path length — Snowball S3 adapter rejects > 1024 bytes
                if len(filepath.encode('utf-8')) > 1024:
                    skipped_files.append(filepath)
                    continue

                try:
                    size = os.path.getsize(filepath)
                    total_size += size

                    sha256 = hashlib.sha256()
                    with open(filepath, 'rb') as f:
                        for chunk in iter(lambda: f.read(8192), b''):
                            sha256.update(chunk)

                    s3_key = os.path.relpath(filepath, source_dir)

                    manifest.append(ManifestEntry(
                        file_path=filepath,
                        s3_key=f"{self.config.s3_prefix}/{s3_key}",
                        size_bytes=size,
                        sha256_hash=sha256.hexdigest(),
                        last_modified=datetime.fromtimestamp(
                            os.path.getmtime(filepath)
                        ).isoformat(),
                    ))
                except OSError as e:
                    skipped_files.append(f"{filepath}: {e}")

        if skipped_files:
            print(f"WARNING: {len(skipped_files)} files skipped (path too long or inaccessible)")
            for f in skipped_files[:10]:
                print(f"  - {f}")

        return manifest

    def validate_pre_ship(self, manifest: List[ManifestEntry]) -> dict:
        """Run all pre-ship validations. Returns validation result."""
        errors = []
        warnings = []

        # 1. Validate KMS key exists and is accessible
        try:
            key_response = self.kms_client.describe_key(
                KeyId=self.config.kms_key_arn
            )
            key_state = key_response['KeyMetadata']['KeyState']
            if key_state != 'Enabled':
                errors.append(f"KMS key is not enabled: {key_state}")
        except Exception as e:
            errors.append(f"KMS key validation failed: {e}")

        # 2. Validate IAM role has required permissions
        required_actions = ['s3:PutObject', 's3:CreateBucket', 'kms:GenerateDataKey']
        for action in required_actions:
            try:
                response = self.iam_client.simulate_principal_policy(
                    PolicySourceArn=self.config.iam_role_arn,
                    ActionNames=[action],
                    ResourceArns=[f"arn:aws:s3:::{self.config.s3_bucket}/*"],
                )
                decision = response['EvaluationResults'][0]['EvalDecision']
                if decision != 'allowed':
                    errors.append(f"IAM role missing {action} permission")
            except Exception as e:
                errors.append(f"IAM permission check failed for {action}: {e}")

        # 3. Validate manifest has entries
        if not manifest:
            errors.append("Manifest is empty — no files to transfer")

        # 4. Validate total size fits on device
        total_gb = sum(e.size_bytes for e in manifest) / (1024 ** 3)
        device_limits = {
            'STANDARD': 80,
            'EDGE_STORAGE': 210,
            'EDGE_COMPUTE': 100,
        }
        limit = device_limits.get(self.config.device_type, 80)
        if total_gb > limit:
            errors.append(
                f"Data size {total_gb:.1f}TB exceeds device capacity {limit}TB"
            )
        elif total_gb > limit * 0.9:
            warnings.append(
                f"Data size {total_gb:.1f}TB is >90% of device capacity {limit}TB. "
                f"Consider leaving headroom."
            )

        # 5. Validate no files have paths > 1024 bytes
        long_paths = [e for e in manifest if len(e.file_path.encode('utf-8')) > 1024]
        if long_paths:
            errors.append(
                f"{len(long_paths)} files have paths > 1024 bytes. "
                f"S3 adapter will reject them."
            )

        return {
            'valid': len(errors) == 0,
            'errors': errors,
            'warnings': warnings,
            'manifest_entries': len(manifest),
            'total_size_gb': round(total_gb, 2),
            'device_type': self.config.device_type,
            'device_capacity_tb': limit,
        }

    def save_manifest(self, manifest: List[ManifestEntry], output_path: str):
        """Save manifest to JSON for post-transfer verification."""
        manifest_data = {
            'job_id': self.config.job_id,
            'generated_at': datetime.utcnow().isoformat(),
            'total_entries': len(manifest),
            'total_bytes': sum(e.size_bytes for e in manifest),
            'entries': [
                {
                    'file_path': e.file_path,
                    's3_key': e.s3_key,
                    'size_bytes': e.size_bytes,
                    'sha256': e.sha256_hash,
                    'last_modified': e.last_modified,
                }
                for e in manifest
            ],
        }

        with open(output_path, 'w') as f:
            json.dump(manifest_data, f, indent=2)

        print(f"Manifest saved: {len(manifest)} entries, "
              f"{manifest_data['total_bytes'] / (1024**3):.2f}GB")

    def verify_post_load(self, manifest_path: str) -> dict:
        """Verify S3 contents match manifest after AWS loading."""
        with open(manifest_path, 'r') as f:
            manifest_data = json.load(f)

        mismatches = []
        missing = []

        for entry in manifest_data['entries']:
            try:
                response = self.s3_client.head_object(
                    Bucket=self.config.s3_bucket,
                    Key=entry['s3_key'],
                )
                s3_size = response['ContentLength']
                if s3_size != entry['size_bytes']:
                    mismatches.append({
                        'key': entry['s3_key'],
                        'expected_size': entry['size_bytes'],
                        'actual_size': s3_size,
                    })
            except self.s3_client.exceptions.ClientError:
                missing.append(entry['s3_key'])

        return {
            'verified': len(mismatches) == 0 and len(missing) == 0,
            'total_checked': len(manifest_data['entries']),
            'missing': len(missing),
            'size_mismatches': len(mismatches),
            'missing_files': missing[:20],
            'mismatched_files': mismatches[:20],
        }

package io.thecodeforge.aws;

import boto3
import json
from dataclasses import dataclass
from typing import Optional


@dataclass
class EdgeComputeConfig:
    """Configuration for Snowball Edge compute workloads."""
    device_id: str
    instance_type: str  # sbg1.medium, sbg1.large, sbg1.xlarge
    ami_id: str
    vcpus: int
    memory_gb: int
    ebs_volumes: list  # list of {size_gb, mount_point}


class SnowballEdgeManager:
    """Manages Snowball Edge compute instances and edge workloads."""

    def __init__(self, device_ip: str, unlock_code: str, manifest_path: str):
        self.device_ip = device_ip
        self.unlock_code = unlock_code
        self.manifest_path = manifest_path
        # Snowball Edge exposes a local EC2-compatible endpoint
        self.ec2_client = boto3.client(
            'ec2',
            endpoint_url=f"https://{device_ip}:8443",
            verify=False,  # Snowball Edge uses self-signed certs
        )
        self.s3_client = boto3.client(
            's3',
            endpoint_url=f"https://{device_ip}:8443",
            verify=False,
        )

    def list_local_instances(self) -> list:
        """List EC2 instances running on the Snowball Edge device."""
        response = self.ec2_client.describe_instances()
        instances = []
        for reservation in response['Reservations']:
            for instance in reservation['Instances']:
                instances.append({
                    'instance_id': instance['InstanceId'],
                    'state': instance['State']['Name'],
                    'type': instance['InstanceType'],
                    'private_ip': instance.get('PrivateIpAddress', 'N/A'),
                })
        return instances

    def launch_edge_instance(self, config: EdgeComputeConfig) -> dict:
        """Launch an EC2 instance on the Snowball Edge device."""
        response = self.ec2_client.run_instances(
            ImageId=config.ami_id,
            InstanceType=config.instance_type,
            MinCount=1,
            MaxCount=1,
        )
        instance = response['Instances'][0]
        return {
            'instance_id': instance['InstanceId'],
            'state': instance['State']['Name'],
        }

    def list_local_buckets(self) -> list:
        """List S3 buckets on the Snowball Edge device."""
        response = self.s3_client.list_buckets()
        return [bucket['Name'] for bucket in response['Buckets']]

    def get_device_storage_status(self) -> dict:
        """Get storage utilization on the Snowball Edge device."""
        # Note: actual API varies by Snowball Edge SDK version
        # This is a simplified example
        return {
            'note': 'Use snowballEdge describe-device --device-id <id> CLI command',
            'command': f'snowballEdge describe-device --device-id {self.device_ip}',
        }

    def validate_edge_workload(self, config: EdgeComputeConfig) -> list:
        """Validate that the edge workload fits on the device."""
        errors = []

        # Check instance type availability
        valid_types = ['sbg1.medium', 'sbg1.large', 'sbg1.xlarge']
        if config.instance_type not in valid_types:
            errors.append(f"Invalid instance type: {config.instance_type}. Valid: {valid_types}")

        # Check EBS volume sizes
        total_ebs_gb = sum(v['size_gb'] for v in config.ebs_volumes)
        if total_ebs_gb > 28000:  # 28TB max for EBS volumes
            errors.append(f"Total EBS size {total_ebs_gb}GB exceeds 28TB limit")

        return errors

package io.thecodeforge.aws;

import boto3
import json
from dataclasses import dataclass


@dataclass
class KMSValidationResult:
    key_arn: str
    key_exists: bool
    key_enabled: bool
    key_account: str
    job_account: str
    account_match: bool
    permissions_valid: bool
    bucket_policy_compatible: bool
    errors: list


class SnowballKMSValidator:
    """Validates KMS key configuration for Snowball jobs."""

    def __init__(self, job_account_region: str = 'us-east-1'):
        self.kms_client = boto3.client('kms', region_name=job_account_region)
        self.iam_client = boto3.client('iam', region_name=job_account_region)
        self.s3_client = boto3.client('s3', region_name=job_account_region)
        self.sts_client = boto3.client('sts', region_name=job_account_region)

    def validate_kms_for_snowball(
        self,
        kms_key_arn: str,
        iam_role_arn: str,
        s3_bucket: str,
    ) -> KMSValidationResult:
        """Comprehensive KMS validation for a Snowball job."""
        errors = []

        # Get current account
        job_account = self.sts_client.get_caller_identity()['Account']

        # 1. Verify KMS key exists and is enabled
        key_exists = False
        key_enabled = False
        key_account = ''
        try:
            key_response = self.kms_client.describe_key(KeyId=kms_key_arn)
            key_metadata = key_response['KeyMetadata']
            key_exists = True
            key_enabled = key_metadata['KeyState'] == 'Enabled'
            key_account = key_metadata['Arn'].split(':')[4]

            if not key_enabled:
                errors.append(f"KMS key is not enabled. State: {key_metadata['KeyState']}")

            if key_metadata.get('KeyManager') == 'AWS':
                errors.append(
                    "AWS-managed keys (aws/s3) are not recommended for Snowball. "
                    "Use a customer-managed key for cross-account compatibility."
                )
        except Exception as e:
            errors.append(f"KMS key not found or not accessible: {e}")

        # 2. Check cross-account permissions
        account_match = key_account == job_account
        permissions_valid = False

        if not account_match:
            errors.append(
                f"KMS key is in account {key_account} but Snowball job is in {job_account}. "
                f"Cross-account KMS permissions required."
            )

            # Check if the IAM role has permissions on the cross-account key
            try:
                for action in ['kms:Decrypt', 'kms:GenerateDataKey', 'kms:DescribeKey']:
                    response = self.iam_client.simulate_principal_policy(
                        PolicySourceArn=iam_role_arn,
                        ActionNames=[action],
                        ResourceArns=[kms_key_arn],
                    )
                    decision = response['EvaluationResults'][0]['EvalDecision']
                    if decision != 'allowed':
                        errors.append(
                            f"IAM role missing {action} on cross-account KMS key. "
                            f"Add permission to IAM policy AND KMS key policy."
                        )
                    else:
                        permissions_valid = True
            except Exception as e:
                errors.append(f"Permission simulation failed: {e}")
        else:
            permissions_valid = True

        # 3. Check S3 bucket policy compatibility
        bucket_policy_compatible = True
        try:
            bucket_policy = self.s3_client.get_bucket_policy(Bucket=s3_bucket)
            policy = json.loads(bucket_policy['Policy'])

            # Check if bucket policy requires a specific KMS key
            for statement in policy.get('Statement', []):
                if statement.get('Effect') == 'Deny':
                    condition = statement.get('Condition', {})
                    kms_condition = condition.get('StringNotEquals', {}).get(
                        's3:x-amz-server-side-encryption-aws-kms-key-id', ''
                    )
                    if kms_condition and kms_condition != kms_key_arn:
                        errors.append(
                            f"Bucket policy denies writes not encrypted with {kms_condition}. "
                            f"Snowball job uses {kms_key_arn}. Mismatch."
                        )
                        bucket_policy_compatible = False
        except self.s3_client.exceptions.NoSuchBucketPolicy:
            pass  # No bucket policy is fine
        except Exception as e:
            errors.append(f"Bucket policy check failed: {e}")

        return KMSValidationResult(
            key_arn=kms_key_arn,
            key_exists=key_exists,
            key_enabled=key_enabled,
            key_account=key_account,
            job_account=job_account,
            account_match=account_match,
            permissions_valid=permissions_valid,
            bucket_policy_compatible=bucket_policy_compatible,
            errors=errors,
        )

    def generate_cross_account_kms_policy(
        self,
        job_account_id: str,
        job_role_arn: str,
    ) -> str:
        """Generate a KMS key policy for cross-account Snowball access."""
        policy = {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Sid": "AllowSnowballJobAccount",
                    "Effect": "Allow",
                    "Principal": {
                        "AWS": job_role_arn
                    },
                    "Action": [
                        "kms:Decrypt",
                        "kms:GenerateDataKey",
                        "kms:DescribeKey",
                    ],
                    "Resource": "*",
                },
                {
                    "Sid": "AllowRootAccountFullAccess",
                    "Effect": "Allow",
                    "Principal": {
                        "AWS": f"arn:aws:iam::{job_account_id}:root"
                    },
                    "Action": "kms:*",
                    "Resource": "*",
                },
            ],
        }
        return json.dumps(policy, indent=2)

package io.thecodeforge.aws;

import os
import subprocess
import tarfile
import tempfile
from dataclasses import dataclass
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path


@dataclass
class CopyStrategy:
    """Copy strategy based on file size distribution."""
    small_file_threshold_mb: float = 1.0
    tar_batch_size_mb: float = 1024.0  # 1GB tar archives
    parallel_sessions: int = 4
    exclude_patterns: list = None

    def __post_init__(self):
        if self.exclude_patterns is None:
            self.exclude_patterns = [
                '*.tmp', '*.temp', '*.log', '*.cache',
                '.DS_Store', 'Thumbs.db', '__pycache__',
                '.git', 'node_modules',
            ]


class SnowballCopyOptimizer:
    """Optimizes data copy to Snowball device for maximum throughput."""

    def __init__(self, source_dir: str, snowball_endpoint: str, bucket: str, prefix: str):
        self.source_dir = Path(source_dir)
        self.snowball_endpoint = snowball_endpoint
        self.bucket = bucket
        self.prefix = prefix

    def analyze_source(self) -> dict:
        """Analyze source directory to determine optimal copy strategy."""
        total_files = 0
        total_size = 0
        small_files = 0
        small_files_size = 0
        large_files = 0
        large_files_size = 0
        max_path_length = 0
        long_path_files = []

        for dirpath, dirnames, filenames in os.walk(self.source_dir):
            for filename in filenames:
                filepath = os.path.join(dirpath, filename)
                try:
                    size = os.path.getsize(filepath)
                except OSError:
                    continue

                total_files += 1
                total_size += size

                if len(filepath) > max_path_length:
                    max_path_length = len(filepath)

                if len(filepath) > 1024:
                    long_path_files.append(filepath)

                if size < 1 * 1024 * 1024:  # < 1MB
                    small_files += 1
                    small_files_size += size
                else:
                    large_files += 1
                    large_files_size += size

        return {
            'total_files': total_files,
            'total_size_gb': total_size / (1024**3),
            'small_files': small_files,
            'small_files_pct': (small_files / total_files * 100) if total_files > 0 else 0,
            'small_files_size_gb': small_files_size / (1024**3),
            'large_files': large_files,
            'large_files_size_gb': large_files_size / (1024**3),
            'max_path_length': max_path_length,
            'long_path_files': len(long_path_files),
            'recommendation': self._recommend_strategy(small_files, total_files),
        }

    def _recommend_strategy(self, small_files: int, total_files: int) -> str:
        """Recommend copy strategy based on file size distribution."""
        small_pct = (small_files / total_files * 100) if total_files > 0 else 0

        if small_pct > 50:
            return (
                "HIGH SMALL FILE COUNT. "
                "Tar small files into large archives before copying. "
                "Expected speedup: 10-50x."
            )
        elif small_pct > 20:
            return (
                "MODERATE SMALL FILE COUNT. "
                "Consider tarring files < 1MB. "
                "Run 4+ parallel copy sessions."
            )
        else:
            return (
                "LARGE FILE DOMINANT. "
                "Direct copy with parallel sessions should achieve near-max throughput. "
                "Use 4-8 parallel s3 cp sessions."
            )

    def create_tar_archives(self, strategy: CopyStrategy) -> list:
        """Create tar archives of small files for efficient transfer."""
        archives = []
        current_batch = []
        current_batch_size = 0
        batch_num = 0

        for dirpath, dirnames, filenames in os.walk(self.source_dir):
            for filename in filenames:
                filepath = os.path.join(dirpath, filename)
                try:
                    size = os.path.getsize(filepath)
                except OSError:
                    continue

                if size < strategy.small_file_threshold_mb * 1024 * 1024:
                    current_batch.append(filepath)
                    current_batch_size += size

                    if current_batch_size >= strategy.tar_batch_size_mb * 1024 * 1024:
                        archive_path = self._write_tar(current_batch, batch_num)
                        archives.append(archive_path)
                        current_batch = []
                        current_batch_size = 0
                        batch_num += 1

        # Write remaining batch
        if current_batch:
            archive_path = self._write_tar(current_batch, batch_num)
            archives.append(archive_path)

        return archives

    def _write_tar(self, files: list, batch_num: int) -> str:
        """Write a batch of files to a tar archive."""
        archive_name = f"small_files_batch_{batch_num:04d}.tar.gz"
        archive_path = os.path.join(tempfile.gettempdir(), archive_name)

        with tarfile.open(archive_path, 'w:gz') as tar:
            for filepath in files:
                arcname = os.path.relpath(filepath, self.source_dir)
                tar.add(filepath, arcname=arcname)

        return archive_path

    def parallel_copy(self, source_dirs: list, strategy: CopyStrategy) -> list:
        """Copy multiple directories in parallel to Snowball device."""
        results = []

        with ThreadPoolExecutor(max_workers=strategy.parallel_sessions) as executor:
            futures = {}
            for i, source in enumerate(source_dirs):
                s3_key = f"{self.prefix}/batch_{i:04d}"
                future = executor.submit(
                    self._copy_to_snowball, source, s3_key
                )
                futures[future] = source

            for future in as_completed(futures):
                source = futures[future]
                try:
                    result = future.result()
                    results.append({'source': source, 'status': 'success', 'result': result})
                except Exception as e:
                    results.append({'source': source, 'status': 'error', 'error': str(e)})

        return results

    def _copy_to_snowball(self, source: str, s3_key: str) -> dict:
        """Copy a single directory to the Snowball device via S3 adapter."""
        cmd = [
            'aws', 's3', 'cp',
            source,
            f"s3://{self.bucket}/{s3_key}/",
            '--recursive',
            '--endpoint-url', f"https://{self.snowball_endpoint}:8443",
            '--no-verify-ssl',
        ]

        result = subprocess.run(cmd, capture_output=True, text=True)

        if result.returncode != 0:
            raise RuntimeError(f"Copy failed: {result.stderr}")

        return {'command': ' '.join(cmd), 'stdout': result.stdout}