gitscope
https://dev.kronoz.win/r/marzneshin.git
app/routes/node.py · v0.3.1
Source
1 512d183 Mrdn 2y ago import json
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 01cdf2e Dawsh 2y ago NodesUsageResponse,
29 01cdf2e Dawsh 2y ago )
30 f7af2b3 Dawsh 2y ago
31 f7af2b3 Dawsh 2y ago logger = logging.getLogger(__name__)
32 01cdf2e Dawsh 2y ago router = APIRouter(prefix="/nodes", tags=["Node"])
33 a47eb4c Dawsh 2y ago
34 a47eb4c Dawsh 2y ago
35 d4ec019 Dawsh 2y ago @router.get("", response_model=Page[NodeResponse])
36 01cdf2e Dawsh 2y ago def get_nodes(
37 01cdf2e Dawsh 2y ago db: DBDep,
38 01cdf2e Dawsh 2y ago admin: SudoAdminDep,
39 01cdf2e Dawsh 2y ago status: list[NodeStatus] = Query(None),
40 01cdf2e Dawsh 2y ago name: str = Query(None),
41 01cdf2e Dawsh 2y ago ):
42 d4ec019 Dawsh 2y ago query = db.query(Node)
43 d4ec019 Dawsh 2y ago
44 d4ec019 Dawsh 2y ago if name:
45 d4ec019 Dawsh 2y ago query = query.filter(Node.name.ilike(f"%{name}%"))
46 d4ec019 Dawsh 2y ago
47 d4ec019 Dawsh 2y ago if status:
48 d4ec019 Dawsh 2y ago query = query.filter(Node.status.in_(status))
49 d4ec019 Dawsh 2y ago
50 d4ec019 Dawsh 2y ago return paginate(db, query)
51 d4ec019 Dawsh 2y ago
52 d4ec019 Dawsh 2y ago
53 d4ec019 Dawsh 2y ago @router.post("", response_model=NodeResponse)
54 01cdf2e Dawsh 2y ago async def add_node(new_node: NodeCreate, db: DBDep, admin: SudoAdminDep):
55 d4ec019 Dawsh 2y ago try:
56 d4ec019 Dawsh 2y ago db_node = crud.create_node(db, new_node)
57 d4ec019 Dawsh 2y ago except sqlalchemy.exc.IntegrityError:
58 d4ec019 Dawsh 2y ago db.rollback()
59 01cdf2e Dawsh 2y ago raise HTTPException(
60 01cdf2e Dawsh 2y ago status_code=409, detail=f'Node "{new_node.name}" already exists'
61 01cdf2e Dawsh 2y ago )
62 d4ec019 Dawsh 2y ago certificate = get_tls_certificate(db)
63 d4ec019 Dawsh 2y ago
64 d4ec019 Dawsh 2y ago await marznode.operations.add_node(db_node, certificate)
65 d4ec019 Dawsh 2y ago
66 d4ec019 Dawsh 2y ago logger.info("New node `%s` added", db_node.name)
67 d4ec019 Dawsh 2y ago return db_node
68 d4ec019 Dawsh 2y ago
69 d4ec019 Dawsh 2y ago
70 2de9d72 Dawsh 2y ago @router.get("/settings", response_model=NodeSettings)
71 01cdf2e Dawsh 2y ago def get_node_settings(db: DBDep, admin: SudoAdminDep):
72 a47eb4c Dawsh 2y ago tls = crud.get_tls_certificate(db)
73 a47eb4c Dawsh 2y ago
74 01cdf2e Dawsh 2y ago return NodeSettings(certificate=tls.certificate)
75 a47eb4c Dawsh 2y ago
76 a47eb4c Dawsh 2y ago
77 512d183 Mrdn 2y ago @router.get("/usage", response_model=NodesUsageResponse)
78 01cdf2e Dawsh 2y ago def get_usage(
79 f1ab16c Dawsh 2y ago db: DBDep,
80 f1ab16c Dawsh 2y ago admin: SudoAdminDep,
81 f1ab16c Dawsh 2y ago start_date: StartDateDep,
82 f1ab16c Dawsh 2y ago end_date: EndDateDep,
83 01cdf2e Dawsh 2y ago ):
84 512d183 Mrdn 2y ago """
85 512d183 Mrdn 2y ago Get nodes usage
86 512d183 Mrdn 2y ago """
87 512d183 Mrdn 2y ago usages = crud.get_nodes_usage(db, start_date, end_date)
88 512d183 Mrdn 2y ago
89 512d183 Mrdn 2y ago return {"usages": usages}
90 512d183 Mrdn 2y ago
91 512d183 Mrdn 2y ago
92 2de9d72 Dawsh 2y ago @router.get("/{node_id}", response_model=NodeResponse)
93 01cdf2e Dawsh 2y ago def get_node(node_id: int, db: DBDep, admin: SudoAdminDep):
94 c13a1aa Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
95 c13a1aa Dawsh 2y ago if not db_node:
96 a47eb4c Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
97 a47eb4c Dawsh 2y ago
98 c13a1aa Dawsh 2y ago return db_node
99 a47eb4c Dawsh 2y ago
100 a47eb4c Dawsh 2y ago
101 512d183 Mrdn 2y ago @router.websocket("/{node_id}/logs")
102 01cdf2e Dawsh 2y ago async def node_logs(
103 01cdf2e Dawsh 2y ago node_id: int, websocket: WebSocket, db: DBDep, include_buffer: bool = True
104 01cdf2e Dawsh 2y ago ):
105 01cdf2e Dawsh 2y ago token = websocket.query_params.get("token") or websocket.headers.get(
106 01cdf2e Dawsh 2y ago "Authorization", ""
107 01cdf2e Dawsh 2y ago ).removeprefix("Bearer ")
108 512d183 Mrdn 2y ago admin = get_admin(db, token)
109 a47eb4c Dawsh 2y ago
110 512d183 Mrdn 2y ago if not admin or not admin.is_sudo:
111 a47eb4c Dawsh 2y ago return await websocket.close(reason="You're not allowed", code=4403)
112 a47eb4c Dawsh 2y ago
113 12b255b Dawsh 2y ago if not marznode.nodes.get(node_id):
114 a47eb4c Dawsh 2y ago return await websocket.close(reason="Node not found", code=4404)
115 a47eb4c Dawsh 2y ago
116 512d183 Mrdn 2y ago # if not await marznode.nodes[node_id].is_healthy:
117 512d183 Mrdn 2y ago # return await websocket.close(reason="Node is not connected", code=4400)
118 a47eb4c Dawsh 2y ago
119 a47eb4c Dawsh 2y ago await websocket.accept()
120 5488fc0 Dawsh 2y ago try:
121 5488fc0 Dawsh 2y ago async for line in marznode.nodes[node_id].get_logs(
122 5488fc0 Dawsh 2y ago include_buffer=include_buffer
123 5488fc0 Dawsh 2y ago ):
124 5488fc0 Dawsh 2y ago try:
125 5488fc0 Dawsh 2y ago await websocket.send_text(line)
126 5488fc0 Dawsh 2y ago except WebSocketDisconnect:
127 5488fc0 Dawsh 2y ago break
128 5488fc0 Dawsh 2y ago except:
129 5488fc0 Dawsh 2y ago await websocket.close()
130 a47eb4c Dawsh 2y ago
131 a47eb4c Dawsh 2y ago
132 2de9d72 Dawsh 2y ago @router.put("/{node_id}", response_model=NodeResponse)
133 01cdf2e Dawsh 2y ago async def modify_node(
134 01cdf2e Dawsh 2y ago node_id: int, modified_node: NodeModify, db: DBDep, admin: SudoAdminDep
135 01cdf2e Dawsh 2y ago ):
136 12b255b Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
137 12b255b Dawsh 2y ago if not db_node:
138 a47eb4c Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
139 a47eb4c Dawsh 2y ago
140 12b255b Dawsh 2y ago db_node = crud.update_node(db, db_node, modified_node)
141 a47eb4c Dawsh 2y ago
142 12b255b Dawsh 2y ago await marznode.operations.remove_node(db_node.id)
143 12b255b Dawsh 2y ago if db_node.status != NodeStatus.disabled:
144 12b255b Dawsh 2y ago certificate = get_tls_certificate(db)
145 f6a32ea Dawsh 2y ago await marznode.operations.add_node(db_node, certificate)
146 a47eb4c Dawsh 2y ago
147 f7af2b3 Dawsh 2y ago logger.info("Node `%s` modified", db_node.name)
148 12b255b Dawsh 2y ago return db_node
149 a47eb4c Dawsh 2y ago
150 a47eb4c Dawsh 2y ago
151 d4ec019 Dawsh 2y ago @router.delete("/{node_id}")
152 01cdf2e Dawsh 2y ago async def remove_node(node_id: int, db: DBDep, admin: SudoAdminDep):
153 d4ec019 Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
154 d4ec019 Dawsh 2y ago if not db_node:
155 d4ec019 Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
156 d4ec019 Dawsh 2y ago
157 d4ec019 Dawsh 2y ago crud.remove_node(db, db_node)
158 d4ec019 Dawsh 2y ago await marznode.operations.remove_node(db_node.id)
159 d4ec019 Dawsh 2y ago
160 d4ec019 Dawsh 2y ago logger.info(f"Node `%s` deleted", db_node.name)
161 d4ec019 Dawsh 2y ago return {}
162 d4ec019 Dawsh 2y ago
163 d4ec019 Dawsh 2y ago
164 2de9d72 Dawsh 2y ago @router.post("/{node_id}/resync")
165 01cdf2e Dawsh 2y ago async def reconnect_node(node_id: int, db: DBDep, admin: SudoAdminDep):
166 12b255b Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
167 12b255b Dawsh 2y ago if not db_node:
168 a47eb4c Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
169 a47eb4c Dawsh 2y ago
170 a47eb4c Dawsh 2y ago return {}
171 a47eb4c Dawsh 2y ago
172 a47eb4c Dawsh 2y ago
173 512d183 Mrdn 2y ago @router.get("/{node_id}/xray_config")
174 01cdf2e Dawsh 2y ago async def get_node_xray_config(node_id: int, db: DBDep, admin: SudoAdminDep):
175 512d183 Mrdn 2y ago if not (node := marznode.nodes.get(node_id)):
176 512d183 Mrdn 2y ago raise HTTPException(status_code=404, detail="Node not found")
177 5488fc0 Dawsh 2y ago
178 5488fc0 Dawsh 2y ago try:
179 5488fc0 Dawsh 2y ago config = await node.get_xray_config()
180 5488fc0 Dawsh 2y ago except:
181 5488fc0 Dawsh 2y ago raise HTTPException(status_code=502, detail="Node isn't responsive")
182 5488fc0 Dawsh 2y ago else:
183 5488fc0 Dawsh 2y ago return json.loads(config)
184 512d183 Mrdn 2y ago
185 512d183 Mrdn 2y ago
186 512d183 Mrdn 2y ago @router.put("/{node_id}/xray_config")
187 01cdf2e Dawsh 2y ago async def alter_node_xray_config(
188 01cdf2e Dawsh 2y ago node_id: int, db: DBDep, admin: SudoAdminDep, body: Annotated[dict, Body()]
189 01cdf2e Dawsh 2y ago ):
190 512d183 Mrdn 2y ago if not (node := marznode.nodes.get(node_id)):
191 512d183 Mrdn 2y ago raise HTTPException(status_code=404, detail="Node not found")
192 512d183 Mrdn 2y ago xray_config = json.dumps(body)
193 5488fc0 Dawsh 2y ago try:
194 5488fc0 Dawsh 2y ago await node.restart_xray(xray_config)
195 5488fc0 Dawsh 2y ago except:
196 5488fc0 Dawsh 2y ago raise HTTPException(status_code=502, detail="Node isn't responsive")
197 512d183 Mrdn 2y ago return {}