gitscope
https://dev.kronoz.win/r/marzneshin.git
app/routes/node.py · v0.3.0
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 d4ec019 Dawsh 2y ago from app.dependencies import DBDep, SudoAdminDep, EndDateDep, StartDateDep, get_admin
16 01cdf2e Dawsh 2y ago from app.models.node import (
17 01cdf2e Dawsh 2y ago NodeCreate,
18 01cdf2e Dawsh 2y ago NodeModify,
19 01cdf2e Dawsh 2y ago NodeResponse,
20 01cdf2e Dawsh 2y ago NodeSettings,
21 01cdf2e Dawsh 2y ago NodeStatus,
22 01cdf2e Dawsh 2y ago NodesUsageResponse,
23 01cdf2e Dawsh 2y ago )
24 f7af2b3 Dawsh 2y ago
25 f7af2b3 Dawsh 2y ago logger = logging.getLogger(__name__)
26 01cdf2e Dawsh 2y ago router = APIRouter(prefix="/nodes", tags=["Node"])
27 a47eb4c Dawsh 2y ago
28 a47eb4c Dawsh 2y ago
29 d4ec019 Dawsh 2y ago @router.get("", response_model=Page[NodeResponse])
30 01cdf2e Dawsh 2y ago def get_nodes(
31 01cdf2e Dawsh 2y ago db: DBDep,
32 01cdf2e Dawsh 2y ago admin: SudoAdminDep,
33 01cdf2e Dawsh 2y ago status: list[NodeStatus] = Query(None),
34 01cdf2e Dawsh 2y ago name: str = Query(None),
35 01cdf2e Dawsh 2y ago ):
36 d4ec019 Dawsh 2y ago query = db.query(Node)
37 d4ec019 Dawsh 2y ago
38 d4ec019 Dawsh 2y ago if name:
39 d4ec019 Dawsh 2y ago query = query.filter(Node.name.ilike(f"%{name}%"))
40 d4ec019 Dawsh 2y ago
41 d4ec019 Dawsh 2y ago if status:
42 d4ec019 Dawsh 2y ago query = query.filter(Node.status.in_(status))
43 d4ec019 Dawsh 2y ago
44 d4ec019 Dawsh 2y ago return paginate(db, query)
45 d4ec019 Dawsh 2y ago
46 d4ec019 Dawsh 2y ago
47 d4ec019 Dawsh 2y ago @router.post("", response_model=NodeResponse)
48 01cdf2e Dawsh 2y ago async def add_node(new_node: NodeCreate, db: DBDep, admin: SudoAdminDep):
49 d4ec019 Dawsh 2y ago try:
50 d4ec019 Dawsh 2y ago db_node = crud.create_node(db, new_node)
51 d4ec019 Dawsh 2y ago except sqlalchemy.exc.IntegrityError:
52 d4ec019 Dawsh 2y ago db.rollback()
53 01cdf2e Dawsh 2y ago raise HTTPException(
54 01cdf2e Dawsh 2y ago status_code=409, detail=f'Node "{new_node.name}" already exists'
55 01cdf2e Dawsh 2y ago )
56 d4ec019 Dawsh 2y ago certificate = get_tls_certificate(db)
57 d4ec019 Dawsh 2y ago
58 d4ec019 Dawsh 2y ago await marznode.operations.add_node(db_node, certificate)
59 d4ec019 Dawsh 2y ago
60 d4ec019 Dawsh 2y ago logger.info("New node `%s` added", db_node.name)
61 d4ec019 Dawsh 2y ago return db_node
62 d4ec019 Dawsh 2y ago
63 d4ec019 Dawsh 2y ago
64 2de9d72 Dawsh 2y ago @router.get("/settings", response_model=NodeSettings)
65 01cdf2e Dawsh 2y ago def get_node_settings(db: DBDep, admin: SudoAdminDep):
66 a47eb4c Dawsh 2y ago tls = crud.get_tls_certificate(db)
67 a47eb4c Dawsh 2y ago
68 01cdf2e Dawsh 2y ago return NodeSettings(certificate=tls.certificate)
69 a47eb4c Dawsh 2y ago
70 a47eb4c Dawsh 2y ago
71 512d183 Mrdn 2y ago @router.get("/usage", response_model=NodesUsageResponse)
72 01cdf2e Dawsh 2y ago def get_usage(
73 01cdf2e Dawsh 2y ago db: DBDep, admin: SudoAdminDep, start_date: StartDateDep, end_date: EndDateDep
74 01cdf2e Dawsh 2y ago ):
75 512d183 Mrdn 2y ago """
76 512d183 Mrdn 2y ago Get nodes usage
77 512d183 Mrdn 2y ago """
78 512d183 Mrdn 2y ago usages = crud.get_nodes_usage(db, start_date, end_date)
79 512d183 Mrdn 2y ago
80 512d183 Mrdn 2y ago return {"usages": usages}
81 512d183 Mrdn 2y ago
82 512d183 Mrdn 2y ago
83 2de9d72 Dawsh 2y ago @router.get("/{node_id}", response_model=NodeResponse)
84 01cdf2e Dawsh 2y ago def get_node(node_id: int, db: DBDep, admin: SudoAdminDep):
85 c13a1aa Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
86 c13a1aa Dawsh 2y ago if not db_node:
87 a47eb4c Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
88 a47eb4c Dawsh 2y ago
89 c13a1aa Dawsh 2y ago return db_node
90 a47eb4c Dawsh 2y ago
91 a47eb4c Dawsh 2y ago
92 512d183 Mrdn 2y ago @router.websocket("/{node_id}/logs")
93 01cdf2e Dawsh 2y ago async def node_logs(
94 01cdf2e Dawsh 2y ago node_id: int, websocket: WebSocket, db: DBDep, include_buffer: bool = True
95 01cdf2e Dawsh 2y ago ):
96 01cdf2e Dawsh 2y 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 512d183 Mrdn 2y ago # if not await marznode.nodes[node_id].is_healthy:
108 512d183 Mrdn 2y ago # return await websocket.close(reason="Node is not connected", code=4400)
109 a47eb4c Dawsh 2y ago
110 a47eb4c Dawsh 2y ago await websocket.accept()
111 5488fc0 Dawsh 2y ago try:
112 5488fc0 Dawsh 2y ago async for line in marznode.nodes[node_id].get_logs(
113 5488fc0 Dawsh 2y ago include_buffer=include_buffer
114 5488fc0 Dawsh 2y ago ):
115 5488fc0 Dawsh 2y ago try:
116 5488fc0 Dawsh 2y ago await websocket.send_text(line)
117 5488fc0 Dawsh 2y ago except WebSocketDisconnect:
118 5488fc0 Dawsh 2y ago break
119 5488fc0 Dawsh 2y ago except:
120 5488fc0 Dawsh 2y 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 512d183 Mrdn 2y ago @router.get("/{node_id}/xray_config")
165 01cdf2e Dawsh 2y ago async def get_node_xray_config(node_id: int, db: DBDep, admin: SudoAdminDep):
166 512d183 Mrdn 2y ago if not (node := marznode.nodes.get(node_id)):
167 512d183 Mrdn 2y ago raise HTTPException(status_code=404, detail="Node not found")
168 5488fc0 Dawsh 2y ago
169 5488fc0 Dawsh 2y ago try:
170 5488fc0 Dawsh 2y ago config = await node.get_xray_config()
171 5488fc0 Dawsh 2y ago except:
172 5488fc0 Dawsh 2y ago raise HTTPException(status_code=502, detail="Node isn't responsive")
173 5488fc0 Dawsh 2y ago else:
174 5488fc0 Dawsh 2y ago return json.loads(config)
175 512d183 Mrdn 2y ago
176 512d183 Mrdn 2y ago
177 512d183 Mrdn 2y ago @router.put("/{node_id}/xray_config")
178 01cdf2e Dawsh 2y ago async def alter_node_xray_config(
179 01cdf2e Dawsh 2y ago node_id: int, db: DBDep, admin: SudoAdminDep, body: Annotated[dict, Body()]
180 01cdf2e Dawsh 2y ago ):
181 512d183 Mrdn 2y ago if not (node := marznode.nodes.get(node_id)):
182 512d183 Mrdn 2y ago raise HTTPException(status_code=404, detail="Node not found")
183 512d183 Mrdn 2y ago xray_config = json.dumps(body)
184 5488fc0 Dawsh 2y ago try:
185 5488fc0 Dawsh 2y ago await node.restart_xray(xray_config)
186 5488fc0 Dawsh 2y ago except:
187 5488fc0 Dawsh 2y ago raise HTTPException(status_code=502, detail="Node isn't responsive")
188 512d183 Mrdn 2y ago return {}