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