clonebox 1.1.4__py3-none-any.whl → 1.1.6__py3-none-any.whl

clonebox/p2p.py

@@ -14,8 +14,10 @@ class P2PManager:
 
     def __init__(self, ssh_options: Optional[list] = None):
         self.ssh_options = ssh_options or [
-            "-o", "StrictHostKeyChecking=no",
-            "-o", "UserKnownHostsFile=/dev/null",
+            "-o",
+            "StrictHostKeyChecking=no",
+            "-o",
+            "UserKnownHostsFile=/dev/null",
         ]
 
     def _run_ssh(self, host: str, command: str) -> subprocess.CompletedProcess:

clonebox/resource_monitor.py

@@ -0,0 +1,162 @@
+"""Resource monitoring system for CloneBox."""
+
+import xml.etree.ElementTree as ET
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+try:
+    import libvirt
+except ImportError:
+    libvirt = None
+
+
+@dataclass
+class ResourceUsage:
+    """Current resource usage of a VM."""
+
+    timestamp: datetime
+    cpu_time_ns: int
+    cpu_percent: float
+    memory_used_bytes: int
+    memory_total_bytes: int
+    memory_percent: float
+    swap_used_bytes: int
+    disk_read_bytes: int
+    disk_write_bytes: int
+    disk_read_requests: int
+    disk_write_requests: int
+    net_rx_bytes: int
+    net_tx_bytes: int
+    net_rx_packets: int
+    net_tx_packets: int
+
+
+class ResourceMonitor:
+    """Monitor VM resource usage using libvirt."""
+
+    def __init__(self, conn: Optional[Any] = None):
+        self.conn = conn
+        self._prev_stats: Dict[str, Dict] = {}
+
+    def get_usage(self, vm_name: str) -> ResourceUsage:
+        """Get current resource usage for a VM."""
+        if not self.conn:
+            raise RuntimeError("libvirt connection not available")
+
+        domain = self.conn.lookupByName(vm_name)
+        if not domain.isActive():
+            raise RuntimeError(f"VM '{vm_name}' is not running")
+
+        # CPU stats
+        cpu_stats = domain.getCPUStats(True)[0]
+        cpu_time = cpu_stats.get("cpu_time", 0)
+        cpu_percent = self._calculate_cpu_percent(vm_name, cpu_time)
+
+        # Memory stats
+        # Need to ensure memory balloon driver is active for accurate stats
+        mem_stats = domain.memoryStats()
+        memory_used = mem_stats.get("rss", 0) * 1024  # RSS is often most accurate for host view
+        memory_total = mem_stats.get("actual", 1) * 1024
+        if "unused" in mem_stats:
+            memory_used = (mem_stats["actual"] - mem_stats["unused"]) * 1024
+        
+        memory_percent = (memory_used / memory_total * 100) if memory_total else 0
+        swap_used = mem_stats.get("swap_in", 0) * 1024
+
+        # Block and Network stats
+        disk_stats = self._get_disk_stats(domain)
+        net_stats = self._get_network_stats(domain)
+
+        return ResourceUsage(
+            timestamp=datetime.now(),
+            cpu_time_ns=cpu_time,
+            cpu_percent=cpu_percent,
+            memory_used_bytes=memory_used,
+            memory_total_bytes=memory_total,
+            memory_percent=memory_percent,
+            swap_used_bytes=swap_used,
+            **disk_stats,
+            **net_stats,
+        )
+
+    def _calculate_cpu_percent(self, vm_name: str, cpu_time: int) -> float:
+        """Calculate CPU percentage from time delta."""
+        import time
+
+        now = time.time()
+        prev = self._prev_stats.get(vm_name, {})
+        prev_time = prev.get("cpu_time", cpu_time)
+        prev_timestamp = prev.get("timestamp", now)
+
+        # Update stored stats
+        self._prev_stats[vm_name] = {
+            "cpu_time": cpu_time,
+            "timestamp": now,
+        }
+
+        time_delta = now - prev_timestamp
+        if time_delta <= 0:
+            return 0.0
+
+        cpu_delta = cpu_time - prev_time
+        # cpu_time is in nanoseconds, time_delta in seconds
+        # (delta_ns / (delta_sec * 1e9)) * 100
+        return (cpu_delta / (time_delta * 1e9)) * 100
+
+    def _get_disk_stats(self, domain) -> Dict[str, int]:
+        """Get aggregated disk stats."""
+        stats = {
+            "disk_read_bytes": 0,
+            "disk_write_bytes": 0,
+            "disk_read_requests": 0,
+            "disk_write_requests": 0,
+        }
+
+        xml = domain.XMLDesc()
+        tree = ET.fromstring(xml)
+
+        for disk in tree.findall(".//disk"):
+            target = disk.find("target")
+            if target is not None:
+                dev = target.get("dev")
+                try:
+                    # blockStats returns: (read_req, read_bytes, write_req, write_bytes, errs)
+                    ds = domain.blockStats(dev)
+                    stats["disk_read_requests"] += ds[0]
+                    stats["disk_read_bytes"] += ds[1]
+                    stats["disk_write_requests"] += ds[2]
+                    stats["disk_write_bytes"] += ds[3]
+                except Exception:
+                    continue
+
+        return stats
+
+    def _get_network_stats(self, domain) -> Dict[str, int]:
+        """Get aggregated network stats."""
+        stats = {
+            "net_rx_bytes": 0,
+            "net_tx_bytes": 0,
+            "net_rx_packets": 0,
+            "net_tx_packets": 0,
+        }
+
+        xml = domain.XMLDesc()
+        tree = ET.fromstring(xml)
+
+        for iface in tree.findall(".//interface"):
+            target = iface.find("target")
+            if target is not None:
+                dev = target.get("dev")
+                try:
+                    # interfaceStats returns: (rx_bytes, rx_packets, rx_errs, rx_drop, 
+                    #                          tx_bytes, tx_packets, tx_errs, tx_drop)
+                    ns = domain.interfaceStats(dev)
+                    stats["net_rx_bytes"] += ns[0]
+                    stats["net_rx_packets"] += ns[1]
+                    stats["net_tx_bytes"] += ns[4]
+                    stats["net_tx_packets"] += ns[5]
+                except Exception:
+                    continue
+
+        return stats

