mimic-big/backend/src/mimic/api/hosts.py

"""Host CRUD endpoints (scoped under an engagement)."""

from __future__ import annotations

from flask import Blueprint, abort, jsonify
from flask.typing import ResponseReturnValue
from sqlalchemy import select

from mimic.api._helpers import jsonify_model, parse_body, parse_uuid
from mimic.db.models import Engagement, Host
from mimic.db.types import HostStatus
from mimic.extensions import db
from mimic.rbac import Permission, require_perm
from mimic.schemas import HostCreate, HostRead, HostUpdate

bp = Blueprint("hosts", __name__)


def _engagement_or_404(eid: str) -> Engagement:
    engagement = db.session.get(Engagement, parse_uuid(eid, field="engagement id"))
    if engagement is None:
        abort(404)
    return engagement


@bp.get("/engagements/<eid>/hosts")
@require_perm(Permission.HOST_CRUD)
def list_hosts(eid: str) -> ResponseReturnValue:
    engagement = _engagement_or_404(eid)
    stmt = select(Host).where(Host.engagement_id == engagement.id).order_by(Host.hostname)
    rows = db.session.execute(stmt).scalars().all()
    return jsonify([HostRead.model_validate(row).model_dump(mode="json") for row in rows])


@bp.post("/engagements/<eid>/hosts")
@require_perm(Permission.HOST_CRUD)
def create_host(eid: str) -> ResponseReturnValue:
    engagement = _engagement_or_404(eid)
    payload = parse_body(HostCreate)
    host = Host(
        engagement_id=engagement.id,
        hostname=payload.hostname,
        ip=payload.ip,
        os=payload.os,
        c2_session_id=payload.c2_session_id,
        c2_type=payload.c2_type,
        status=HostStatus.UNKNOWN,
    )
    db.session.add(host)
    db.session.commit()
    return jsonify_model(HostRead.model_validate(host), status=201)


@bp.put("/engagements/<eid>/hosts/<hid>")
@require_perm(Permission.HOST_CRUD)
def update_host(eid: str, hid: str) -> ResponseReturnValue:
    engagement = _engagement_or_404(eid)
    host = db.session.get(Host, parse_uuid(hid, field="host id"))
    if host is None or host.engagement_id != engagement.id:
        abort(404)
    payload = parse_body(HostUpdate)
    for field, value in payload.model_dump(exclude_unset=True).items():
        setattr(host, field, value)
    db.session.commit()
    return jsonify_model(HostRead.model_validate(host))


@bp.delete("/engagements/<eid>/hosts/<hid>")
@require_perm(Permission.HOST_CRUD)
def delete_host(eid: str, hid: str) -> ResponseReturnValue:
    engagement = _engagement_or_404(eid)
    host = db.session.get(Host, parse_uuid(hid, field="host id"))
    if host is None or host.engagement_id != engagement.id:
        abort(404)
    db.session.delete(host)
    db.session.commit()
    return "", 204