1import asyncio
  2import logging
  3from typing import Annotated
  4
  5import sqlalchemy
  6from fastapi import APIRouter, Body, Query
  7from fastapi import HTTPException, WebSocket
  8from fastapi_pagination.ext.sqlalchemy import paginate
  9from fastapi_pagination.links import Page
 10from starlette.websockets import WebSocketDisconnect, WebSocketState
 11from grpclib.exceptions import StreamTerminatedError, GRPCError
 12
 13from app import marznode
 14from app.db import crud, get_tls_certificate
 15from app.db.models import Node
 16from app.dependencies import (
 17    DBDep,
 18    SudoAdminDep,
 19    EndDateDep,
 20    StartDateDep,
 21    get_admin,
 22)
 23from app.models.node import (
 24    NodeCreate,
 25    NodeModify,
 26    NodeResponse,
 27    NodeSettings,
 28    NodeStatus,
 29    BackendConfig,
 30    BackendStats,
 31)
 32from app.models.system import TrafficUsageSeries
 33
 34logger = logging.getLogger(__name__)
 35router = APIRouter(prefix="/nodes", tags=["Node"])
 36
 37
 38@router.get("", response_model=Page[NodeResponse])
 39def get_nodes(
 40    db: DBDep,
 41    admin: SudoAdminDep,
 42    status: list[NodeStatus] = Query(None),
 43    name: str = Query(None),
 44):
 45    query = db.query(Node)
 46
 47    if name:
 48        query = query.filter(Node.name.ilike(f"%{name}%"))
 49
 50    if status:
 51        query = query.filter(Node.status.in_(status))
 52
 53    return paginate(db, query)
 54
 55
 56@router.post("", response_model=NodeResponse)
 57async def add_node(new_node: NodeCreate, db: DBDep, admin: SudoAdminDep):
 58    try:
 59        db_node = crud.create_node(db, new_node)
 60    except sqlalchemy.exc.IntegrityError:
 61        db.rollback()
 62        raise HTTPException(
 63            status_code=409, detail=f'Node "{new_node.name}" already exists'
 64        )
 65    certificate = get_tls_certificate(db)
 66
 67    await marznode.operations.add_node(db_node, certificate)
 68
 69    logger.info("New node `%s` added", db_node.name)
 70    return db_node
 71
 72
 73@router.get("/settings", response_model=NodeSettings)
 74def get_node_settings(db: DBDep, admin: SudoAdminDep):
 75    tls = crud.get_tls_certificate(db)
 76
 77    return NodeSettings(certificate=tls.certificate)
 78
 79
 80@router.get("/{node_id}", response_model=NodeResponse)
 81def get_node(node_id: int, db: DBDep, admin: SudoAdminDep):
 82    db_node = crud.get_node_by_id(db, node_id)
 83    if not db_node:
 84        raise HTTPException(status_code=404, detail="Node not found")
 85
 86    return db_node
 87
 88
 89@router.websocket("/{node_id}/{backend}/logs")
 90async def node_logs(
 91    node_id: int,
 92    backend: str,
 93    websocket: WebSocket,
 94    db: DBDep,
 95    include_buffer: bool = True,
 96):
 97    token = websocket.query_params.get("token", "") or websocket.headers.get(
 98        "Authorization", ""
 99    ).removeprefix("Bearer ")
