gitscope
https://dev.kronoz.win/r/marzneshin.git
app/routes/node.py · v0.5.6
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 01cdf2e Dawsh 2y ago NodesUsageResponse,
28 f5c168d Mardin 2y ago BackendConfig,
29 f5c168d Mardin 2y ago BackendStats,
30 01cdf2e Dawsh 2y ago )
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 512d183 Mrdn 2y ago @router.get("/usage", response_model=NodesUsageResponse)
79 01cdf2e Dawsh 2y ago def get_usage(
80 f1ab16c Dawsh 2y ago db: DBDep,
81 f1ab16c Dawsh 2y ago admin: SudoAdminDep,
82 f1ab16c Dawsh 2y ago start_date: StartDateDep,
83 f1ab16c Dawsh 2y ago end_date: EndDateDep,
84 01cdf2e Dawsh 2y ago ):
85 512d183 Mrdn 2y ago """
86 512d183 Mrdn 2y ago Get nodes usage
87 512d183 Mrdn 2y ago """
88 512d183 Mrdn 2y ago usages = crud.get_nodes_usage(db, start_date, end_date)
89 512d183 Mrdn 2y ago
90 512d183 Mrdn 2y ago return {"usages": usages}
91 512d183 Mrdn 2y ago
92 512d183 Mrdn 2y ago
93 2de9d72 Dawsh 2y ago @router.get("/{node_id}", response_model=NodeResponse)
94 01cdf2e Dawsh 2y ago def get_node(node_id: int, db: DBDep, admin: SudoAdminDep):
95 c13a1aa Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
96 c13a1aa Dawsh 2y ago if not db_node:
97 a47eb4c Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
98 a47eb4c Dawsh 2y ago
99 c13a1aa Dawsh 2y ago return db_node
100 a47eb4c Dawsh 2y ago
101 a47eb4c Dawsh 2y ago
102 f5c168d Mardin 2y ago @router.websocket("/{node_id}/{backend}/logs")
103 01cdf2e Dawsh 2y ago async def node_logs(
104 f5c168d Mardin 2y ago node_id: int,
105 f5c168d Mardin 2y ago backend: str,
106 f5c168d Mardin 2y ago websocket: WebSocket,
107 f5c168d Mardin 2y ago db: DBDep,
108 f5c168d Mardin 2y ago include_buffer: bool = True,
109 01cdf2e Dawsh 2y ago ):
110 01cdf2e Dawsh 2y ago token = websocket.query_params.get("token") or websocket.headers.get(
111 01cdf2e Dawsh 2y ago "Authorization", ""
112 01cdf2e Dawsh 2y ago ).removeprefix("Bearer ")
113 512d183 Mrdn 2y ago admin = get_admin(db, token)
114 a47eb4c Dawsh 2y ago
115 512d183 Mrdn 2y ago if not admin or not admin.is_sudo:
116 a47eb4c Dawsh 2y ago return await websocket.close(reason="You're not allowed", code=4403)
117 a47eb4c Dawsh 2y ago
118 12b255b Dawsh 2y ago if not marznode.nodes.get(node_id):
119 a47eb4c Dawsh 2y ago return await websocket.close(reason="Node not found", code=4404)
120 a47eb4c Dawsh 2y ago
121 a47eb4c Dawsh 2y ago await websocket.accept()
122 5488fc0 Dawsh 2y ago try:
123 5488fc0 Dawsh 2y ago async for line in marznode.nodes[node_id].get_logs(
124 f5c168d Mardin 2y ago name=backend, include_buffer=include_buffer
125 5488fc0 Dawsh 2y ago ):
126 5488fc0 Dawsh 2y ago try:
127 5488fc0 Dawsh 2y ago await websocket.send_text(line)
128 5488fc0 Dawsh 2y ago except WebSocketDisconnect:
129 5488fc0 Dawsh 2y ago break
130 5488fc0 Dawsh 2y ago except:
131 5488fc0 Dawsh 2y ago await websocket.close()
132 a47eb4c Dawsh 2y ago
133 a47eb4c Dawsh 2y ago
134 2de9d72 Dawsh 2y ago @router.put("/{node_id}", response_model=NodeResponse)
135 01cdf2e Dawsh 2y ago async def modify_node(
136 01cdf2e Dawsh 2y ago node_id: int, modified_node: NodeModify, db: DBDep, admin: SudoAdminDep
137 01cdf2e Dawsh 2y ago ):
138 12b255b Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
139 12b255b Dawsh 2y ago if not db_node:
140 a47eb4c Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
141 a47eb4c Dawsh 2y ago
142 12b255b Dawsh 2y ago db_node = crud.update_node(db, db_node, modified_node)
143 a47eb4c Dawsh 2y ago
144 12b255b Dawsh 2y ago await marznode.operations.remove_node(db_node.id)
145 12b255b Dawsh 2y ago if db_node.status != NodeStatus.disabled:
146 12b255b Dawsh 2y ago certificate = get_tls_certificate(db)
147 f6a32ea Dawsh 2y ago await marznode.operations.add_node(db_node, certificate)
148 a47eb4c Dawsh 2y ago
149 f7af2b3 Dawsh 2y ago logger.info("Node `%s` modified", db_node.name)
150 12b255b Dawsh 2y ago return db_node
151 a47eb4c Dawsh 2y ago
152 a47eb4c Dawsh 2y ago
153 d4ec019 Dawsh 2y ago @router.delete("/{node_id}")
154 01cdf2e Dawsh 2y ago async def remove_node(node_id: int, db: DBDep, admin: SudoAdminDep):
155 d4ec019 Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
156 d4ec019 Dawsh 2y ago if not db_node:
157 d4ec019 Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
158 d4ec019 Dawsh 2y ago
159 d4ec019 Dawsh 2y ago crud.remove_node(db, db_node)
160 d4ec019 Dawsh 2y ago await marznode.operations.remove_node(db_node.id)
161 d4ec019 Dawsh 2y ago
162 d4ec019 Dawsh 2y ago logger.info(f"Node `%s` deleted", db_node.name)
163 d4ec019 Dawsh 2y ago return {}
164 d4ec019 Dawsh 2y ago
165 d4ec019 Dawsh 2y ago
166 2de9d72 Dawsh 2y ago @router.post("/{node_id}/resync")
167 01cdf2e Dawsh 2y ago async def reconnect_node(node_id: int, db: DBDep, admin: SudoAdminDep):
168 12b255b Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
169 12b255b Dawsh 2y ago if not db_node:
170 a47eb4c Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
171 a47eb4c Dawsh 2y ago
172 a47eb4c Dawsh 2y ago return {}
173 a47eb4c Dawsh 2y ago
174 a47eb4c Dawsh 2y ago
175 f5c168d Mardin 2y ago @router.get("/{node_id}/{backend}/stats", response_model=BackendStats)
176 f5c168d Mardin 2y ago async def get_backend_stats(
177 f5c168d Mardin 2y ago node_id: int, backend: str, db: DBDep, admin: SudoAdminDep
178 f5c168d Mardin 2y ago ):
179 512d183 Mrdn 2y ago if not (node := marznode.nodes.get(node_id)):
180 512d183 Mrdn 2y ago raise HTTPException(status_code=404, detail="Node not found")
181 5488fc0 Dawsh 2y ago
182 5488fc0 Dawsh 2y ago try:
183 f5c168d Mardin 2y ago stats = await node.get_backend_stats(backend)
184 f5c168d Mardin 2y ago except Exception:
185 f5c168d Mardin 2y ago raise HTTPException(502)
186 f5c168d Mardin 2y ago else:
187 f5c168d Mardin 2y ago return BackendStats(running=stats.running)
188 f5c168d Mardin 2y ago
189 f5c168d Mardin 2y ago
190 f5c168d Mardin 2y ago @router.get("/{node_id}/{backend}/config", response_model=BackendConfig)
191 f5c168d Mardin 2y ago async def get_node_xray_config(
192 f5c168d Mardin 2y ago node_id: int, backend: str, db: DBDep, admin: SudoAdminDep
193 f5c168d Mardin 2y ago ):
194 f5c168d Mardin 2y ago if not (node := marznode.nodes.get(node_id)):
195 f5c168d Mardin 2y ago raise HTTPException(status_code=404, detail="Node not found")
196 f5c168d Mardin 2y ago
197 f5c168d Mardin 2y ago try:
198 f5c168d Mardin 2y ago config, config_format = await node.get_backend_config(name=backend)
199 f5c168d Mardin 2y ago except Exception:
200 5488fc0 Dawsh 2y ago raise HTTPException(status_code=502, detail="Node isn't responsive")
201 5488fc0 Dawsh 2y ago else:
202 f5c168d Mardin 2y ago return {"config": config, "format": config_format}
203 512d183 Mrdn 2y ago
204 512d183 Mrdn 2y ago
205 f5c168d Mardin 2y ago @router.put("/{node_id}/{backend}/config")
206 01cdf2e Dawsh 2y ago async def alter_node_xray_config(
207 f5c168d Mardin 2y ago node_id: int,
208 f5c168d Mardin 2y ago backend: str,
209 f5c168d Mardin 2y ago db: DBDep,
210 f5c168d Mardin 2y ago admin: SudoAdminDep,
211 f5c168d Mardin 2y ago config: Annotated[BackendConfig, Body()],
212 01cdf2e Dawsh 2y ago ):
213 512d183 Mrdn 2y ago if not (node := marznode.nodes.get(node_id)):
214 512d183 Mrdn 2y ago raise HTTPException(status_code=404, detail="Node not found")
215 5488fc0 Dawsh 2y ago try:
216 f5c168d Mardin 2y ago await node.restart_backend(
217 f5c168d Mardin 2y ago name=backend,
218 f5c168d Mardin 2y ago config=config.config,
219 f5c168d Mardin 2y ago config_format=config.format.value,
220 f5c168d Mardin 2y ago )
221 5488fc0 Dawsh 2y ago except:
222 5488fc0 Dawsh 2y ago raise HTTPException(status_code=502, detail="Node isn't responsive")
223 512d183 Mrdn 2y ago return {}