1import logging
  2from datetime import datetime
  3from typing import List
  4
  5import sqlalchemy
  6from fastapi import APIRouter, Depends
  7from fastapi import HTTPException, WebSocket
  8
  9from app import marznode
 10from app.db import crud, get_tls_certificate
 11from app.dependencies import DBDep, SudoAdminDep, sudo_admin, EndDateDep, StartDateDep
 12from app.marznode import MarzNodeGRPCIO
 13from app.models.admin import Admin
 14from app.models.node import (NodeCreate, NodeModify, NodeResponse,
 15                             NodeSettings, NodeStatus, NodesUsageResponse)
 16
 17logger = logging.getLogger(__name__)
 18router = APIRouter(prefix="/nodes", tags=['Node'], dependencies=[Depends(sudo_admin)])
 19
 20
 21@router.get("/settings", response_model=NodeSettings)
 22def get_node_settings(db: DBDep,
 23                      admin: SudoAdminDep):
 24    tls = crud.get_tls_certificate(db)
 25
 26    return NodeSettings(
 27        certificate=tls.certificate
 28    )
 29
 30
 31@router.post("", response_model=NodeResponse)
 32async def add_node(new_node: NodeCreate,
 33                   db: DBDep,
 34                   admin: SudoAdminDep):
 35    try:
 36        db_node = crud.create_node(db, new_node)
 37    except sqlalchemy.exc.IntegrityError:
 38        db.rollback()
 39        raise HTTPException(status_code=409, detail=f"Node \"{new_node.name}\" already exists")
 40    certificate = get_tls_certificate(db)
 41
 42    await marznode.operations.add_node(db_node, certificate)
 43
 44    logger.info("New node `%s` added", db_node.name)
 45    return db_node
 46
 47
 48@router.get("/{node_id}", response_model=NodeResponse)
 49def get_node(node_id: int,
 50             db: DBDep,
 51             admin: SudoAdminDep):
 52    db_node = crud.get_node_by_id(db, node_id)
 53    if not db_node:
 54        raise HTTPException(status_code=404, detail="Node not found")
 55
 56    return db_node
 57
 58
 59@router.websocket("/{node_id}/logs", dependencies=None)
 60async def node_logs(node_id: int,
 61                    websocket: WebSocket,
 62                    db: DBDep):
 63    token = (
 64        websocket.query_params.get('token')
 65        or
 66        websocket.headers.get('Authorization', '').removeprefix("Bearer ")
 67    )
 68    admin = Admin.get_admin(token, db)
 69    if not admin:
 70        return await websocket.close(reason="Unauthorized", code=4401)
 71
 72    if not admin.is_sudo:
 73        return await websocket.close(reason="You're not allowed", code=4403)
 74
 75    if not marznode.nodes.get(node_id):
 76        return await websocket.close(reason="Node not found", code=4404)
 77
 78    if not await xray.nodes[node_id].is_healthy():
 79        return await websocket.close(reason="Node is not connected", code=4400)
 80
 81    interval = websocket.query_params.get('interval')
 82    if interval:
 83        try:
 84            interval = float(interval)
 85        except ValueError:
 86            return await websocket.close(reason="Invalid interval value", code=4400)
 87        if interval > 10:
 88            return await websocket.close(reason="Interval must be more than 0 and at most 10 seconds", code=4400)
 89
 90    await websocket.accept()
 91    async for l in xray.nodes[node_id].get_logs():
 92        await websocket.send_text(l)
 93
 94
 95@router.get("", response_model=List[NodeResponse])
 96def get_nodes(db: DBDep,
 97              admin: SudoAdminDep):
 98    return crud.get_nodes(db)
 99
100
101@router.put("/{node_id}", response_model=NodeResponse)
102async def modify_node(node_id: int,
103                      modified_node: NodeModify,
104                      db: DBDep,
105                      admin: SudoAdminDep):
106    db_node = crud.get_node_by_id(db, node_id)
107    if not db_node:
108        raise HTTPException(status_code=404, detail="Node not found")
109
110    db_node = crud.update_node(db, db_node, modified_node)
111
112    await marznode.operations.remove_node(db_node.id)
113    if db_node.status != NodeStatus.disabled:
114        certificate = get_tls_certificate(db)
115        await marznode.operations.add_node(db_node, certificate)
116
117    logger.info("Node `%s` modified", db_node.name)
118    return db_node
119
120
121@router.post("/{node_id}/resync")
122async def reconnect_node(node_id: int,
123                         db: DBDep,
124                         admin: SudoAdminDep):
125    db_node = crud.get_node_by_id(db, node_id)
126    if not db_node:
127        raise HTTPException(status_code=404, detail="Node not found")
128
129    return {}
130
131
132@router.delete("/{node_id}")
133async def remove_node(node_id: int,
134                      db: DBDep,
135                      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.get("/usage", response_model=NodesUsageResponse)
148def get_usage(db: DBDep,
149              admin: SudoAdminDep,
150              start_date: StartDateDep,
151              end_date: EndDateDep):
152    """
153    Get nodes usage
154    """
155    usages = crud.get_nodes_usage(db, start_date, end_date)
156
157    return {"usages": usages}