gitscope
https://dev.kronoz.win/r/marzneshin.git
app/routes/node.py · v0.7.2
Source
1 31cdd86 Dawsh 1y ago import asyncio
2 f7af2b3 Dawsh 2y ago import logging
3 d4ec019 Dawsh 2y ago from typing import Annotated
4 a47eb4c Dawsh 2y ago
5 a47eb4c Dawsh 2y ago import sqlalchemy
6 d4ec019 Dawsh 2y ago from fastapi import APIRouter, Body, Query
7 c13a1aa Dawsh 2y ago from fastapi import HTTPException, WebSocket
8 d4ec019 Dawsh 2y ago from fastapi_pagination.ext.sqlalchemy import paginate
9 d4ec019 Dawsh 2y ago from fastapi_pagination.links import Page
10 5488fc0 Dawsh 2y ago from starlette.websockets import WebSocketDisconnect
11 a47eb4c Dawsh 2y ago
12 f7af2b3 Dawsh 2y ago from app import marznode
13 c13a1aa Dawsh 2y ago from app.db import crud, get_tls_certificate
14 d4ec019 Dawsh 2y ago from app.db.models import Node
15 f1ab16c Dawsh 2y ago from app.dependencies import (
16 f1ab16c Dawsh 2y ago DBDep,
17 f1ab16c Dawsh 2y ago SudoAdminDep,
18 f1ab16c Dawsh 2y ago EndDateDep,
19 f1ab16c Dawsh 2y ago StartDateDep,
20 f1ab16c Dawsh 2y ago get_admin,
21 f1ab16c Dawsh 2y ago )
22 01cdf2e Dawsh 2y ago from app.models.node import (
23 01cdf2e Dawsh 2y ago NodeCreate,
24 01cdf2e Dawsh 2y ago NodeModify,
25 01cdf2e Dawsh 2y ago NodeResponse,
26 01cdf2e Dawsh 2y ago NodeSettings,
27 01cdf2e Dawsh 2y ago NodeStatus,
28 f5c168d Mardin 2y ago BackendConfig,
29 f5c168d Mardin 2y ago BackendStats,
30 01cdf2e Dawsh 2y ago )
31 78a3e64 Dawsh 2y ago from app.models.system import TrafficUsageSeries
32 f7af2b3 Dawsh 2y ago
33 f7af2b3 Dawsh 2y ago logger = logging.getLogger(__name__)
34 01cdf2e Dawsh 2y ago router = APIRouter(prefix="/nodes", tags=["Node"])
35 a47eb4c Dawsh 2y ago
36 a47eb4c Dawsh 2y ago
37 d4ec019 Dawsh 2y ago @router.get("", response_model=Page[NodeResponse])
38 01cdf2e Dawsh 2y ago def get_nodes(
39 01cdf2e Dawsh 2y ago db: DBDep,
40 01cdf2e Dawsh 2y ago admin: SudoAdminDep,
41 01cdf2e Dawsh 2y ago status: list[NodeStatus] = Query(None),
42 01cdf2e Dawsh 2y ago name: str = Query(None),
43 01cdf2e Dawsh 2y ago ):
44 d4ec019 Dawsh 2y ago query = db.query(Node)
45 d4ec019 Dawsh 2y ago
46 d4ec019 Dawsh 2y ago if name:
47 d4ec019 Dawsh 2y ago query = query.filter(Node.name.ilike(f"%{name}%"))
48 d4ec019 Dawsh 2y ago
49 d4ec019 Dawsh 2y ago if status:
50 d4ec019 Dawsh 2y ago query = query.filter(Node.status.in_(status))
51 d4ec019 Dawsh 2y ago
52 d4ec019 Dawsh 2y ago return paginate(db, query)
53 d4ec019 Dawsh 2y ago
54 d4ec019 Dawsh 2y ago
55 d4ec019 Dawsh 2y ago @router.post("", response_model=NodeResponse)
56 01cdf2e Dawsh 2y ago async def add_node(new_node: NodeCreate, db: DBDep, admin: SudoAdminDep):
57 d4ec019 Dawsh 2y ago try:
58 d4ec019 Dawsh 2y ago db_node = crud.create_node(db, new_node)
59 d4ec019 Dawsh 2y ago except sqlalchemy.exc.IntegrityError:
60 d4ec019 Dawsh 2y ago db.rollback()
61 01cdf2e Dawsh 2y ago raise HTTPException(
62 01cdf2e Dawsh 2y ago status_code=409, detail=f'Node "{new_node.name}" already exists'
63 01cdf2e Dawsh 2y ago )
64 d4ec019 Dawsh 2y ago certificate = get_tls_certificate(db)
65 d4ec019 Dawsh 2y ago
66 d4ec019 Dawsh 2y ago await marznode.operations.add_node(db_node, certificate)
67 d4ec019 Dawsh 2y ago
68 d4ec019 Dawsh 2y ago logger.info("New node `%s` added", db_node.name)
69 d4ec019 Dawsh 2y ago return db_node
70 d4ec019 Dawsh 2y ago
71 d4ec019 Dawsh 2y ago
72 2de9d72 Dawsh 2y ago @router.get("/settings", response_model=NodeSettings)
73 01cdf2e Dawsh 2y ago def get_node_settings(db: DBDep, admin: SudoAdminDep):
74 a47eb4c Dawsh 2y ago tls = crud.get_tls_certificate(db)
75 a47eb4c Dawsh 2y ago
76 01cdf2e Dawsh 2y ago return NodeSettings(certificate=tls.certificate)
77 a47eb4c Dawsh 2y ago
78 a47eb4c Dawsh 2y ago
79 2de9d72 Dawsh 2y ago @router.get("/{node_id}", response_model=NodeResponse)
80 01cdf2e Dawsh 2y ago def get_node(node_id: int, db: DBDep, admin: SudoAdminDep):
81 c13a1aa Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
82 c13a1aa Dawsh 2y ago if not db_node:
83 a47eb4c Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
84 a47eb4c Dawsh 2y ago
85 c13a1aa Dawsh 2y ago return db_node
86 a47eb4c Dawsh 2y ago
87 a47eb4c Dawsh 2y ago
88 f5c168d Mardin 2y ago @router.websocket("/{node_id}/{backend}/logs")
89 01cdf2e Dawsh 2y ago async def node_logs(
90 f5c168d Mardin 2y ago node_id: int,
91 f5c168d Mardin 2y ago backend: str,
92 f5c168d Mardin 2y ago websocket: WebSocket,
93 f5c168d Mardin 2y ago db: DBDep,
94 f5c168d Mardin 2y ago include_buffer: bool = True,
95 01cdf2e Dawsh 2y ago ):
96 31cdd86 Dawsh 1y ago token = websocket.query_params.get("token", "") or websocket.headers.get(
97 01cdf2e Dawsh 2y ago "Authorization", ""
98 01cdf2e Dawsh 2y ago ).removeprefix("Bearer ")
99 512d183 Mrdn 2y ago admin = get_admin(db, token)
100 a47eb4c Dawsh 2y ago
101 512d183 Mrdn 2y ago if not admin or not admin.is_sudo:
102 a47eb4c Dawsh 2y ago return await websocket.close(reason="You're not allowed", code=4403)
103 a47eb4c Dawsh 2y ago
104 12b255b Dawsh 2y ago if not marznode.nodes.get(node_id):
105 a47eb4c Dawsh 2y ago return await websocket.close(reason="Node not found", code=4404)
106 a47eb4c Dawsh 2y ago
107 a47eb4c Dawsh 2y ago await websocket.accept()
108 5488fc0 Dawsh 2y ago try:
109 5488fc0 Dawsh 2y ago async for line in marznode.nodes[node_id].get_logs(
110 f5c168d Mardin 2y ago name=backend, include_buffer=include_buffer
111 5488fc0 Dawsh 2y ago ):
112 5488fc0 Dawsh 2y ago try:
113 5488fc0 Dawsh 2y ago await websocket.send_text(line)
114 5488fc0 Dawsh 2y ago except WebSocketDisconnect:
115 5488fc0 Dawsh 2y ago break
116 31cdd86 Dawsh 1y ago finally:
117 5488fc0 Dawsh 2y ago await websocket.close()
118 a47eb4c Dawsh 2y ago
119 a47eb4c Dawsh 2y ago
120 2de9d72 Dawsh 2y ago @router.put("/{node_id}", response_model=NodeResponse)
121 01cdf2e Dawsh 2y ago async def modify_node(
122 01cdf2e Dawsh 2y ago node_id: int, modified_node: NodeModify, db: DBDep, admin: SudoAdminDep
123 01cdf2e Dawsh 2y ago ):
124 12b255b Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
125 12b255b Dawsh 2y ago if not db_node:
126 a47eb4c Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
127 a47eb4c Dawsh 2y ago
128 12b255b Dawsh 2y ago db_node = crud.update_node(db, db_node, modified_node)
129 a47eb4c Dawsh 2y ago
130 12b255b Dawsh 2y ago await marznode.operations.remove_node(db_node.id)
131 12b255b Dawsh 2y ago if db_node.status != NodeStatus.disabled:
132 12b255b Dawsh 2y ago certificate = get_tls_certificate(db)
133 f6a32ea Dawsh 2y ago await marznode.operations.add_node(db_node, certificate)
134 a47eb4c Dawsh 2y ago
135 f7af2b3 Dawsh 2y ago logger.info("Node `%s` modified", db_node.name)
136 12b255b Dawsh 2y ago return db_node
137 a47eb4c Dawsh 2y ago
138 a47eb4c Dawsh 2y ago
139 d4ec019 Dawsh 2y ago @router.delete("/{node_id}")
140 01cdf2e Dawsh 2y ago async def remove_node(node_id: int, db: DBDep, admin: SudoAdminDep):
141 d4ec019 Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
142 d4ec019 Dawsh 2y ago if not db_node:
143 d4ec019 Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
144 d4ec019 Dawsh 2y ago
145 d4ec019 Dawsh 2y ago crud.remove_node(db, db_node)
146 d4ec019 Dawsh 2y ago await marznode.operations.remove_node(db_node.id)
147 d4ec019 Dawsh 2y ago
148 d4ec019 Dawsh 2y ago logger.info(f"Node `%s` deleted", db_node.name)
149 d4ec019 Dawsh 2y ago return {}
150 d4ec019 Dawsh 2y ago
151 d4ec019 Dawsh 2y ago
152 2de9d72 Dawsh 2y ago @router.post("/{node_id}/resync")
153 01cdf2e Dawsh 2y ago async def reconnect_node(node_id: int, db: DBDep, admin: SudoAdminDep):
154 12b255b Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
155 12b255b Dawsh 2y ago if not db_node:
156 a47eb4c Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
157 a47eb4c Dawsh 2y ago
158 a47eb4c Dawsh 2y ago return {}
159 a47eb4c Dawsh 2y ago
160 a47eb4c Dawsh 2y ago
161 78a3e64 Dawsh 2y ago @router.get("/{node_id}/usage", response_model=TrafficUsageSeries)
162 78a3e64 Dawsh 2y ago def get_usage(
163 78a3e64 Dawsh 2y ago node_id: int,
164 78a3e64 Dawsh 2y ago db: DBDep,
165 78a3e64 Dawsh 2y ago admin: SudoAdminDep,
166 78a3e64 Dawsh 2y ago start_date: StartDateDep,
167 78a3e64 Dawsh 2y ago end_date: EndDateDep,
168 78a3e64 Dawsh 2y ago ):
169 78a3e64 Dawsh 2y ago """
170 78a3e64 Dawsh 2y ago Get nodes usage
171 78a3e64 Dawsh 2y ago """
172 78a3e64 Dawsh 2y ago node = crud.get_node_by_id(db, node_id)
173 78a3e64 Dawsh 2y ago if not node:
174 78a3e64 Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
175 78a3e64 Dawsh 2y ago
176 78a3e64 Dawsh 2y ago return crud.get_node_usage(db, start_date, end_date, node)
177 78a3e64 Dawsh 2y ago
178 78a3e64 Dawsh 2y ago
179 f5c168d Mardin 2y ago @router.get("/{node_id}/{backend}/stats", response_model=BackendStats)
180 f5c168d Mardin 2y ago async def get_backend_stats(
181 f5c168d Mardin 2y ago node_id: int, backend: str, db: DBDep, admin: SudoAdminDep
182 f5c168d Mardin 2y ago ):
183 512d183 Mrdn 2y ago if not (node := marznode.nodes.get(node_id)):
184 512d183 Mrdn 2y ago raise HTTPException(status_code=404, detail="Node not found")
185 5488fc0 Dawsh 2y ago
186 5488fc0 Dawsh 2y ago try:
187 f5c168d Mardin 2y ago stats = await node.get_backend_stats(backend)
188 f5c168d Mardin 2y ago except Exception:
189 f5c168d Mardin 2y ago raise HTTPException(502)
190 f5c168d Mardin 2y ago else:
191 f5c168d Mardin 2y ago return BackendStats(running=stats.running)
192 f5c168d Mardin 2y ago
193 f5c168d Mardin 2y ago
194 f5c168d Mardin 2y ago @router.get("/{node_id}/{backend}/config", response_model=BackendConfig)
195 f5c168d Mardin 2y ago async def get_node_xray_config(
196 31cdd86 Dawsh 1y ago node_id: int, backend: str, admin: SudoAdminDep
197 f5c168d Mardin 2y ago ):
198 f5c168d Mardin 2y ago if not (node := marznode.nodes.get(node_id)):
199 f5c168d Mardin 2y ago raise HTTPException(status_code=404, detail="Node not found")
200 f5c168d Mardin 2y ago
201 f5c168d Mardin 2y ago try:
202 f5c168d Mardin 2y ago config, config_format = await node.get_backend_config(name=backend)
203 f5c168d Mardin 2y ago except Exception:
204 5488fc0 Dawsh 2y ago raise HTTPException(status_code=502, detail="Node isn't responsive")
205 5488fc0 Dawsh 2y ago else:
206 f5c168d Mardin 2y ago return {"config": config, "format": config_format}
207 512d183 Mrdn 2y ago
208 512d183 Mrdn 2y ago
209 f5c168d Mardin 2y ago @router.put("/{node_id}/{backend}/config")
210 01cdf2e Dawsh 2y ago async def alter_node_xray_config(
211 f5c168d Mardin 2y ago node_id: int,
212 f5c168d Mardin 2y ago backend: str,
213 f5c168d Mardin 2y ago admin: SudoAdminDep,
214 f5c168d Mardin 2y ago config: Annotated[BackendConfig, Body()],
215 01cdf2e Dawsh 2y ago ):
216 512d183 Mrdn 2y ago if not (node := marznode.nodes.get(node_id)):
217 512d183 Mrdn 2y ago raise HTTPException(status_code=404, detail="Node not found")
218 31cdd86 Dawsh 1y ago
219 5488fc0 Dawsh 2y ago try:
220 31cdd86 Dawsh 1y ago await asyncio.wait_for(
221 31cdd86 Dawsh 1y ago node.restart_backend(
222 31cdd86 Dawsh 1y ago name=backend,
223 31cdd86 Dawsh 1y ago config=config.config,
224 31cdd86 Dawsh 1y ago config_format=config.format.value,
225 31cdd86 Dawsh 1y ago ),
226 31cdd86 Dawsh 1y ago 5,
227 f5c168d Mardin 2y ago )
228 5488fc0 Dawsh 2y ago except:
229 31cdd86 Dawsh 1y ago raise HTTPException(
230 31cdd86 Dawsh 1y ago status_code=502, detail="No response from the node."
231 31cdd86 Dawsh 1y ago )
232 512d183 Mrdn 2y ago return {}