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