htcli 1.1.0__py3-none-any.whl

src/htcli/utils/wallet/migration.py

@@ -0,0 +1,159 @@
+"""
+Wallet layout migration utilities.
+
+Used by both the CLI command (`htcli wallet upgrade`) and the standalone helper
+script in `scripts/migrate_wallet_layout.py`.
+"""
+
+from __future__ import annotations
+
+import json
+import shutil
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable, Iterable, Optional
+
+from .crypto import (
+    COLDKEY_FILE_SUFFIX,
+    HOTKEY_FILE_SUFFIX,
+    get_wallet_directory,
+)
+
+LogFn = Optional[Callable[[str], None]]
+
+
+def _load_json(path: Path) -> Optional[dict]:
+    try:
+        with open(path) as f:
+            return json.load(f)
+    except Exception:
+        return None
+
+
+def _ensure_backup(path: Path, backup_enabled: bool):
+    if not backup_enabled:
+        return
+    backup_path = path.with_suffix(path.suffix + ".bak")
+    if not backup_path.exists():
+        shutil.copy2(path, backup_path)
+
+
+@dataclass
+class MigrationReport:
+    migrated: list[str] = field(default_factory=list)
+    skipped: list[str] = field(default_factory=list)
+    warnings: list[str] = field(default_factory=list)
+
+    def as_dict(self) -> dict:
+        return {
+            "migrated": self.migrated,
+            "skipped": self.skipped,
+            "warnings": self.warnings,
+        }
+
+
+def migrate_wallet_layout(
+    wallet_dir: Optional[Path] = None,
+    *,
+    dry_run: bool = False,
+    force: bool = False,
+    backup: bool = True,
+    only_wallets: Optional[Iterable[str]] = None,
+    log: LogFn = None,
+) -> MigrationReport:
+    """Migrate flat wallet files into the hierarchical layout."""
+
+    def emit(message: str):
+        if log:
+            log(message)
+
+    target_dir = (wallet_dir or get_wallet_directory()).expanduser()
+    report = MigrationReport()
+
+    if not target_dir.exists():
+        warning = f"Wallet directory '{target_dir}' does not exist."
+        report.warnings.append(warning)
+        emit(f"[yellow]{warning}[/yellow]")
+        return report
+
+    selected = {name for name in only_wallets} if only_wallets else None
+
+    legacy_coldkeys = sorted(target_dir.glob(f"*{COLDKEY_FILE_SUFFIX}"))
+    legacy_hotkeys = sorted(target_dir.glob(f"*{HOTKEY_FILE_SUFFIX}"))
+    coldkey_address_map: dict[str, str] = {}
+
+    for coldkey_file in legacy_coldkeys:
+        data = _load_json(coldkey_file)
+        if not data:
+            warning = f"Unable to read coldkey file '{coldkey_file.name}', skipping."
+            report.warnings.append(warning)
+            emit(f"[yellow]{warning}[/yellow]")
+            continue
+
+        coldkey_name = data.get("name") or coldkey_file.stem.replace(
+            COLDKEY_FILE_SUFFIX, ""
+        )
+        if selected and coldkey_name not in selected:
+            continue
+
+        coldkey_address = data.get("ss58_address") or data.get("address")
+        if coldkey_address:
+            coldkey_address_map[coldkey_address] = coldkey_name
+
+        target = target_dir / coldkey_name / coldkey_file.name
+        if target.exists() and not force:
+            msg = f"Target coldkey '{target}' exists. Use --force to overwrite."
+            report.skipped.append(str(target))
+            emit(f"[cyan]{msg}[/cyan]")
+            continue
+
+        emit(f"[green]Move[/green] {coldkey_file.name} -> {target}")
+        report.migrated.append(f"{coldkey_file} -> {target}")
+        if dry_run:
+            continue
+        target.parent.mkdir(parents=True, exist_ok=True)
+        _ensure_backup(coldkey_file, backup)
+        shutil.move(str(coldkey_file), target)
+
+    for hotkey_file in legacy_hotkeys:
+        data = _load_json(hotkey_file)
+        if not data:
+            warning = f"Unable to read hotkey file '{hotkey_file.name}', skipping."
+            report.warnings.append(warning)
+            emit(f"[yellow]{warning}[/yellow]")
+            continue
+
+        owner_address = data.get("owner_address")
+        coldkey_name = coldkey_address_map.get(owner_address)
+        if selected and coldkey_name and coldkey_name not in selected:
+            continue
+
+        if not coldkey_name:
+            warning = (
+                f"Cannot resolve coldkey for hotkey '{hotkey_file.name}'. "
+                "Ensure the coldkey is migrated first."
+            )
+            report.warnings.append(warning)
+            emit(f"[yellow]{warning}[/yellow]")
+            continue
+
+        target = target_dir / coldkey_name / "hotkeys" / hotkey_file.name
+        if target.exists() and not force:
+            msg = f"Target hotkey '{target}' exists. Use --force to overwrite."
+            report.skipped.append(str(target))
+            emit(f"[cyan]{msg}[/cyan]")
+            continue
+
+        emit(f"[green]Move[/green] {hotkey_file.name} -> {target}")
+        report.migrated.append(f"{hotkey_file} -> {target}")
+        if dry_run:
+            continue
+        target.parent.mkdir(parents=True, exist_ok=True)
+        _ensure_backup(hotkey_file, backup)
+        shutil.move(str(hotkey_file), target)
+
+    if not report.migrated and not report.skipped and not report.warnings:
+        emit("[cyan]No legacy wallet files found. Nothing to do.[/cyan]")
+
+    return report
+