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

import base64
import uuid

from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from src.auth.dependencies import _get_user, get_current_user_id, get_db_session
from src.config import Settings
from src.models import SSHKey
from src.schemas.project import (
    SSHKeyCreate,
    SSHKeyResponse,
    SignPayloadRequest,
    SignatureResponse,
    VerifySignatureRequest,
    VerifySignatureResponse,
)

router = APIRouter(prefix="/ssh-keys", tags=["ssh-keys"])


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

    settings = Settings()
    # Derive a 32-byte key from the session secret using SHA256
    key_bytes = hashlib.sha256(settings.session_secret.encode()).digest()
    # Base64 encode it for Fernet (must be 32 url-safe base64-encoded bytes)
    key = base64.urlsafe_b64encode(key_bytes)
    return Fernet(key)


def generate_ssh_key_pair() -> tuple[str, str]:
    """Generate a new Ed25519 SSH key pair.

    Returns:
        Tuple of (private_key, public_key) as strings.
    """
    private_key = Ed25519PrivateKey.generate()
    public_key = private_key.public_key()

    private_bytes = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.OpenSSH,
        encryption_algorithm=serialization.NoEncryption(),
    )

    public_bytes = public_key.public_bytes(
        encoding=serialization.Encoding.OpenSSH,
        format=serialization.PublicFormat.OpenSSH,
    )

    return private_bytes.decode("utf-8"), public_bytes.decode("utf-8")


@router.post(
    "",
    response_model=SSHKeyResponse,
    status_code=status.HTTP_201_CREATED,
    summary="Create SSH key",
    description="Generate a new Ed25519 SSH key pair for the authenticated user.",
)
async def create_ssh_key(
    data: SSHKeyCreate,
    user_id: uuid.UUID = Depends(get_current_user_id),
    session: AsyncSession = Depends(get_db_session),
) -> SSHKey:
    """Create a new SSH key pair.

    Args:
        data: SSH key creation data including the key name.
        user_id: ID of the authenticated user.
        session: Database session.

    Returns:
        The newly created SSH key with public key exposed.
    """
    user = await _get_user(session, user_id)
    private_key, public_key = generate_ssh_key_pair()

    fernet = _get_fernet()
    encrypted_private = fernet.encrypt(private_key.encode()).decode()

    ssh_key = SSHKey(
        name=data.name,
        public_key=public_key,
        private_key_encrypted=encrypted_private,
        user_id=user.id,
    )
    session.add(ssh_key)
    await session.commit()
    await session.refresh(ssh_key)
    return ssh_key


@router.get(
    "",
    response_model=list[SSHKeyResponse],
    summary="List SSH keys",
    description="List all SSH keys for the authenticated user.",
)
async def list_ssh_keys(
    user_id: uuid.UUID = Depends(get_current_user_id),
    session: AsyncSession = Depends(get_db_session),
) -> list[SSHKey]:
    """List all SSH keys for the authenticated user.

    Args:
        user_id: ID of the authenticated user.
        session: Database session.

    Returns:
        List of SSH keys owned by the user.
    """
    user = await _get_user(session, user_id)
    result = await session.execute(select(SSHKey).where(SSHKey.user_id == user.id))
    return list(result.scalars().all())


@router.delete(
    "/{key_id}",
    status_code=status.HTTP_204_NO_CONTENT,
    summary="Delete SSH key",
    description="Delete an SSH key by ID.",
)
async def delete_ssh_key(
    key_id: uuid.UUID,
    user_id: uuid.UUID = Depends(get_current_user_id),
    session: AsyncSession = Depends(get_db_session),
) -> None:
    """Delete an SSH key.

    Args:
        key_id: UUID of the SSH key to delete.
        user_id: ID of the authenticated user.
        session: Database session.

    Returns:
        None with 204 status code.
    """
    user = await _get_user(session, user_id)
    ssh_key = await session.get(SSHKey, key_id)
    if ssh_key is None or ssh_key.user_id != user.id:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="ssh key not found"
        )

    await session.delete(ssh_key)
    await session.commit()


@router.post(
    "/{key_id}/sign",
    response_model=SignatureResponse,
    summary="Sign payload",
    description="Sign a payload using the SSH private key.",
)
async def sign_payload(
    key_id: uuid.UUID,
    data: SignPayloadRequest,
    user_id: uuid.UUID = Depends(get_current_user_id),
    session: AsyncSession = Depends(get_db_session),
) -> SignatureResponse:
    """Sign a payload with an SSH key.

    Args:
        key_id: UUID of the SSH key to use for signing.
        data: Sign request containing the payload string.
        user_id: ID of the authenticated user.
        session: Database session.

    Returns:
        Base64-encoded Ed25519 signature.
    """
    user = await _get_user(session, user_id)
    ssh_key = await session.get(SSHKey, key_id)
    if ssh_key is None or ssh_key.user_id != user.id:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="ssh key not found"
        )

    fernet = _get_fernet()
    private_key_pem = fernet.decrypt(ssh_key.private_key_encrypted.encode()).decode()

    private_key = serialization.load_ssh_private_key(
        private_key_pem.encode(), password=None
    )

    signature = private_key.sign(data.payload.encode())
    return SignatureResponse(signature=base64.b64encode(signature).decode())


@router.post(
    "/{key_id}/verify",
    response_model=VerifySignatureResponse,
    summary="Verify signature",
    description="Verify a signature against a payload using the SSH public key.",
)
async def verify_signature(
    key_id: uuid.UUID,
    data: VerifySignatureRequest,
    user_id: uuid.UUID = Depends(get_current_user_id),
    session: AsyncSession = Depends(get_db_session),
) -> VerifySignatureResponse:
    """Verify a signature with an SSH key's public key.

    Args:
        key_id: UUID of the SSH key to use for verification.
        data: Verify request containing payload and base64-encoded signature.
        user_id: ID of the authenticated user.
        session: Database session.

    Returns:
        Whether the signature is valid.
    """
    user = await _get_user(session, user_id)
    ssh_key = await session.get(SSHKey, key_id)
    if ssh_key is None or ssh_key.user_id != user.id:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="ssh key not found"
        )

    public_key = serialization.load_ssh_public_key(ssh_key.public_key.encode())

    try:
        signature = base64.b64decode(data.signature)
        public_key.verify(signature, data.payload.encode())
        return VerifySignatureResponse(valid=True)
    except Exception:
        return VerifySignatureResponse(valid=False)