clonebox/resources.py

@@ -0,0 +1,222 @@
+"""
+Resource limits and monitoring for CloneBox VMs.
+"""
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import List, Optional, Union
+
+
+def parse_size(value: Union[str, int]) -> int:
+    """Parse size string like '8G', '512M' to bytes."""
+    if isinstance(value, int):
+        return value
+        
+    value = str(value).strip().upper()
+    multipliers = {
+        "K": 1024,
+        "M": 1024**2,
+        "G": 1024**3,
+        "T": 1024**4,
+    }
+
+    if value[-1] in multipliers:
+        try:
+            return int(float(value[:-1]) * multipliers[value[-1]])
+        except ValueError:
+            return 0
+    
+    try:
+        return int(value)
+    except ValueError:
+        return 0
+
+
+def parse_bandwidth(value: Union[str, int]) -> int:
+    """Parse bandwidth like '100Mbps' to bits/sec."""
+    if isinstance(value, int):
+        return value
+        
+    value = str(value).strip().lower()
+
+    if value.endswith("gbps"):
+        return int(float(value[:-4]) * 1_000_000_000)
+    elif value.endswith("mbps"):
+        return int(float(value[:-4]) * 1_000_000)
+    elif value.endswith("kbps"):
+        return int(float(value[:-4]) * 1_000)
+    elif value.endswith("bps"):
+        return int(float(value[:-3]))
+
+    try:
+        return int(value)
+    except ValueError:
+        return 0
+
+
+@dataclass
+class CPULimits:
+    """CPU resource limits."""
+
+    vcpus: int = 2
+    shares: int = 1024  # CFS shares (weight)
+    period: int = 100000  # CFS period (microseconds)
+    quota: Optional[int] = None  # CFS quota (microseconds)
+    pin: Optional[List[int]] = None  # CPU pinning
+
+    def get_max_percent(self) -> float:
+        """Get max CPU percentage."""
+        if self.quota:
+            return (self.quota / self.period) * 100
+        return self.vcpus * 100
+
+    def to_libvirt_xml(self) -> str:
+        """Generate libvirt cputune XML."""
+        elements = []
+        elements.append(f"    <shares>{self.shares}</shares>")
+
+        if self.quota:
+            elements.append(f"    <period>{self.period}</period>")
+            elements.append(f"    <quota>{self.quota}</quota>")
+
+        if self.pin:
+            for idx, cpu in enumerate(self.pin):
+                elements.append(f'    <vcpupin vcpu="{idx}" cpuset="{cpu}"/>')
+
+        if not elements:
+            return ""
+            
+        return f"  <cputune>\n{chr(10).join(elements)}\n  </cputune>"
+
+
+@dataclass
+class MemoryLimits:
+    """Memory resource limits."""
+
+    limit: str = "4G"  # Hard limit
+    soft_limit: Optional[str] = None  # Soft limit
+    swap: Optional[str] = None  # Swap limit
+    hugepages: bool = False  # Use hugepages
+
+    @property
+    def limit_bytes(self) -> int:
+        return parse_size(self.limit)
+
+    @property
+    def soft_limit_bytes(self) -> Optional[int]:
+        return parse_size(self.soft_limit) if self.soft_limit else None
+
+    @property
+    def swap_bytes(self) -> Optional[int]:
+        return parse_size(self.swap) if self.swap else None
+
+    def to_libvirt_xml(self) -> str:
+        """Generate libvirt memtune XML."""
+        elements = []
+
+        # Convert to KiB for libvirt
+        limit_kib = self.limit_bytes // 1024
+        elements.append(f"    <hard_limit unit='KiB'>{limit_kib}</hard_limit>")
+
+        if self.soft_limit_bytes:
+            soft_kib = self.soft_limit_bytes // 1024
+            elements.append(f"    <soft_limit unit='KiB'>{soft_kib}</soft_limit>")
+
+        if self.swap_bytes:
+            swap_kib = self.swap_bytes // 1024
+            elements.append(f"    <swap_hard_limit unit='KiB'>{swap_kib}</swap_hard_limit>")
+
+        return f"  <memtune>\n{chr(10).join(elements)}\n  </memtune>"
+
+
+@dataclass
+class DiskLimits:
+    """Disk I/O limits."""
+
+    read_bps: Optional[str] = None  # Read bytes/sec
+    write_bps: Optional[str] = None  # Write bytes/sec
+    read_iops: Optional[int] = None  # Read IOPS
+    write_iops: Optional[int] = None  # Write IOPS
+
+    @property
+    def read_bps_bytes(self) -> Optional[int]:
+        return parse_size(self.read_bps) if self.read_bps else None
+
+    @property
+    def write_bps_bytes(self) -> Optional[int]:
+        return parse_size(self.write_bps) if self.write_bps else None
+
+    def to_libvirt_xml(self) -> str:
+        """Generate libvirt iotune XML for disk device."""
+        elements = []
+
+        if self.read_bps_bytes:
+            elements.append(f"      <read_bytes_sec>{self.read_bps_bytes}</read_bytes_sec>")
+        if self.write_bps_bytes:
+            elements.append(f"      <write_bytes_sec>{self.write_bps_bytes}</write_bytes_sec>")
+        if self.read_iops:
+            elements.append(f"      <read_iops_sec>{self.read_iops}</read_iops_sec>")
+        if self.write_iops:
+            elements.append(f"      <write_iops_sec>{self.write_iops}</write_iops_sec>")
+
+        if elements:
+            return f"    <iotune>\n{chr(10).join(elements)}\n    </iotune>"
+        return ""
+
+
+@dataclass
+class NetworkLimits:
+    """Network bandwidth limits."""
+
+    inbound: Optional[str] = None  # Inbound bandwidth
+    outbound: Optional[str] = None  # Outbound bandwidth
+
+    @property
+    def inbound_kbps(self) -> Optional[int]:
+        if self.inbound:
+            return parse_bandwidth(self.inbound) // 1000  # Convert to kbps
+        return None
+
+    @property
+    def outbound_kbps(self) -> Optional[int]:
+        if self.outbound:
+            return parse_bandwidth(self.outbound) // 1000
+        return None
+
+    def to_libvirt_xml(self) -> str:
+        """Generate libvirt bandwidth XML for interface."""
+        elements = []
+
+        if self.inbound_kbps:
+            # average in KB/s (libvirt uses kbps but sometimes expects KB/s depending on version,
+            # usually it's average in kbytes/s)
+            avg_kbs = self.inbound_kbps // 8
+            elements.append(f'      <inbound average="{avg_kbs}"/>')
+
+        if self.outbound_kbps:
+            avg_kbs = self.outbound_kbps // 8
+            elements.append(f'      <outbound average="{avg_kbs}"/>')
+
+        if elements:
+            return f"    <bandwidth>\n{chr(10).join(elements)}\n    </bandwidth>"
+        return ""
+
+
+@dataclass
+class ResourceLimits:
+    """Combined resource limits for a VM."""
+
+    cpu: CPULimits = field(default_factory=CPULimits)
+    memory: MemoryLimits = field(default_factory=MemoryLimits)
+    disk: DiskLimits = field(default_factory=DiskLimits)
+    network: NetworkLimits = field(default_factory=NetworkLimits)
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "ResourceLimits":
+        """Create from config dict."""
+        return cls(
+            cpu=CPULimits(**data.get("cpu", {})),
+            memory=MemoryLimits(**data.get("memory", {})),
+            disk=DiskLimits(**data.get("disk", {})),
+            network=NetworkLimits(**data.get("network", {})),
+        )

