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
10from starlette.websockets import WebSocketDisconnect
11
12from app import marznode
13from app.db import crud, get_tls_certificate
14from app.db.models import Node
15from app.dependencies import DBDep, SudoAdminDep, EndDateDep, StartDateDep, get_admin
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 try:
112 async for line in marznode.nodes[node_id].get_logs(
113 include_buffer=include_buffer
114 ):
115 try:
116 await websocket.send_text(line)
117 except WebSocketDisconnect:
118 break
119 except:
120 await websocket.close()
121
122
123@router.put("/{node_id}", response_model=NodeResponse)
124async def modify_node(
125 node_id: int, modified_node: NodeModify, db: DBDep, admin: SudoAdminDep
126):
127 db_node = crud.get_node_by_id(db, node_id)
128 if not db_node:
129 raise HTTPException(status_code=404, detail="Node not found")
130
131 db_node = crud.update_node(db, db_node, modified_node)
132
133 await marznode.operations.remove_node(db_node.id)
134 if db_node.status != NodeStatus.disabled:
135 certificate = get_tls_certificate(db)
136 await marznode.operations.add_node(db_node, certificate)
137
138 logger.info("Node `%s` modified", db_node.name)
139 return db_node
140
141
142@router.delete("/{node_id}")
143async def remove_node(node_id: int, db: DBDep, admin: SudoAdminDep):
144 db_node = crud.get_node_by_id(db, node_id)
145 if not db_node:
146 raise HTTPException(status_code=404, detail="Node not found")
147
148 crud.remove_node(db, db_node)
149 await marznode.operations.remove_node(db_node.id)
150
151 logger.info(f"Node `%s` deleted", db_node.name)
152 return {}
153
154
155@router.post("/{node_id}/resync")
156async def reconnect_node(node_id: int, db: DBDep, admin: SudoAdminDep):
157 db_node = crud.get_node_by_id(db, node_id)
158 if not db_node:
159 raise HTTPException(status_code=404, detail="Node not found")
160
161 return {}
162
163
164@router.get("/{node_id}/xray_config")
165async def get_node_xray_config(node_id: int, db: DBDep, admin: SudoAdminDep):
166 if not (node := marznode.nodes.get(node_id)):
167 raise HTTPException(status_code=404, detail="Node not found")
168
169 try:
170 config = await node.get_xray_config()
171 except:
172 raise HTTPException(status_code=502, detail="Node isn't responsive")
173 else:
174 return json.loads(config)
175
176
177@router.put("/{node_id}/xray_config")
178async def alter_node_xray_config(
179 node_id: int, db: DBDep, admin: SudoAdminDep, body: Annotated[dict, Body()]
180):
181 if not (node := marznode.nodes.get(node_id)):
182 raise HTTPException(status_code=404, detail="Node not found")
183 xray_config = json.dumps(body)
184 try:
185 await node.restart_xray(xray_config)
186 except:
187 raise HTTPException(status_code=502, detail="Node isn't responsive")
188 return {}