gitscope
https://dev.kronoz.win/r/marzneshin.git
app/routes/node.py · v0.3.0-alpha9
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 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 d4ec019 Dawsh 2y ago from app.dependencies import DBDep, SudoAdminDep, EndDateDep, StartDateDep, get_admin
15 d4ec019 Dawsh 2y ago
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 512d183 Mrdn 2y ago async for line in marznode.nodes[node_id].get_logs(include_buffer=include_buffer):
112 512d183 Mrdn 2y ago await websocket.send_text(line)
113 a47eb4c Dawsh 2y ago
114 a47eb4c Dawsh 2y ago
115 2de9d72 Dawsh 2y ago @router.put("/{node_id}", response_model=NodeResponse)
116 01cdf2e Dawsh 2y ago async def modify_node(
117 01cdf2e Dawsh 2y ago node_id: int, modified_node: NodeModify, db: DBDep, admin: SudoAdminDep
118 01cdf2e Dawsh 2y ago ):
119 12b255b Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
120 12b255b Dawsh 2y ago if not db_node:
121 a47eb4c Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
122 a47eb4c Dawsh 2y ago
123 12b255b Dawsh 2y ago db_node = crud.update_node(db, db_node, modified_node)
124 a47eb4c Dawsh 2y ago
125 12b255b Dawsh 2y ago await marznode.operations.remove_node(db_node.id)
126 12b255b Dawsh 2y ago if db_node.status != NodeStatus.disabled:
127 12b255b Dawsh 2y ago certificate = get_tls_certificate(db)
128 f6a32ea Dawsh 2y ago await marznode.operations.add_node(db_node, certificate)
129 a47eb4c Dawsh 2y ago
130 f7af2b3 Dawsh 2y ago logger.info("Node `%s` modified", db_node.name)
131 12b255b Dawsh 2y ago return db_node
132 a47eb4c Dawsh 2y ago
133 a47eb4c Dawsh 2y ago
134 d4ec019 Dawsh 2y ago @router.delete("/{node_id}")
135 01cdf2e Dawsh 2y ago async def remove_node(node_id: int, db: DBDep, admin: SudoAdminDep):
136 d4ec019 Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
137 d4ec019 Dawsh 2y ago if not db_node:
138 d4ec019 Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
139 d4ec019 Dawsh 2y ago
140 d4ec019 Dawsh 2y ago crud.remove_node(db, db_node)
141 d4ec019 Dawsh 2y ago await marznode.operations.remove_node(db_node.id)
142 d4ec019 Dawsh 2y ago
143 d4ec019 Dawsh 2y ago logger.info(f"Node `%s` deleted", db_node.name)
144 d4ec019 Dawsh 2y ago return {}
145 d4ec019 Dawsh 2y ago
146 d4ec019 Dawsh 2y ago
147 2de9d72 Dawsh 2y ago @router.post("/{node_id}/resync")
148 01cdf2e Dawsh 2y ago async def reconnect_node(node_id: int, db: DBDep, admin: SudoAdminDep):
149 12b255b Dawsh 2y ago db_node = crud.get_node_by_id(db, node_id)
150 12b255b Dawsh 2y ago if not db_node:
151 a47eb4c Dawsh 2y ago raise HTTPException(status_code=404, detail="Node not found")
152 a47eb4c Dawsh 2y ago
153 a47eb4c Dawsh 2y ago return {}
154 a47eb4c Dawsh 2y ago
155 a47eb4c Dawsh 2y ago
156 512d183 Mrdn 2y ago @router.get("/{node_id}/xray_config")
157 01cdf2e Dawsh 2y ago async def get_node_xray_config(node_id: int, db: DBDep, admin: SudoAdminDep):
158 512d183 Mrdn 2y ago if not (node := marznode.nodes.get(node_id)):
159 512d183 Mrdn 2y ago raise HTTPException(status_code=404, detail="Node not found")
160 512d183 Mrdn 2y ago return json.loads(await node.get_xray_config())
161 512d183 Mrdn 2y ago
162 512d183 Mrdn 2y ago
163 512d183 Mrdn 2y ago @router.put("/{node_id}/xray_config")
164 01cdf2e Dawsh 2y ago async def alter_node_xray_config(
165 01cdf2e Dawsh 2y ago node_id: int, db: DBDep, admin: SudoAdminDep, body: Annotated[dict, Body()]
166 01cdf2e Dawsh 2y ago ):
167 512d183 Mrdn 2y ago if not (node := marznode.nodes.get(node_id)):
168 512d183 Mrdn 2y ago raise HTTPException(status_code=404, detail="Node not found")
169 512d183 Mrdn 2y ago xray_config = json.dumps(body)
170 512d183 Mrdn 2y ago await node.restart_xray(xray_config)
171 512d183 Mrdn 2y ago return {}