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