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