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    NodesUsageResponse,
 28    BackendConfig,
 29    BackendStats,
 30)
 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("/usage", response_model=NodesUsageResponse)
 79def get_usage(
 80    db: DBDep,
 81    admin: SudoAdminDep,
 82    start_date: StartDateDep,
 83    end_date: EndDateDep,
 84):
 85    """
 86    Get nodes usage
 87    """
 88    usages = crud.get_nodes_usage(db, start_date, end_date)
 89
 90    return {"usages": usages}
 91
 92
 93@router.get("/{node_id}", response_model=NodeResponse)
 94def get_node(node_id: int, db: DBDep, admin: SudoAdminDep):
 95    db_node = crud.get_node_by_id(db, node_id)
 96    if not db_node:
 97        raise HTTPException(status_code=404, detail="Node not found")
 98
 99    return db_node
100
101
102@router.websocket("/{node_id}/{backend}/logs")
103async def node_logs(
104    node_id: int,
105    backend: str,
106    websocket: WebSocket,
107    db: DBDep,
108    include_buffer: bool = True,
109):
110    token = websocket.query_params.get("token") or websocket.headers.get(
111        "Authorization", ""
112    ).removeprefix("Bearer ")
113    admin = get_admin(db, token)
114
115    if not admin or not admin.is_sudo:
116        return await websocket.close(reason="You're not allowed", code=4403)
117
118    if not marznode.nodes.get(node_id):
119        return await websocket.close(reason="Node not found", code=4404)
120
121    await websocket.accept()
122    try:
123        async for line in marznode.nodes[node_id].get_logs(
124            name=backend, include_buffer=include_buffer
125        ):
126            try:
127                await websocket.send_text(line)
128            except WebSocketDisconnect:
129                break
130    except:
131        await websocket.close()
132
133
134@router.put("/{node_id}", response_model=NodeResponse)
135async def modify_node(
136    node_id: int, modified_node: NodeModify, db: DBDep, admin: SudoAdminDep
137):
138    db_node = crud.get_node_by_id(db, node_id)
139    if not db_node:
140        raise HTTPException(status_code=404, detail="Node not found")
141
142    db_node = crud.update_node(db, db_node, modified_node)
143
144    await marznode.operations.remove_node(db_node.id)
145    if db_node.status != NodeStatus.disabled:
146        certificate = get_tls_certificate(db)
147        await marznode.operations.add_node(db_node, certificate)
148
149    logger.info("Node `%s` modified", db_node.name)
150    return db_node
151
152
153@router.delete("/{node_id}")
154async def remove_node(node_id: int, db: DBDep, admin: SudoAdminDep):
155    db_node = crud.get_node_by_id(db, node_id)
156    if not db_node:
157        raise HTTPException(status_code=404, detail="Node not found")
158
159    crud.remove_node(db, db_node)
160    await marznode.operations.remove_node(db_node.id)
161
162    logger.info(f"Node `%s` deleted", db_node.name)
163    return {}
164
165
166@router.post("/{node_id}/resync")
167async def reconnect_node(node_id: int, db: DBDep, admin: SudoAdminDep):
168    db_node = crud.get_node_by_id(db, node_id)
169    if not db_node:
170        raise HTTPException(status_code=404, detail="Node not found")
171
172    return {}
173
174
175@router.get("/{node_id}/{backend}/stats", response_model=BackendStats)
176async def get_backend_stats(
177    node_id: int, backend: str, db: DBDep, admin: SudoAdminDep
178):
179    if not (node := marznode.nodes.get(node_id)):
180        raise HTTPException(status_code=404, detail="Node not found")
181
182    try:
183        stats = await node.get_backend_stats(backend)
184    except Exception:
185        raise HTTPException(502)
186    else:
187        return BackendStats(running=stats.running)
188
189
190@router.get("/{node_id}/{backend}/config", response_model=BackendConfig)
191async def get_node_xray_config(
192    node_id: int, backend: str, db: DBDep, admin: SudoAdminDep
193):
194    if not (node := marznode.nodes.get(node_id)):
195        raise HTTPException(status_code=404, detail="Node not found")
196
197    try:
198        config, config_format = await node.get_backend_config(name=backend)
199    except Exception:
200        raise HTTPException(status_code=502, detail="Node isn't responsive")
201    else:
202        return {"config": config, "format": config_format}
203
204
205@router.put("/{node_id}/{backend}/config")
206async def alter_node_xray_config(
207    node_id: int,
208    backend: str,
209    db: DBDep,
210    admin: SudoAdminDep,
211    config: Annotated[BackendConfig, Body()],
212):
213    if not (node := marznode.nodes.get(node_id)):
214        raise HTTPException(status_code=404, detail="Node not found")
215    try:
216        await node.restart_backend(
217            name=backend,
218            config=config.config,
219            config_format=config.format.value,
220        )
221    except:
222        raise HTTPException(status_code=502, detail="Node isn't responsive")
223    return {}