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