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