pyssh-client/app/key_manager.py

"""Gestion des cles SSH : generation, import, export, registre local."""
from __future__ import annotations

import json
import shutil
import uuid
from dataclasses import dataclass, field, asdict
from pathlib import Path

import paramiko
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa, ed25519

from app.paths import DATA_DIR  # noqa: E402  (chemins valables aussi en .exe)
KEYS_FILE = DATA_DIR / "keys.json"
DEFAULT_KEY_DIR = DATA_DIR / "keys"

KEY_TYPE_ED25519 = "ed25519"
KEY_TYPE_RSA = "rsa"


@dataclass
class KeyEntry:
    id: str = field(default_factory=lambda: str(uuid.uuid4()))
    name: str = ""
    key_type: str = KEY_TYPE_ED25519
    private_path: str = ""
    public_path: str = ""
    fingerprint: str = ""
    public_key_text: str = ""

    def to_dict(self) -> dict:
        return asdict(self)

    @staticmethod
    def from_dict(data: dict) -> "KeyEntry":
        e = KeyEntry()
        for k in KeyEntry.__dataclass_fields__:
            if k in data:
                setattr(e, k, data[k])
        return e


class KeyManager:
    def __init__(self, registry_path: Path = KEYS_FILE):
        self.registry_path = registry_path
        self.keys: list[KeyEntry] = []
        self.load()

    # ---------- registre ----------
    def load(self) -> None:
        DATA_DIR.mkdir(parents=True, exist_ok=True)
        DEFAULT_KEY_DIR.mkdir(parents=True, exist_ok=True)
        if not self.registry_path.exists():
            self.keys = []
            self.save()
            return
        try:
            raw = json.loads(self.registry_path.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, OSError):
            raw = {"keys": []}
        self.keys = [KeyEntry.from_dict(k) for k in raw.get("keys", [])]

    def save(self) -> None:
        DATA_DIR.mkdir(parents=True, exist_ok=True)
        payload = {"keys": [k.to_dict() for k in self.keys]}
        self.registry_path.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8")

    def remove(self, key_id: str, delete_files: bool = False) -> None:
        entry = next((k for k in self.keys if k.id == key_id), None)
        if entry and delete_files:
            for p in (entry.private_path, entry.public_path):
                try:
                    if p and Path(p).exists():
                        Path(p).unlink()
                except OSError:
                    pass
        self.keys = [k for k in self.keys if k.id != key_id]
        self.save()

    # ---------- generation ----------
    def generate_key(
        self,
        name: str,
        key_type: str = KEY_TYPE_ED25519,
        rsa_bits: int = 4096,
        passphrase: str | None = None,
        save_dir: Path | None = None,
        comment: str = "",
    ) -> KeyEntry:
        save_dir = save_dir or DEFAULT_KEY_DIR
        save_dir.mkdir(parents=True, exist_ok=True)
        safe_name = "".join(c for c in name if c.isalnum() or c in ("-", "_")) or "id_key"
        filename = f"id_{key_type}_{safe_name}"
        priv_path = save_dir / filename
        pub_path = save_dir / f"{filename}.pub"

        if key_type == KEY_TYPE_ED25519:
            private_key = ed25519.Ed25519PrivateKey.generate()
        else:
            private_key = rsa.generate_private_key(public_exponent=65537, key_size=rsa_bits)

        encryption = (
            serialization.BestAvailableEncryption(passphrase.encode("utf-8"))
            if passphrase
            else serialization.NoEncryption()
        )
        priv_bytes = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.OpenSSH,
            encryption_algorithm=encryption,
        )
        pub_bytes = private_key.public_key().public_bytes(
            encoding=serialization.Encoding.OpenSSH,
            format=serialization.PublicFormat.OpenSSH,
        )
        comment_suffix = f" {comment}".rstrip() if comment else ""
        pub_text = pub_bytes.decode("utf-8") + comment_suffix + "\n"

        priv_path.write_bytes(priv_bytes)
        pub_path.write_text(pub_text, encoding="utf-8")
        try:
            priv_path.chmod(0o600)
        except OSError:
            pass

        fingerprint = self._fingerprint(pub_bytes)
        entry = KeyEntry(
            name=name,
            key_type=key_type,
            private_path=str(priv_path),
            public_path=str(pub_path),
            fingerprint=fingerprint,
            public_key_text=pub_text.strip(),
        )
        self.keys.append(entry)
        self.save()
        return entry

    # ---------- import ----------
    def import_key(self, name: str, private_key_path: str, public_key_path: str | None = None) -> KeyEntry:
        priv_path = Path(private_key_path)
        key_type, pub_text = self._inspect_private_key(priv_path)

        pub_path_str = ""
        if public_key_path and Path(public_key_path).exists():
            pub_path_str = public_key_path
            pub_text = Path(public_key_path).read_text(encoding="utf-8").strip()
        else:
            guess = priv_path.with_suffix(priv_path.suffix + ".pub")
            if guess.exists():
                pub_path_str = str(guess)
                pub_text = guess.read_text(encoding="utf-8").strip()

        fingerprint = pub_text.split()[1] if pub_text and len(pub_text.split()) > 1 else ""
        entry = KeyEntry(
            name=name,
            key_type=key_type,
            private_path=str(priv_path),
            public_path=pub_path_str,
            fingerprint=fingerprint[:16],
            public_key_text=pub_text,
        )
        self.keys.append(entry)
        self.save()
        return entry

    def _inspect_private_key(self, priv_path: Path) -> tuple[str, str]:
        """Tente de determiner le type de cle sans mot de passe (best effort)."""
        for loader, ktype in (
            (paramiko.Ed25519Key, KEY_TYPE_ED25519),
            (paramiko.RSAKey, KEY_TYPE_RSA),
        ):
            try:
                pkey = loader.from_private_key_file(str(priv_path))
                pub_text = f"{pkey.get_name()} {pkey.get_base64()}"
                return ktype, pub_text
            except paramiko.PasswordRequiredException:
                return ktype, ""
            except (paramiko.SSHException, OSError):
                continue
        return "unknown", ""

    # ---------- export ----------
    def export_public_key(self, entry: KeyEntry, destination: str) -> None:
        Path(destination).write_text(entry.public_key_text + "\n", encoding="utf-8")

    def export_private_key(self, entry: KeyEntry, destination: str) -> None:
        shutil.copyfile(entry.private_path, destination)

    @staticmethod
    def _fingerprint(pub_bytes: bytes) -> str:
        import base64
        import hashlib

        parts = pub_bytes.split()
        if len(parts) < 2:
            return ""
        raw = base64.b64decode(parts[1])
        digest = hashlib.sha256(raw).digest()
        return "SHA256:" + base64.b64encode(digest).decode("utf-8").rstrip("=")