headquarter/apps/api/src/services/ssh_keys.py

"""SSH key service utilities for preparing keys for container use."""

import logging
import os
import re
from pathlib import Path

from cryptography.fernet import Fernet

from src.config import Settings

logger = logging.getLogger(__name__)


def _get_fernet() -> Fernet:
    """Generate a valid Fernet key from the session secret."""
    import base64
    import hashlib

    settings = Settings()
    key_bytes = hashlib.sha256(settings.session_secret.encode()).digest()
    key = base64.urlsafe_b64encode(key_bytes)
    return Fernet(key)


def _sanitize_filename(name: str) -> str:
    """Sanitize a string for use as a filename.

    Replaces non-alphanumeric characters with underscores and strips
    leading/trailing underscores.
    """
    sanitized = re.sub(r"[^a-zA-Z0-9_-]", "_", name)
    sanitized = sanitized.strip("_")
    # Ensure it's not empty
    if not sanitized:
        sanitized = "key"
    return sanitized


def prepare_ssh_key_files(
    instance_dir: str,
    ssh_key,
    subdir: str = ".ssh",
    uid: int | None = None,
    gid: int | None = None,
    key_filename: str = "id_ed25519",
    write_config: bool = True,
) -> str:
    """Decrypt and write SSH key files to instance directory for container mounting.

    Args:
        instance_dir: Path to instance directory
        ssh_key: SSHKey model instance with encrypted private key
        subdir: Subdirectory within instance_dir to write to (default: ".ssh")
        uid: Optional UID to own the files (for bind-mount into non-root container)
        gid: Optional GID to own the files
        key_filename: Base filename for the key pair (default: "id_ed25519").
            The private key will be named "{key_filename}" and the public key
            "{key_filename}.pub".
        write_config: Whether to write an SSH config file (default: True).
            Set to False when combining multiple keys into one directory,
            then call write_ssh_config() separately.

    Returns:
        Path to the .ssh directory
    """
    ssh_dir = Path(instance_dir) / subdir
    ssh_dir.mkdir(parents=True, exist_ok=True)

    # Decrypt private key
    fernet = _get_fernet()
    private_key = fernet.decrypt(ssh_key.private_key_encrypted.encode()).decode()

    # Write private key with restricted permissions
    private_key_path = ssh_dir / key_filename
    private_key_path.write_text(private_key)
    os.chmod(private_key_path, 0o600)

    # Write public key
    public_key_path = ssh_dir / f"{key_filename}.pub"
    public_key_path.write_text(ssh_key.public_key)
    os.chmod(public_key_path, 0o644)

    # Write SSH config (only if requested)
    if write_config:
        config_path = ssh_dir / "config"
        config_content = f"""Host *
    StrictHostKeyChecking no
    UserKnownHostsFile /dev/null
    IdentityFile ~/.ssh/{key_filename}
    IdentitiesOnly yes
"""
        config_path.write_text(config_content)
        os.chmod(config_path, 0o644)

        # Set ownership to target container user if requested
        if uid is not None or gid is not None:
            effective_uid = uid if uid is not None else -1
            effective_gid = gid if gid is not None else -1
            try:
                os.chown(ssh_dir, effective_uid, effective_gid)
                os.chown(private_key_path, effective_uid, effective_gid)
                os.chown(public_key_path, effective_uid, effective_gid)
                os.chown(config_path, effective_uid, effective_gid)
                logger.debug(
                    "Set SSH key ownership to uid=%s gid=%s for %s",
                    effective_uid,
                    effective_gid,
                    ssh_dir,
                )
            except PermissionError as exc:
                logger.warning(
                    "Cannot chown SSH keys to uid=%s gid=%s (running as uid=%s): %s",
                    effective_uid,
                    effective_gid,
                    os.getuid(),
                    exc,
                )
    else:
        # Still chown the key files even if we didn't write config
        if uid is not None or gid is not None:
            effective_uid = uid if uid is not None else -1
            effective_gid = gid if gid is not None else -1
            try:
                os.chown(private_key_path, effective_uid, effective_gid)
                os.chown(public_key_path, effective_uid, effective_gid)
            except PermissionError:
                pass

    return str(ssh_dir)


def write_ssh_config(
    ssh_dir: str,
    key_filenames: list[str],
    uid: int | None = None,
    gid: int | None = None,
) -> None:
    """Write an SSH config file that includes multiple IdentityFile entries.

    Args:
        ssh_dir: Path to the .ssh directory
        key_filenames: List of key filenames (without .pub extension)
        uid: Optional UID to own the config file
        gid: Optional GID to own the config file
    """
    ssh_dir_path = Path(ssh_dir)
    ssh_dir_path.mkdir(parents=True, exist_ok=True)

    config_path = ssh_dir_path / "config"
    lines = ["Host *"]
    lines.append("    StrictHostKeyChecking no")
    lines.append("    UserKnownHostsFile /dev/null")
    lines.append("    IdentitiesOnly yes")
    for filename in key_filenames:
        lines.append(f"    IdentityFile ~/.ssh/{filename}")
    lines.append("")

    config_content = "\n".join(lines)
    config_path.write_text(config_content)
    os.chmod(config_path, 0o644)

    if uid is not None or gid is not None:
        effective_uid = uid if uid is not None else -1
        effective_gid = gid if gid is not None else -1
        try:
            os.chown(config_path, effective_uid, effective_gid)
        except PermissionError:
            pass


def cleanup_ssh_key_files(instance_dir: str) -> None:
    """Remove temporary SSH key files from instance directory.

    Args:
        instance_dir: Path to instance directory
    """
    ssh_dir = Path(instance_dir) / ".ssh"
    if ssh_dir.exists():
        for file_path in ssh_dir.iterdir():
            file_path.unlink()
        ssh_dir.rmdir()