1import json
  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
 10
 11from app import marznode
 12from app.db import crud, get_tls_certificate
 13from app.db.models import Node
 14from app.dependencies import DBDep, SudoAdminDep, EndDateDep, StartDateDep, get_admin
 15
 16from app.models.node import (
 17    NodeCreate,
 18    NodeModify,
 19    NodeResponse,
 20    NodeSettings,
 21    NodeStatus,
 22    NodesUsageResponse,
 23)
 24
 25logger = logging.getLogger(__name__)
 26router = APIRouter(prefix="/nodes", tags=["Node"])
 27
 28
 29@router.get("", response_model=Page[NodeResponse])
 30def get_nodes(
 31    db: DBDep,
 32    admin: SudoAdminDep,
 33    status: list[NodeStatus] = Query(None),
 34    name: str = Query(None),
 35):
 36    query = db.query(Node)
 37
 38    if name:
 39        query = query.filter(Node.name.ilike(f"%{name}%"))
 40
 41    if status:
 42        query = query.filter(Node.status.in_(status))
 43
 44    return paginate(db, query)
 45
 46
 47@router.post("", response_model=NodeResponse)
 48async def add_node(new_node: NodeCreate, db: DBDep, admin: SudoAdminDep):
 49    try:
 50        db_node = crud.create_node(db, new_node)
 51    except sqlalchemy.exc.IntegrityError:
 52        db.rollback()
 53        raise HTTPException(
 54            status_code=409, detail=f'Node "{new_node.name}" already exists'
 55        )
 56    certificate = get_tls_certificate(db)
 57
 58    await marznode.operations.add_node(db_node, certificate)
 59
 60    logger.info("New node `%s` added", db_node.name)
 61    return db_node
 62
 63
 64@router.get("/settings", response_model=NodeSettings)
 65def get_node_settings(db: DBDep, admin: SudoAdminDep):
 66    tls = crud.get_tls_certificate(db)
 67
 68    return NodeSettings(certificate=tls.certificate)
 69
 70
 71@router.get("/usage", response_model=NodesUsageResponse)
 72def get_usage(
 73    db: DBDep, admin: SudoAdminDep, start_date: StartDateDep, end_date: EndDateDep
 74):
 75    """
 76    Get nodes usage
 77    """
 78    usages = crud.get_nodes_usage(db, start_date, end_date)
 79
 80    return {"usages": usages}
 81
 82
 83@router.get("/{node_id}", response_model=NodeResponse)
 84def get_node(node_id: int, db: DBDep, admin: SudoAdminDep):
 85    db_node = crud.get_node_by_id(db, node_id)
 86    if not db_node:
 87        raise HTTPException(status_code=404, detail="Node not found")
 88
 89    return db_node
 90
 91
 92@router.websocket("/{node_id}/logs")
 93async def node_logs(
 94    node_id: int, websocket: WebSocket, db: DBDep, 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    # if not await marznode.nodes[node_id].is_healthy:
108    #    return await websocket.close(reason="Node is not connected", code=4400)
109
110    await websocket.accept()
111    async for line in marznode.nodes[node_id].get_logs(include_buffer=include_buffer):
112        await websocket.send_text(line)
113
114
115@router.put("/{node_id}", response_model=NodeResponse)
116async def modify_node(
117    node_id: int, modified_node: NodeModify, db: DBDep, admin: SudoAdminDep
118):
119    db_node = crud.get_node_by_id(db, node_id)
120    if not db_node:
121        raise HTTPException(status_code=404, detail="Node not found")
122
123    db_node = crud.update_node(db, db_node, modified_node)
124
125    await marznode.operations.remove_node(db_node.id)
126    if db_node.status != NodeStatus.disabled:
127        certificate = get_tls_certificate(db)
128        await marznode.operations.add_node(db_node, certificate)
129
130    logger.info("Node `%s` modified", db_node.name)
131    return db_node
132
133
134@router.delete("/{node_id}")
135async def remove_node(node_id: int, db: DBDep, admin: SudoAdminDep):
136    db_node = crud.get_node_by_id(db, node_id)
137    if not db_node:
138        raise HTTPException(status_code=404, detail="Node not found")
139
140    crud.remove_node(db, db_node)
141    await marznode.operations.remove_node(db_node.id)
142
143    logger.info(f"Node `%s` deleted", db_node.name)
144    return {}
145
146
147@router.post("/{node_id}/resync")
148async def reconnect_node(node_id: int, db: DBDep, admin: SudoAdminDep):
149    db_node = crud.get_node_by_id(db, node_id)
150    if not db_node:
151        raise HTTPException(status_code=404, detail="Node not found")
152
153    return {}
154
155
156@router.get("/{node_id}/xray_config")
157async def get_node_xray_config(node_id: int, db: DBDep, admin: SudoAdminDep):
158    if not (node := marznode.nodes.get(node_id)):
159        raise HTTPException(status_code=404, detail="Node not found")
160    return json.loads(await node.get_xray_config())
161
162
163@router.put("/{node_id}/xray_config")
164async def alter_node_xray_config(
165    node_id: int, db: DBDep, admin: SudoAdminDep, body: Annotated[dict, Body()]
166):
167    if not (node := marznode.nodes.get(node_id)):
168        raise HTTPException(status_code=404, detail="Node not found")
169    xray_config = json.dumps(body)
170    await node.restart_xray(xray_config)
171    return {}