Cuttlefish Environment Setup Script

init.py

import os
import sys
import time
import logging
import shutil
from pathlib import Path
from typing import List, Dict, Optional
from dataclasses import dataclass
from git import Repo
import docker
import podman
import subprocess
import requests
from bs4 import BeautifulSoup

# Constants
CUTTLEFISH_REPO = "https://github.com/google/android-cuttlefish.git"
DEFAULT_LOGGING_FORMAT = '%(asctime)s - %(levelname)s - %(message)s'

@dataclass
class AndroidConfig:
    version: int
    disk_size: str = "32G"
    memory_size: str = "6G"
    num_cpus: int = 4
    device_type: str = "cuttlefish"
    enable_webrtc: bool = True
    base_ports: Dict[str, int] = None

    def __post_init__(self):
        self.base_ports = {
            "http": 1080,
            "https": 1443,
            "alt_http": 2080,
            "alt_https": 2443,
            "webrtc_start": 15550,
            "webrtc_end": 15560,
            "adb_start": 6520,
            "adb_end": 6620
        }

class CuttlefishSetup:
    def __init__(self):
        self._setup_logging()
        self.logger = logging.getLogger(__name__)
        self.supported_versions = [12, 13]

    @staticmethod
    def _setup_logging():
        logging.basicConfig(level=logging.INFO, format=DEFAULT_LOGGING_FORMAT)

    def handle_error(func):
        def wrapper(self, *args, **kwargs):
            try:
                return func(self, *args, **kwargs)
            except Exception as e:
                self.logger.error(f"Error in {func.__name__}: {e}")
                sys.exit(1)
        return wrapper

    @handle_error
    def get_container_runtime(self) -> Optional[object]:
        """Get available container runtime client (Docker or Podman)."""
        if shutil.which("docker"):
            return docker.from_env()
        elif shutil.which("podman"):
            return podman.Client()
        raise RuntimeError("Neither Docker nor Podman is installed")

    @handle_error
    def build_cuttlefish_image(self) -> bool:
        """
        Build the Cuttlefish container image.
        Returns True if image already existed, False if it needed to be built.
        """
        self.logger.info("Checking Cuttlefish image...")
        
        if self._image_exists("cuttlefish-orchestration"):
            self.logger.info("Cuttlefish image already exists")
            return True
            
        self.logger.info("Building Cuttlefish image...")
        repo_path = Path("android-cuttlefish")
        
        if not repo_path.is_dir():
            self.logger.info("Cloning the Cuttlefish repository...")
            Repo.clone_from(CUTTLEFISH_REPO, str(repo_path))
        
        subprocess.run(["bash", str(repo_path / "docker/image-builder.sh"), "-d"], check=True)
        return False

    def _image_exists(self, image_name: str) -> bool:
        """Check if a Docker image exists."""
        result = subprocess.run(["docker", "images"], capture_output=True, text=True)
        return image_name in result.stdout

    def _generate_service_config(self, android_config: AndroidConfig) -> dict:
        """Generate Docker Compose service configuration for a specific Android version."""
        port_offset = 0 if android_config.version == 12 else 1
        
        return {
            f"cuttlefish-android{android_config.version}": {
                "image": "cuttlefish-orchestration:latest",
                "privileged": True,
                "user": "root",
                "volumes": [
                    f"./cuttlefish-android{android_config.version}.json:/etc/cuttlefish/cuttlefish.json",
                    "/dev/kvm:/dev/kvm",
                    "/dev/vhost-net:/dev/vhost-net",
                    "/dev/vhost-vsock:/dev/vhost-vsock",
                    "/lib/modules:/lib/modules:ro",
                    "/usr/src:/usr/src:ro"
                ],
                "ports": self._generate_port_mappings(android_config, port_offset),
                "environment": self._generate_environment_vars(android_config),
                "command": "bash -c 'apt update && apt install -y kmod && exec /init'",
                "cap_add": ["SYS_ADMIN", "MKNOD"]
            }
        }

    def _generate_port_mappings(self, config: AndroidConfig, offset: int) -> List[str]:
        """Generate port mappings for the container."""
        base_ports = config.base_ports
        return [
            f"{base_ports['http'] + offset}:{base_ports['http']}",
            f"{base_ports['https'] + offset}:{base_ports['https']}",
            f"{base_ports['alt_http'] + offset}:{base_ports['alt_http']}",
            f"{base_ports['alt_https'] + offset}:{base_ports['alt_https']}",
            f"{base_ports['webrtc_start'] + (offset * 21)}-{base_ports['webrtc_end'] + (offset * 21)}:{base_ports['webrtc_start']}-{base_ports['webrtc_end']}",
            f"{base_ports['adb_start'] + (offset * 101)}-{base_ports['adb_end'] + (offset * 101)}:{base_ports['adb_start']}-{base_ports['adb_end']}"
        ]

    @staticmethod
    def _generate_environment_vars(config: AndroidConfig) -> Dict[str, str]:
        """Generate environment variables for the container."""
        return {
            "ANDROID_VERSION": str(config.version),
            "DEVICE_TYPE": config.device_type,
            "DISK_SIZE": config.disk_size,
            "MEMORY_SIZE": config.memory_size,
            "NUM_CPUS": str(config.num_cpus),
            "BOOT_ARGS": "androidboot.hardware=android,androidboot.serialno=emulator",
            "ENABLE_WEBRTC": str(config.enable_webrtc).lower()
        }

    @handle_error
    def create_configuration_files(self):
        """Create Docker Compose and configuration files."""
        services = {}
        for version in self.supported_versions:
            config = AndroidConfig(version=version)
            services.update(self._generate_service_config(config))
            
            self._write_android_config(config)
        
        self._write_docker_compose(services)

    def _write_docker_compose(self, services: Dict):
        """Write Docker Compose configuration file."""
        compose_config = {
            "version": "3.8",
            "services": services
        }
        with open("docker-compose.yml", "w") as f:
            import yaml
            yaml.dump(compose_config, f, default_flow_style=False)

    def _write_android_config(self, config: AndroidConfig):
        """Write Android-specific configuration file."""
        config_content = {
            "android_version": str(config.version),
            "device_type": config.device_type,
            "disk_size": config.disk_size,
            "memory_size": config.memory_size,
            "num_cpus": config.num_cpus,
            "boot_args": [
                "androidboot.hardware=android",
                "androidboot.serialno=emulator"
            ],
            "enable_webrtc": config.enable_webrtc
        }
        
        with open(f"cuttlefish-android{config.version}.json", "w") as f:
            import json
            json.dump(config_content, f, indent=2)

    @handle_error
    def start_containers(self):
        """Start the Docker Compose services."""
        compose_file = Path("docker-compose.yml")
        
        if not compose_file.exists():
            self.logger.error("docker-compose.yml not found. Run create_configuration_files() first.")
            return
            
        self.logger.info("Starting Docker Compose services...")
        try:
            subprocess.run(["docker-compose", "up", "-d"], check=True)
            
            # Wait a few seconds and check container status
            time.sleep(5)
            result = subprocess.run(["docker-compose", "ps"], capture_output=True, text=True)
            
            if "Exit" in result.stdout or "Error" in result.stdout:
                self.logger.warning("Some containers may have failed to start properly. Check logs with 'docker-compose logs'")
            else:
                self.logger.info("Containers started successfully")
                
        except subprocess.CalledProcessError as e:
            self.logger.error(f"Failed to start containers: {e}")
            raise

    @handle_error
    def check_container_health(self):
        """Check the health of running containers."""
        try:
            result = subprocess.run(
                ["docker-compose", "ps"], 
                capture_output=True, 
                text=True,
                check=True
            )
            
            containers = result.stdout.split('\n')[1:]  # Skip header
            for container in containers:
                if container.strip():
                    if "Up" not in container:
                        self.logger.warning(f"Container may not be healthy: {container}")
                        # Get container logs
                        container_name = container.split()[0]
                        logs = subprocess.run(
                            ["docker-compose", "logs", container_name],
                            capture_output=True,
                            text=True
                        )
                        self.logger.info(f"Container logs for {container_name}:\n{logs.stdout}")
                        
        except subprocess.CalledProcessError as e:
            self.logger.error(f"Failed to check container health: {e}")
            raise

    @staticmethod
    def get_latest_android_12_artifact() -> str:
        url = "https://ci.android.com/builds/branches/aosp-android12-gsi/grid"
        response = requests.get(url)
        response.raise_for_status()

        soup = BeautifulSoup(response.text, 'html.parser')
        artifact_links = soup.find_all('a', href=True)

        for link in artifact_links:
            if 'artifact' in link['href']:
                artifact_url = link['href']
                break
        else:
            raise RuntimeError("No artifact found for Android 12")

        return artifact_url

    @staticmethod
    def get_latest_android_13_artifact() -> str:
        url = "https://ci.android.com/builds/branches/aosp-android13-gsi/grid"
        response = requests.get(url)
        response.raise_for_status()

        soup = BeautifulSoup(response.text, 'html.parser')
        artifact_links = soup.find_all('a', href=True)

        for link in artifact_links:
            if 'artifact' in link['href']:
                artifact_url = link['href']
                break
        else:
            raise RuntimeError("No artifact found for Android 13")

        return artifact_url

    @staticmethod
    def download_artifact(url: str, destination: str):
        response = requests.get(url, stream=True)
        response.raise_for_status()

        with open(destination, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)

        logging.info(f"Downloaded artifact to {destination}")

    @handle_error
    def fetch_and_download_latest_artifact(self, android_version: int):
        self.logger.info(f"Fetching the latest Android {android_version} artifact...")
        if android_version == 12:
            artifact_url = self.get_latest_android_12_artifact()
        elif android_version == 13:
            artifact_url = self.get_latest_android_13_artifact()
        else:
            raise ValueError("Unsupported Android version")
        
        self.logger.info(f"Artifact URL: {artifact_url}")
        self.download_artifact(artifact_url, f"android{android_version}_artifact.zip")

def main():
    setup = CuttlefishSetup()
    
    if setup.build_cuttlefish_image():
        setup.create_configuration_files()
        setup.start_containers()
        setup.check_container_health()
    else:
        setup.create_configuration_files()
        setup.logger.info("Image built successfully. You can now start the containers with 'docker-compose up -d'")

    setup.fetch_and_download_latest_artifact(12)  # Change to 13 if you want to download Android 13 artifact

if __name__ == "__main__":
    main()