clonebox/rollback.py

@@ -0,0 +1,172 @@
+"""
+Transactional rollback support for CloneBox operations.
+"""
+
+import logging
+import shutil
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable, List, Optional
+
+log = logging.getLogger(__name__)
+
+
+@dataclass
+class RollbackAction:
+    """A single rollback action."""
+
+    description: str
+    action: Callable[[], None]
+    critical: bool = True  # If True, failure stops rollback chain
+
+
+@dataclass
+class RollbackContext:
+    """
+    Context manager for transactional operations with automatic rollback.
+
+    Usage:
+        with RollbackContext("create VM") as ctx:
+            ctx.add_file(disk_path)  # Will be deleted on error
+            ctx.add_directory(vm_dir)  # Will be deleted on error
+            ctx.add_action("stop VM", lambda: cloner.stop_vm(name))
+
+            # If any exception occurs, all registered items are cleaned up
+            do_risky_operation()
+    """
+
+    operation_name: str
+    _files: List[Path] = field(default_factory=list)
+    _directories: List[Path] = field(default_factory=list)
+    _actions: List[RollbackAction] = field(default_factory=list)
+    _committed: bool = False
+    _console: Optional[Any] = None
+
+    def add_file(self, path: Path, description: Optional[str] = None) -> Path:
+        """Register a file for cleanup on rollback."""
+        self._files.append(path)
+        log.debug(f"Registered file for rollback: {path}")
+        return path
+
+    def add_directory(self, path: Path, description: Optional[str] = None) -> Path:
+        """Register a directory for cleanup on rollback."""
+        self._directories.append(path)
+        log.debug(f"Registered directory for rollback: {path}")
+        return path
+
+    def add_action(
+        self, description: str, action: Callable[[], None], critical: bool = False
+    ) -> None:
+        """Register a custom rollback action."""
+        self._actions.append(
+            RollbackAction(description=description, action=action, critical=critical)
+        )
+        log.debug(f"Registered action for rollback: {description}")
+
+    def add_libvirt_domain(self, conn, domain_name: str) -> None:
+        """Register a libvirt domain for cleanup."""
+
+        def cleanup_domain():
+            try:
+                dom = conn.lookupByName(domain_name)
+                if dom.isActive():
+                    dom.destroy()
+                dom.undefine()
+            except Exception as e:
+                log.warning(f"Failed to cleanup domain {domain_name}: {e}")
+
+        self._actions.append(
+            RollbackAction(
+                description=f"undefine domain {domain_name}",
+                action=cleanup_domain,
+                critical=False,
+            )
+        )
+
+    def commit(self) -> None:
+        """Mark operation as successful, preventing rollback."""
+        self._committed = True
+        log.info(f"Operation '{self.operation_name}' committed successfully")
+
+    def rollback(self) -> List[str]:
+        """Execute rollback actions. Returns list of errors."""
+        errors = []
+
+        if self._console:
+            self._console.print(
+                f"[yellow]Rolling back '{self.operation_name}'...[/yellow]"
+            )
+
+        # Execute custom actions first (in reverse order)
+        for action in reversed(self._actions):
+            try:
+                log.info(f"Rollback action: {action.description}")
+                action.action()
+            except Exception as e:
+                error_msg = f"Rollback action '{action.description}' failed: {e}"
+                errors.append(error_msg)
+                log.error(error_msg)
+                if action.critical:
+                    break
+
+        # Delete files
+        for path in reversed(self._files):
+            try:
+                if path.exists():
+                    path.unlink()
+                    log.info(f"Deleted file: {path}")
+            except Exception as e:
+                errors.append(f"Failed to delete {path}: {e}")
+
+        # Delete directories
+        for path in reversed(self._directories):
+            try:
+                if path.exists():
+                    shutil.rmtree(path)
+                    log.info(f"Deleted directory: {path}")
+            except Exception as e:
+                errors.append(f"Failed to delete {path}: {e}")
+
+        return errors
+
+    def __enter__(self) -> "RollbackContext":
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
+        if exc_type is not None and not self._committed:
+            errors = self.rollback()
+            if errors and self._console:
+                self._console.print("[red]Rollback completed with errors:[/red]")
+                for error in errors:
+                    self._console.print(f"  [dim]- {error}[/dim]")
+        return False  # Don't suppress the exception
+
+
+@contextmanager
+def vm_creation_transaction(cloner: Any, config: Any, console: Optional[Any] = None):
+    """
+    Context manager for VM creation with automatic rollback.
+
+    Usage:
+        with vm_creation_transaction(cloner, config, console) as ctx:
+            vm_dir = ctx.add_directory(images_dir / config.name)
+            vm_dir.mkdir(parents=True, exist_ok=True)
+
+            disk_path = ctx.add_file(vm_dir / "root.qcow2")
+            create_disk(disk_path)
+
+            ctx.add_libvirt_domain(cloner.conn, config.name)
+            cloner.conn.defineXML(xml)
+
+            ctx.commit()  # Success!
+    """
+    ctx = RollbackContext(
+        operation_name=f"create VM '{config.name}'", _console=console
+    )
+    try:
+        yield ctx
+    except Exception:
+        if not ctx._committed:
+            ctx.rollback()
+        raise