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