Borealis-Github-Replica/Data/Engine/services/API/devices/tunnel.py

# ======================================================
# Data\Engine\services\API\devices\tunnel.py
# Description: WireGuard VPN tunnel API (connect/status/disconnect).
#
# API Endpoints (if applicable):
# - POST /api/tunnel/connect (Token Authenticated) - Issues VPN session material for an agent.
# - GET /api/tunnel/status (Token Authenticated) - Returns VPN status for an agent.
# - DELETE /api/tunnel/disconnect (Token Authenticated) - Tears down VPN session for an agent.
# ======================================================

"""WireGuard VPN tunnel API (Engine side)."""
from __future__ import annotations

import os
from typing import Any, Dict, Optional, Tuple

from flask import Blueprint, jsonify, request, session
from itsdangerous import BadSignature, SignatureExpired, URLSafeTimedSerializer

from ...VPN import VpnTunnelService

if False:  # pragma: no cover - import cycle hint for type checkers
    from .. import EngineServiceAdapters


def _current_user(app) -> Optional[Dict[str, str]]:
    username = session.get("username")
    role = session.get("role") or "User"
    if username:
        return {"username": username, "role": role}

    token = None
    auth_header = request.headers.get("Authorization") or ""
    if auth_header.lower().startswith("bearer "):
        token = auth_header.split(" ", 1)[1].strip()
    if not token:
        token = request.cookies.get("borealis_auth")
    if not token:
        return None

    try:
        serializer = URLSafeTimedSerializer(app.secret_key or "borealis-dev-secret", salt="borealis-auth")
        token_ttl = int(os.environ.get("BOREALIS_TOKEN_TTL_SECONDS", 60 * 60 * 24 * 30))
        data = serializer.loads(token, max_age=token_ttl)
        username = data.get("u")
        role = data.get("r") or "User"
        if username:
            return {"username": username, "role": role}
    except (BadSignature, SignatureExpired, Exception):
        return None
    return None


def _require_login(app) -> Optional[Tuple[Dict[str, Any], int]]:
    user = _current_user(app)
    if not user:
        return {"error": "unauthorized"}, 401
    return None


def _get_tunnel_service(adapters: "EngineServiceAdapters") -> VpnTunnelService:
    service = getattr(adapters.context, "vpn_tunnel_service", None) or getattr(adapters, "_vpn_tunnel_service", None)
    if service is None:
        manager = getattr(adapters.context, "wireguard_server_manager", None)
        if manager is None:
            raise RuntimeError("wireguard_manager_unavailable")
        service = VpnTunnelService(
            context=adapters.context,
            wireguard_manager=manager,
            db_conn_factory=adapters.db_conn_factory,
            socketio=getattr(adapters.context, "socketio", None),
            service_log=adapters.service_log,
            signer=getattr(adapters, "script_signer", None),
        )
        setattr(adapters, "_vpn_tunnel_service", service)
        setattr(adapters.context, "vpn_tunnel_service", service)
    return service


def _normalize_text(value: Any) -> str:
    if value is None:
        return ""
    try:
        return str(value).strip()
    except Exception:
        return ""


def register_tunnel(app, adapters: "EngineServiceAdapters") -> None:
    blueprint = Blueprint("vpn_tunnel", __name__)
    logger = adapters.context.logger.getChild("vpn_tunnel.api")

    @blueprint.route("/api/tunnel/connect", methods=["POST"])
    def connect_tunnel():
        requirement = _require_login(app)
        if requirement:
            payload, status = requirement
            return jsonify(payload), status

        user = _current_user(app) or {}
        operator_id = user.get("username") or None

        body = request.get_json(silent=True) or {}
        agent_id = _normalize_text(body.get("agent_id"))
        if not agent_id:
            return jsonify({"error": "agent_id_required"}), 400

        try:
            tunnel_service = _get_tunnel_service(adapters)
            payload = tunnel_service.connect(agent_id=agent_id, operator_id=operator_id)
        except RuntimeError as exc:
            logger.warning("vpn connect failed for agent_id=%s: %s", agent_id, exc)
            return jsonify({"error": "connect_failed"}), 500

        return jsonify(payload), 200

    @blueprint.route("/api/tunnel/status", methods=["GET"])
    def tunnel_status():
        requirement = _require_login(app)
        if requirement:
            payload, status = requirement
            return jsonify(payload), status

        agent_id = _normalize_text(request.args.get("agent_id") or "")
        if not agent_id:
            return jsonify({"error": "agent_id_required"}), 400

        tunnel_service = _get_tunnel_service(adapters)
        payload = tunnel_service.status(agent_id)
        if not payload:
            return jsonify({"status": "down", "agent_id": agent_id}), 200
        payload["status"] = "up"
        bump = _normalize_text(request.args.get("bump") or "")
        if bump:
            tunnel_service.bump_activity(agent_id)
        return jsonify(payload), 200

    @blueprint.route("/api/tunnel/connect/status", methods=["GET"])
    def tunnel_connect_status():
        return tunnel_status()

    @blueprint.route("/api/tunnel/disconnect", methods=["DELETE"])
    def disconnect_tunnel():
        requirement = _require_login(app)
        if requirement:
            payload, status = requirement
            return jsonify(payload), status

        body = request.get_json(silent=True) or {}
        agent_id = _normalize_text(body.get("agent_id"))
        tunnel_id = _normalize_text(body.get("tunnel_id"))
        reason = _normalize_text(body.get("reason") or "operator_stop")

        tunnel_service = _get_tunnel_service(adapters)
        stopped = False
        if tunnel_id:
            stopped = tunnel_service.disconnect_by_tunnel(tunnel_id, reason=reason)
        elif agent_id:
            stopped = tunnel_service.disconnect(agent_id, reason=reason)
        else:
            return jsonify({"error": "agent_id_required"}), 400

        if not stopped:
            return jsonify({"error": "not_found"}), 404

        return jsonify({"status": "stopped", "reason": reason}), 200

    app.register_blueprint(blueprint)