100    admin = get_admin(db, token)
101
102    if not admin or not admin.is_sudo:
103        return await websocket.close(reason="You're not allowed", code=4403)
104
105    if not marznode.nodes.get(node_id):
106        return await websocket.close(reason="Node not found", code=4404)
107
108    await websocket.accept()
109    try:
110        async for line in marznode.nodes[node_id].get_logs(
111            name=backend, include_buffer=include_buffer
112        ):
113            await websocket.send_text(line)
114    except WebSocketDisconnect:
115        logger.debug("websocket disconnected")
116    except (StreamTerminatedError, GRPCError):
117        logger.info("node %i detached", node_id)
118    finally:
119        if websocket.state == WebSocketState.CONNECTED:
120            await websocket.close()
121
122
123@router.put("/{node_id}", response_model=NodeResponse)
124async def modify_node(
125    node_id: int, modified_node: NodeModify, db: DBDep, admin: SudoAdminDep
126):
127    db_node = crud.get_node_by_id(db, node_id)
128    if not db_node:
129        raise HTTPException(status_code=404, detail="Node not found")
130
131    db_node = crud.update_node(db, db_node, modified_node)
132
133    await marznode.operations.remove_node(db_node.id)
134    if db_node.status != NodeStatus.disabled:
135        certificate = get_tls_certificate(db)
136        await marznode.operations.add_node(db_node, certificate)
137
138    logger.info("Node `%s` modified", db_node.name)
139    return db_node
140
141
142@router.delete("/{node_id}")
143async def remove_node(node_id: int, db: DBDep, admin: SudoAdminDep):
144    db_node = crud.get_node_by_id(db, node_id)
145    if not db_node:
146        raise HTTPException(status_code=404, detail="Node not found")
147
148    crud.remove_node(db, db_node)
149    await marznode.operations.remove_node(db_node.id)
150
151    logger.info(f"Node `%s` deleted", db_node.name)
152    return {}
153
154
155@router.post("/{node_id}/resync")
156async def reconnect_node(node_id: int, db: DBDep, admin: SudoAdminDep):
157    db_node = crud.get_node_by_id(db, node_id)
158    if not db_node:
159        raise HTTPException(status_code=404, detail="Node not found")
160
161    return {}
162
163
164@router.get("/{node_id}/usage", response_model=TrafficUsageSeries)
165def get_usage(
166    node_id: int,
167    db: DBDep,
168    admin: SudoAdminDep,
169    start_date: StartDateDep,
170    end_date: EndDateDep,
171):
172    """
173    Get nodes usage
174    """
175    node = crud.get_node_by_id(db, node_id)
176    if not node:
177        raise HTTPException(status_code=404, detail="Node not found")
178
179    return crud.get_node_usage(db, start_date, end_date, node)
180
181
182@router.get("/{node_id}/{backend}/stats", response_model=BackendStats)
183async def get_backend_stats(
184    node_id: int, backend: str, db: DBDep, admin: SudoAdminDep
185):
186    if not (node := marznode.nodes.get(node_id)):
187        raise HTTPException(status_code=404, detail="Node not found")
188
189    try:
190        stats = await node.get_backend_stats(backend)
191    except Exception:
192        raise HTTPException(502)
193    else:
194        return BackendStats(running=stats.running)
195
196
197@router.get("/{node_id}/{backend}/config", response_model=BackendConfig)
198async def get_node_xray_config(
199    node_id: int, backend: str, admin: SudoAdminDep
200):
201    if not (node := marznode.nodes.get(node_id)):
202        raise HTTPException(status_code=404, detail="Node not found")
203
204    try:
205        config, config_format = await node.get_backend_config(name=backend)
206    except Exception:
207        raise HTTPException(status_code=502, detail="Node isn't responsive")
208    else:
209        return {"config": config, "format": config_format}
210
211
212@router.put("/{node_id}/{backend}/config")
213async def alter_node_xray_config(
214    node_id: int,
215    backend: str,
216    admin: SudoAdminDep,
217    config: Annotated[BackendConfig, Body()],
218):
219    if not (node := marznode.nodes.get(node_id)):
220        raise HTTPException(status_code=404, detail="Node not found")
221
222    try:
223        await asyncio.wait_for(
224            node.restart_backend(
225                name=backend,
226                config=config.config,
227                config_format=config.format.value,
228            ),
229            5,
230        )
231    except:
232        raise HTTPException(
233            status_code=502, detail="No response from the node."
234        )
235    return {}