pi-pm3/app/backend/managers/update_manager.py

"""Update Manager for Dangerous Pi.

Handles checking for updates from GitHub releases, downloading,
and applying updates to the system.
"""
import asyncio
import hashlib
import json
import os
import re
import shutil
import tempfile
from dataclasses import dataclass
from datetime import datetime, timedelta
from enum import Enum
from pathlib import Path
from typing import Optional, Dict, Any, List
import aiohttp
import aiosqlite

from .. import config


class UpdateStatus(str, Enum):
    """Update status enum."""
    IDLE = "idle"
    CHECKING = "checking"
    AVAILABLE = "available"
    DOWNLOADING = "downloading"
    INSTALLING = "installing"
    COMPLETE = "complete"
    FAILED = "failed"


@dataclass
class ReleaseInfo:
    """GitHub release information."""
    version: str
    tag_name: str
    published_at: str
    download_url: str
    changelog: str
    is_prerelease: bool
    asset_name: str
    asset_size: int
    checksum: Optional[str] = None


@dataclass
class UpdateProgress:
    """Update progress information."""
    status: UpdateStatus
    current_version: str
    available_version: Optional[str] = None
    download_progress: float = 0.0
    error_message: Optional[str] = None
    last_check: Optional[str] = None


class UpdateManager:
    """Manages system updates from GitHub releases."""

    def __init__(self):
        """Initialize the update manager."""
        self._current_version = config.VERSION
        self._github_repo = config.GITHUB_REPO
        self._github_api_base = "https://api.github.com"
        self._status = UpdateStatus.IDLE
        self._progress = UpdateProgress(
            status=UpdateStatus.IDLE,
            current_version=self._current_version
        )
        self._latest_release: Optional[ReleaseInfo] = None
        self._update_lock = asyncio.Lock()
        self._download_path: Optional[Path] = None
        self._check_interval = config.UPDATE_CHECK_INTERVAL

    async def start_periodic_checks(self):
        """Start periodic update checks in background."""
        while True:
            try:
                await self.check_for_updates()
                await asyncio.sleep(self._check_interval)
            except asyncio.CancelledError:
                break
            except Exception as e:
                print(f"Error in periodic update check: {e}")
                await asyncio.sleep(self._check_interval)

    async def check_for_updates(self) -> Dict[str, Any]:
        """Check for available updates from GitHub.

        Returns:
            Dict containing update status and info
        """
        async with self._update_lock:
            try:
                self._status = UpdateStatus.CHECKING
                self._progress.status = UpdateStatus.CHECKING

                # Fetch latest release from GitHub
                release = await self._fetch_latest_release()

                if not release:
                    self._status = UpdateStatus.IDLE
                    self._progress.status = UpdateStatus.IDLE
                    self._progress.last_check = datetime.utcnow().isoformat()
                    return {
                        "update_available": False,
                        "current_version": self._current_version,
                        "message": "No releases found"
                    }

                # Compare versions
                if self._is_newer_version(release.version, self._current_version):
                    self._latest_release = release
                    self._status = UpdateStatus.AVAILABLE
                    self._progress.status = UpdateStatus.AVAILABLE
                    self._progress.available_version = release.version

                    return {
                        "update_available": True,
                        "current_version": self._current_version,
                        "latest_version": release.version,
                        "release_date": release.published_at,
                        "changelog": release.changelog,
                        "is_prerelease": release.is_prerelease,
                        "download_size": release.asset_size
                    }
                else:
                    self._status = UpdateStatus.IDLE
                    self._progress.status = UpdateStatus.IDLE
                    self._progress.last_check = datetime.utcnow().isoformat()

                    return {
                        "update_available": False,
                        "current_version": self._current_version,
                        "latest_version": release.version,
                        "message": "System is up to date"
                    }

            except Exception as e:
                self._status = UpdateStatus.FAILED
                self._progress.status = UpdateStatus.FAILED
                self._progress.error_message = str(e)
                raise Exception(f"Failed to check for updates: {e}")

    async def download_update(self) -> bool:
        """Download the latest update.

        Returns:
            True if download successful, False otherwise
        """
        async with self._update_lock:
            if not self._latest_release:
                raise ValueError("No update available to download")

            try:
                self._status = UpdateStatus.DOWNLOADING
                self._progress.status = UpdateStatus.DOWNLOADING
                self._progress.download_progress = 0.0

                # Create temp directory for download
                temp_dir = Path(tempfile.mkdtemp(prefix="dangerous-pi-update-"))
                self._download_path = temp_dir / self._latest_release.asset_name

                # Download the release asset
                async with aiohttp.ClientSession() as session:
                    async with session.get(self._latest_release.download_url) as response:
                        if response.status != 200:
                            raise Exception(f"Download failed with status {response.status}")

                        total_size = int(response.headers.get('content-length', 0))
                        downloaded = 0

                        with open(self._download_path, 'wb') as f:
                            async for chunk in response.content.iter_chunked(8192):
                                f.write(chunk)
                                downloaded += len(chunk)
                                if total_size > 0:
                                    self._progress.download_progress = (downloaded / total_size) * 100

                # Verify checksum if available
                if self._latest_release.checksum:
                    if not await self._verify_checksum(self._download_path, self._latest_release.checksum):
                        raise Exception("Checksum verification failed")

                self._progress.download_progress = 100.0
                return True

            except Exception as e:
                self._status = UpdateStatus.FAILED
                self._progress.status = UpdateStatus.FAILED
                self._progress.error_message = str(e)

                # Cleanup on failure
                if self._download_path and self._download_path.parent.exists():
                    shutil.rmtree(self._download_path.parent, ignore_errors=True)

                raise Exception(f"Failed to download update: {e}")

    async def install_update(self) -> bool:
        """Install the downloaded update.

        Returns:
            True if installation successful, False otherwise
        """
        async with self._update_lock:
            if not self._download_path or not self._download_path.exists():
                raise ValueError("No update downloaded")

            try:
                self._status = UpdateStatus.INSTALLING
                self._progress.status = UpdateStatus.INSTALLING

                # Extract the update archive
                install_dir = Path("/opt/dangerous-pi")
                backup_dir = Path("/opt/dangerous-pi-backup")

                # Create backup of current installation
                if install_dir.exists():
                    if backup_dir.exists():
                        shutil.rmtree(backup_dir)
                    shutil.copytree(install_dir, backup_dir)

                # Extract update (assuming tar.gz format)
                await self._run_command(
                    f"tar -xzf {self._download_path} -C {install_dir.parent}"
                )

                # Run post-install script if exists
                post_install = install_dir / "scripts" / "post-install.sh"
                if post_install.exists():
                    await self._run_command(f"sudo bash {post_install}")

                # Rebuild PM3 client if needed
                await self._rebuild_pm3_client()

                # Update version in config
                await self._update_version_file(self._latest_release.version)

                # Cleanup
                shutil.rmtree(self._download_path.parent, ignore_errors=True)
                self._download_path = None

                self._status = UpdateStatus.COMPLETE
                self._progress.status = UpdateStatus.COMPLETE
                self._current_version = self._latest_release.version
                self._progress.current_version = self._latest_release.version

                return True

            except Exception as e:
                self._status = UpdateStatus.FAILED
                self._progress.status = UpdateStatus.FAILED
                self._progress.error_message = str(e)

                # Restore backup on failure
                if backup_dir.exists():
                    if install_dir.exists():
                        shutil.rmtree(install_dir)
                    shutil.copytree(backup_dir, install_dir)

                raise Exception(f"Failed to install update: {e}")

    async def get_progress(self) -> UpdateProgress:
        """Get current update progress.

        Returns:
            UpdateProgress object
        """
        return self._progress

    async def get_release_notes(self, version: Optional[str] = None) -> str:
        """Get release notes for a specific version.

        Args:
            version: Version to get notes for (latest if None)

        Returns:
            Release notes as markdown string
        """
        try:
            if version:
                release = await self._fetch_release_by_version(version)
            else:
                release = await self._fetch_latest_release()

            return release.changelog if release else "No release notes available"

        except Exception as e:
            raise Exception(f"Failed to get release notes: {e}")

    async def _fetch_latest_release(self) -> Optional[ReleaseInfo]:
        """Fetch latest release from GitHub API.

        Returns:
            ReleaseInfo object or None
        """
        url = f"{self._github_api_base}/repos/{self._github_repo}/releases/latest"

        async with aiohttp.ClientSession() as session:
            async with session.get(url) as response:
                if response.status == 404:
                    return None

                if response.status != 200:
                    raise Exception(f"GitHub API returned status {response.status}")

                data = await response.json()
                return self._parse_release_data(data)

    async def _fetch_release_by_version(self, version: str) -> Optional[ReleaseInfo]:
        """Fetch specific release by version tag.

        Args:
            version: Version tag (e.g., "v1.0.0")

        Returns:
            ReleaseInfo object or None
        """
        tag = version if version.startswith('v') else f'v{version}'
        url = f"{self._github_api_base}/repos/{self._github_repo}/releases/tags/{tag}"

        async with aiohttp.ClientSession() as session:
            async with session.get(url) as response:
                if response.status == 404:
                    return None

                if response.status != 200:
                    raise Exception(f"GitHub API returned status {response.status}")

                data = await response.json()
                return self._parse_release_data(data)

    def _parse_release_data(self, data: Dict[str, Any]) -> ReleaseInfo:
        """Parse GitHub release API response.

        Args:
            data: GitHub API release data

        Returns:
            ReleaseInfo object
        """
        # Find the main release asset (tar.gz)
        assets = data.get('assets', [])
        main_asset = None
        checksum_asset = None

        for asset in assets:
            name = asset['name']
            if name.endswith('.tar.gz'):
                main_asset = asset
            elif name.endswith('.sha256'):
                checksum_asset = asset

        if not main_asset:
            raise ValueError("No suitable release asset found")

        # Get version from tag (remove 'v' prefix)
        version = data['tag_name'].lstrip('v')

        # Get checksum if available
        checksum = None
        if checksum_asset:
            # Would need to download and read checksum file
            # For now, we'll skip this step
            pass

        return ReleaseInfo(
            version=version,
            tag_name=data['tag_name'],
            published_at=data['published_at'],
            download_url=main_asset['browser_download_url'],
            changelog=data.get('body', ''),
            is_prerelease=data.get('prerelease', False),
            asset_name=main_asset['name'],
            asset_size=main_asset['size'],
            checksum=checksum
        )

    def _is_newer_version(self, version1: str, version2: str) -> bool:
        """Compare two semantic versions.

        Args:
            version1: First version (e.g., "1.2.3")
            version2: Second version (e.g., "1.1.0")

        Returns:
            True if version1 is newer than version2
        """
        def parse_version(v: str) -> tuple:
            # Remove 'v' prefix and split
            v = v.lstrip('v')
            parts = re.split(r'[-+]', v)[0]  # Remove pre-release/build metadata
            return tuple(map(int, parts.split('.')))

        try:
            v1_parts = parse_version(version1)
            v2_parts = parse_version(version2)
            return v1_parts > v2_parts
        except (ValueError, AttributeError):
            return False

    async def _verify_checksum(self, file_path: Path, expected_checksum: str) -> bool:
        """Verify file checksum.

        Args:
            file_path: Path to file to verify
            expected_checksum: Expected SHA256 checksum

        Returns:
            True if checksum matches, False otherwise
        """
        sha256 = hashlib.sha256()

        with open(file_path, 'rb') as f:
            while True:
                data = f.read(65536)  # 64KB chunks
                if not data:
                    break
                sha256.update(data)

        actual_checksum = sha256.hexdigest()
        return actual_checksum.lower() == expected_checksum.lower()

    async def _rebuild_pm3_client(self):
        """Rebuild Proxmark3 client after update."""
        pm3_dir = Path("/opt/proxmark3")

        if not pm3_dir.exists():
            print("Proxmark3 directory not found, skipping rebuild")
            return

        # Run make clean and make
        await self._run_command(f"cd {pm3_dir} && make clean && make")

    async def _update_version_file(self, version: str):
        """Update version file with new version.

        Args:
            version: New version string
        """
        version_file = Path("/opt/dangerous-pi/VERSION")
        version_file.write_text(version)

    async def _run_command(self, command: str) -> str:
        """Run a shell command.

        Args:
            command: Command to run

        Returns:
            Command output
        """
        process = await asyncio.create_subprocess_shell(
            command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )

        stdout, stderr = await process.communicate()

        if process.returncode != 0:
            raise Exception(f"Command failed: {stderr.decode()}")

        return stdout.decode()


# Global update manager instance
_update_manager: Optional[UpdateManager] = None


def get_update_manager() -> UpdateManager:
    """Get the global update manager instance."""
    global _update_manager
    if _update_manager is None:
        _update_manager = UpdateManager()
    return _update_manager