gitscope
https://dev.kronoz.win/r/marzneshin.git
app/routes/node.py · v0.2.0
Source
1 512d183 Mrdn 2y ago import json
2 f7af2b3 Dawsh 2y ago import logging
3 a47eb4c Dawsh 2y ago from datetime import datetime
4 512d183 Mrdn 2y ago from typing import List, Annotated
5 a47eb4c Dawsh 2y ago
6 a47eb4c Dawsh 2y ago import sqlalchemy
7 512d183 Mrdn 2y ago from fastapi import APIRouter, Depends, Body
8 c13a1aa Dawsh 2y ago from fastapi import HTTPException, WebSocket
9 512d183 Mrdn 2y ago from starlette.requests import Request
10 512d183 Mrdn 2y ago from starlette.responses import PlainTextResponse
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 512d183 Mrdn 2y ago from app.dependencies import DBDep, SudoAdminDep, sudo_admin, EndDateDep, StartDateDep, get_admin
15 f6a32ea Dawsh 2y ago from app.marznode import MarzNodeGRPCIO
16 a47eb4c Dawsh 2y ago from app.models.admin import Admin
17 a47eb4c Dawsh 2y ago from app.models.node import (NodeCreate, NodeModify, NodeResponse,
18 a47eb4c Dawsh 2y ago NodeSettings, NodeStatus, NodesUsageResponse)
19 f7af2b3 Dawsh 2y ago
20 f7af2b3 Dawsh 2y ago logger = logging.getLogger(__name__)
21 512d183 Mrdn 2y ago router = APIRouter(prefix="/nodes", tags=['Node'])
22 a47eb4c Dawsh 2y ago
23 a47eb4c Dawsh 2y ago
24 2de9d72 Dawsh 2y ago @router.get("/settings", response_model=NodeSettings)
25 c13a1aa Dawsh 2y ago def get_node_settings(db: DBDep,
26 c13a1aa Dawsh 2y ago admin: SudoAdminDep):
27 a47eb4c Dawsh 2y ago tls = crud.get_tls_certificate(db)
28 a47eb4c Dawsh 2y ago
29 a47eb4c Dawsh 2y ago return NodeSettings(
30 a47eb4c Dawsh 2y ago certificate=tls.certificate
31 a47eb4c Dawsh 2y ago )
32 a47eb4c Dawsh 2y ago
33 a47eb4c Dawsh 2y ago
34 512d183 Mrdn 2y ago @router.get("/usage", response_model=NodesUsageResponse)
35 512d183 Mrdn 2y ago def get_usage(db: DBDep,
36 512d183 Mrdn 2y ago admin: SudoAdminDep,
37 512d183 Mrdn 2y ago start_date: StartDateDep,
38 512d183 Mrdn 2y ago end_date: EndDateDep):
39 512d183 Mrdn 2y ago """
40 512d183 Mrdn 2y ago Get nodes usage
41 512d183 Mrdn 2y ago """
42 512d183 Mrdn 2y ago usages = crud.get_nodes_usage(db, start_date, end_date)
43 512d183 Mrdn 2y ago
44 512d183 Mrdn 2y ago return {"usages": usages}
45 512d183 Mrdn 2y ago
46 512d183 Mrdn 2y ago
47 2de9d72 Dawsh 2y ago @router.post("", response_model=NodeResponse)
48 a47eb4c Dawsh 2y ago async def add_node(new_node: NodeCreate,
49 c13a1aa Dawsh 2y ago db: DBDep,
50 c13a1aa Dawsh 2y ago admin: SudoAdminDep):
51 a47eb4c Dawsh 2y ago try:
52 12b255b Dawsh 2y ago db_node = crud.create_node(db, new_node)
53 a47eb4c Dawsh 2y ago except sqlalchemy.exc.IntegrityError:
54 a47eb4c Dawsh 2y ago db.rollback()
55 a47eb4c Dawsh 2y ago raise HTTPException(status_code=409, detail=f"Node \"{new_node.name}\" already exists")
56 12b255b Dawsh 2y ago certificate = get_tls_certificate(db)
57 a47eb4c Dawsh 2y ago
58 f6a32ea Dawsh 2y ago await marznode.operations.add_node(db_node, certificate)
59 a47eb4c Dawsh 2y ago
60 f7af2b3 Dawsh 2y ago logger.info("New node `%s` added", db_node.name)
61 12b255b Dawsh 2y ago return db_node
62 a47eb4c Dawsh 2y ago
63 a47eb4c Dawsh 2y ago
64 2de9d72 Dawsh 2y ago @router.get("/{node_id}", response_model=NodeResponse)
65 a47eb4c Dawsh 2y ago def get_node(node_id: int,
66 c13a1aa Dawsh 2y ago db: DBDep,
67 c13a1aa Dawsh 2y ago admin: SudoAdminDep):
68 c13a1aa Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
69 c13a1aa Dawsh 2y ago if not db_node:
70 a47eb4c Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
71 a47eb4c Dawsh 2y ago
72 c13a1aa Dawsh 2y ago return db_node
73 a47eb4c Dawsh 2y ago
74 a47eb4c Dawsh 2y ago
75 512d183 Mrdn 2y ago @router.websocket("/{node_id}/logs")
76 c13a1aa Dawsh 2y ago async def node_logs(node_id: int,
77 c13a1aa Dawsh 2y ago websocket: WebSocket,
78 512d183 Mrdn 2y ago db: DBDep,
79 512d183 Mrdn 2y ago include_buffer: bool = True):
80 a47eb4c Dawsh 2y ago token = (
81 a47eb4c Dawsh 2y ago websocket.query_params.get('token')
82 a47eb4c Dawsh 2y ago or
83 a47eb4c Dawsh 2y ago websocket.headers.get('Authorization', '').removeprefix("Bearer ")
84 a47eb4c Dawsh 2y ago )
85 512d183 Mrdn 2y ago admin = get_admin(db, token)
86 a47eb4c Dawsh 2y ago
87 512d183 Mrdn 2y ago if not admin or not admin.is_sudo:
88 a47eb4c Dawsh 2y ago return await websocket.close(reason="You're not allowed", code=4403)
89 a47eb4c Dawsh 2y ago
90 12b255b Dawsh 2y ago if not marznode.nodes.get(node_id):
91 a47eb4c Dawsh 2y ago return await websocket.close(reason="Node not found", code=4404)
92 a47eb4c Dawsh 2y ago
93 512d183 Mrdn 2y ago # if not await marznode.nodes[node_id].is_healthy:
94 512d183 Mrdn 2y ago # return await websocket.close(reason="Node is not connected", code=4400)
95 a47eb4c Dawsh 2y ago
96 a47eb4c Dawsh 2y ago await websocket.accept()
97 512d183 Mrdn 2y ago async for line in marznode.nodes[node_id].get_logs(include_buffer=include_buffer):
98 512d183 Mrdn 2y ago await websocket.send_text(line)
99 a47eb4c Dawsh 2y ago
100 a47eb4c Dawsh 2y ago
101 5bb2712 Mardin 2y ago @router.get("", response_model=List[NodeResponse])
102 c13a1aa Dawsh 2y ago def get_nodes(db: DBDep,
103 c13a1aa Dawsh 2y ago admin: SudoAdminDep):
104 a47eb4c Dawsh 2y ago return crud.get_nodes(db)
105 a47eb4c Dawsh 2y ago
106 a47eb4c Dawsh 2y ago
107 2de9d72 Dawsh 2y ago @router.put("/{node_id}", response_model=NodeResponse)
108 a47eb4c Dawsh 2y ago async def modify_node(node_id: int,
109 12b255b Dawsh 2y ago modified_node: NodeModify,
110 c13a1aa Dawsh 2y ago db: DBDep,
111 c13a1aa Dawsh 2y ago admin: SudoAdminDep):
112 12b255b Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
113 12b255b Dawsh 2y ago if not db_node:
114 a47eb4c Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
115 a47eb4c Dawsh 2y ago
116 12b255b Dawsh 2y ago db_node = crud.update_node(db, db_node, modified_node)
117 a47eb4c Dawsh 2y ago
118 12b255b Dawsh 2y ago await marznode.operations.remove_node(db_node.id)
119 12b255b Dawsh 2y ago if db_node.status != NodeStatus.disabled:
120 12b255b Dawsh 2y ago certificate = get_tls_certificate(db)
121 f6a32ea Dawsh 2y ago await marznode.operations.add_node(db_node, certificate)
122 a47eb4c Dawsh 2y ago
123 f7af2b3 Dawsh 2y ago logger.info("Node `%s` modified", db_node.name)
124 12b255b Dawsh 2y ago return db_node
125 a47eb4c Dawsh 2y ago
126 a47eb4c Dawsh 2y ago
127 2de9d72 Dawsh 2y ago @router.post("/{node_id}/resync")
128 a47eb4c Dawsh 2y ago async def reconnect_node(node_id: int,
129 c13a1aa Dawsh 2y ago db: DBDep,
130 c13a1aa Dawsh 2y ago admin: SudoAdminDep):
131 12b255b Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
132 12b255b Dawsh 2y ago if not db_node:
133 a47eb4c Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
134 a47eb4c Dawsh 2y ago
135 a47eb4c Dawsh 2y ago return {}
136 a47eb4c Dawsh 2y ago
137 a47eb4c Dawsh 2y ago
138 512d183 Mrdn 2y ago @router.get("/{node_id}/xray_config")
139 512d183 Mrdn 2y ago async def get_node_xray_config(node_id: int,
140 512d183 Mrdn 2y ago db: DBDep,
141 512d183 Mrdn 2y ago admin: SudoAdminDep):
142 512d183 Mrdn 2y ago if not (node := marznode.nodes.get(node_id)):
143 512d183 Mrdn 2y ago raise HTTPException(status_code=404, detail="Node not found")
144 512d183 Mrdn 2y ago return json.loads(await node.get_xray_config())
145 512d183 Mrdn 2y ago
146 512d183 Mrdn 2y ago
147 512d183 Mrdn 2y ago @router.put("/{node_id}/xray_config")
148 512d183 Mrdn 2y ago async def alter_node_xray_config(node_id: int,
149 512d183 Mrdn 2y ago db: DBDep,
150 512d183 Mrdn 2y ago admin: SudoAdminDep,
151 512d183 Mrdn 2y ago body: Annotated[dict, Body()]):
152 512d183 Mrdn 2y ago if not (node := marznode.nodes.get(node_id)):
153 512d183 Mrdn 2y ago raise HTTPException(status_code=404, detail="Node not found")
154 512d183 Mrdn 2y ago xray_config = json.dumps(body)
155 512d183 Mrdn 2y ago await node.restart_xray(xray_config)
156 512d183 Mrdn 2y ago return {}
157 512d183 Mrdn 2y ago
158 512d183 Mrdn 2y ago
159 2de9d72 Dawsh 2y ago @router.delete("/{node_id}")
160 a47eb4c Dawsh 2y ago async def remove_node(node_id: int,
161 c13a1aa Dawsh 2y ago db: DBDep,
162 c13a1aa Dawsh 2y ago admin: SudoAdminDep):
163 12b255b Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
164 12b255b Dawsh 2y ago if not db_node:
165 a47eb4c Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
166 a47eb4c Dawsh 2y ago
167 12b255b Dawsh 2y ago crud.remove_node(db, db_node)
168 12b255b Dawsh 2y ago await marznode.operations.remove_node(db_node.id)
169 a47eb4c Dawsh 2y ago
170 f7af2b3 Dawsh 2y ago logger.info(f"Node `%s` deleted", db_node.name)
171 a47eb4c Dawsh 2y ago return {}