gitscope
https://dev.kronoz.win/r/marzneshin.git
app/db/crud.py · v0.5.6
Source
1 12b255b Dawsh 2y ago import json
2 34a7612 Dawsh 2y ago import secrets
3 f446990 Dawsh 2y ago from datetime import datetime, timedelta
4 a47eb4c Dawsh 2y ago from enum import Enum
5 4708fcb Dawsh 2y ago from types import NoneType
6 2de9d72 Dawsh 2y ago from typing import List, Optional, Tuple, Union
7 a47eb4c Dawsh 2y ago
8 512d183 Mrdn 2y ago from sqlalchemy import and_, delete, update, select
9 a47eb4c Dawsh 2y ago from sqlalchemy.orm import Session
10 a47eb4c Dawsh 2y ago
11 01cdf2e Dawsh 2y ago from app.db.models import (
12 01cdf2e Dawsh 2y ago JWT,
13 01cdf2e Dawsh 2y ago TLS,
14 01cdf2e Dawsh 2y ago Admin,
15 01cdf2e Dawsh 2y ago Node,
16 01cdf2e Dawsh 2y ago NodeUsage,
17 01cdf2e Dawsh 2y ago NodeUserUsage,
18 01cdf2e Dawsh 2y ago NotificationReminder,
19 01cdf2e Dawsh 2y ago InboundHost,
20 01cdf2e Dawsh 2y ago Service,
21 01cdf2e Dawsh 2y ago Inbound,
22 01cdf2e Dawsh 2y ago System,
23 01cdf2e Dawsh 2y ago User,
24 f5c168d Mardin 2y ago Backend,
25 01cdf2e Dawsh 2y ago )
26 4708fcb Dawsh 2y ago from app.models.admin import AdminCreate, AdminPartialModify
27 f1ab16c Dawsh 2y ago from app.models.node import (
28 f1ab16c Dawsh 2y ago NodeCreate,
29 f1ab16c Dawsh 2y ago NodeModify,
30 f1ab16c Dawsh 2y ago NodeStatus,
31 f1ab16c Dawsh 2y ago NodeUsageResponse,
32 f1ab16c Dawsh 2y ago )
33 a47eb4c Dawsh 2y ago from app.models.proxy import InboundHost as InboundHostModify
34 2de9d72 Dawsh 2y ago from app.models.service import Service as ServiceModify, ServiceCreate
35 01cdf2e Dawsh 2y ago from app.models.user import (
36 01cdf2e Dawsh 2y ago ReminderType,
37 01cdf2e Dawsh 2y ago UserCreate,
38 d377c4c Dawsh 2y ago UserDataUsageResetStrategy,
39 01cdf2e Dawsh 2y ago UserModify,
40 01cdf2e Dawsh 2y ago UserStatus,
41 01cdf2e Dawsh 2y ago UserUsageResponse,
42 d377c4c Dawsh 2y ago UserExpireStrategy,
43 f1ab16c Dawsh 2y ago )
44 a47eb4c Dawsh 2y ago
45 a47eb4c Dawsh 2y ago
46 a47eb4c Dawsh 2y ago def add_default_hosts(db: Session, inbounds: List[Inbound]):
47 01cdf2e Dawsh 2y ago hosts = [
48 01cdf2e Dawsh 2y ago InboundHost(
49 01cdf2e Dawsh 2y ago remark="🚀 Marz ({USERNAME}) [{PROTOCOL} - {TRANSPORT}]",
50 01cdf2e Dawsh 2y ago address="{SERVER_IP}",
51 01cdf2e Dawsh 2y ago inbound=i,
52 01cdf2e Dawsh 2y ago )
53 01cdf2e Dawsh 2y ago for i in inbounds
54 01cdf2e Dawsh 2y ago ]
55 a47eb4c Dawsh 2y ago db.add_all(hosts)
56 a47eb4c Dawsh 2y ago db.commit()
57 a47eb4c Dawsh 2y ago
58 a47eb4c Dawsh 2y ago
59 f5c168d Mardin 2y ago def ensure_node_backends(db: Session, backends, node_id: int):
60 f5c168d Mardin 2y ago old_backends = db.query(Backend).where(Backend.node_id == node_id)
61 f5c168d Mardin 2y ago for backend in old_backends:
62 f5c168d Mardin 2y ago db.delete(backend)
63 f5c168d Mardin 2y ago backends = [
64 f5c168d Mardin 2y ago Backend(
65 f5c168d Mardin 2y ago name=backend.name,
66 f5c168d Mardin 2y ago backend_type=backend.type,
67 f5c168d Mardin 2y ago version=backend.version,
68 f5c168d Mardin 2y ago node_id=node_id,
69 f5c168d Mardin 2y ago )
70 f5c168d Mardin 2y ago for backend in backends
71 f5c168d Mardin 2y ago ]
72 f5c168d Mardin 2y ago db.add_all(backends)
73 f5c168d Mardin 2y ago db.flush()
74 f5c168d Mardin 2y ago
75 f5c168d Mardin 2y ago
76 f5c168d Mardin 2y ago def ensure_node_inbounds(db: Session, inbounds: List[Inbound], node_id: int):
77 01cdf2e Dawsh 2y ago current_tags = [
78 01cdf2e Dawsh 2y ago i[0]
79 01cdf2e Dawsh 2y ago for i in db.execute(
80 01cdf2e Dawsh 2y ago select(Inbound.tag).filter(Inbound.node_id == node_id)
81 01cdf2e Dawsh 2y ago ).all()
82 01cdf2e Dawsh 2y ago ]
83 512d183 Mrdn 2y ago updated_tags = set(i.tag for i in list(inbounds))
84 512d183 Mrdn 2y ago inbound_additions, tag_deletions = list(), set()
85 512d183 Mrdn 2y ago for tag in current_tags:
86 512d183 Mrdn 2y ago if tag not in updated_tags:
87 512d183 Mrdn 2y ago tag_deletions.add(tag)
88 f5c168d Mardin 2y ago removals = db.query(Inbound).where(
89 01cdf2e Dawsh 2y ago and_(Inbound.node_id == node_id, Inbound.tag.in_(tag_deletions))
90 01cdf2e Dawsh 2y ago )
91 f5c168d Mardin 2y ago for i in removals:
92 512d183 Mrdn 2y ago db.delete(i)
93 512d183 Mrdn 2y ago
94 512d183 Mrdn 2y ago for inb in inbounds:
95 512d183 Mrdn 2y ago if inb.tag in current_tags:
96 01cdf2e Dawsh 2y ago stmt = (
97 01cdf2e Dawsh 2y ago update(Inbound)
98 f1ab16c Dawsh 2y ago .where(
99 f1ab16c Dawsh 2y ago and_(Inbound.node_id == node_id, Inbound.tag == inb.tag)
100 f1ab16c Dawsh 2y ago )
101 f1ab16c Dawsh 2y ago .values(
102 f1ab16c Dawsh 2y ago protocol=json.loads(inb.config)["protocol"],
103 f1ab16c Dawsh 2y ago config=inb.config,
104 f1ab16c Dawsh 2y ago )
105 a47eb4c Dawsh 2y ago )
106 512d183 Mrdn 2y ago db.execute(stmt)
107 512d183 Mrdn 2y ago else:
108 512d183 Mrdn 2y ago inbound_additions.append(inb)
109 01cdf2e Dawsh 2y ago new_inbounds = [
110 01cdf2e Dawsh 2y ago Inbound(
111 01cdf2e Dawsh 2y ago tag=inb.tag,
112 01cdf2e Dawsh 2y ago protocol=json.loads(inb.config)["protocol"],
113 01cdf2e Dawsh 2y ago config=inb.config,
114 01cdf2e Dawsh 2y ago node_id=node_id,
115 01cdf2e Dawsh 2y ago )
116 01cdf2e Dawsh 2y ago for inb in inbound_additions
117 01cdf2e Dawsh 2y ago ]
118 a47eb4c Dawsh 2y ago db.add_all(new_inbounds)
119 a47eb4c Dawsh 2y ago db.flush()
120 a47eb4c Dawsh 2y ago for i in new_inbounds:
121 2de9d72 Dawsh 2y ago db.refresh(i)
122 a47eb4c Dawsh 2y ago add_default_hosts(db, new_inbounds)
123 a47eb4c Dawsh 2y ago db.commit()
124 a47eb4c Dawsh 2y ago
125 a47eb4c Dawsh 2y ago
126 01cdf2e Dawsh 2y ago def get_node_users(
127 f1ab16c Dawsh 2y ago db: Session,
128 d377c4c Dawsh 2y ago node_id: int,
129 01cdf2e Dawsh 2y ago ):
130 01cdf2e Dawsh 2y ago query = (
131 01cdf2e Dawsh 2y ago db.query(User.id, User.username, User.key, Inbound)
132 01cdf2e Dawsh 2y ago .distinct()
133 01cdf2e Dawsh 2y ago .join(Inbound.services)
134 01cdf2e Dawsh 2y ago .join(Service.users)
135 01cdf2e Dawsh 2y ago .filter(Inbound.node_id == node_id)
136 d377c4c Dawsh 2y ago .filter(User.activated == True)
137 01cdf2e Dawsh 2y ago )
138 a47eb4c Dawsh 2y ago return query.all()
139 a47eb4c Dawsh 2y ago
140 12b255b Dawsh 2y ago
141 a47eb4c Dawsh 2y ago def get_user_hosts(db: Session, user_id: int):
142 01cdf2e Dawsh 2y ago return (
143 01cdf2e Dawsh 2y ago db.query(InboundHost)
144 01cdf2e Dawsh 2y ago .distinct()
145 01cdf2e Dawsh 2y ago .join(User.services)
146 01cdf2e Dawsh 2y ago .join(Service.inbounds)
147 01cdf2e Dawsh 2y ago .join(Inbound.hosts)
148 01cdf2e Dawsh 2y ago .filter(User.id == user_id)
149 01cdf2e Dawsh 2y ago .all()
150 01cdf2e Dawsh 2y ago )
151 a47eb4c Dawsh 2y ago
152 12b255b Dawsh 2y ago
153 a47eb4c Dawsh 2y ago def get_inbound_hosts(db: Session, inbound_id: int) -> List[InboundHost]:
154 f1ab16c Dawsh 2y ago return (
155 f1ab16c Dawsh 2y ago db.query(InboundHost)
156 f1ab16c Dawsh 2y ago .filter(InboundHost.inbound_id == inbound_id)
157 f1ab16c Dawsh 2y ago .all()
158 f1ab16c Dawsh 2y ago )
159 a47eb4c Dawsh 2y ago
160 12b255b Dawsh 2y ago
161 a47eb4c Dawsh 2y ago def get_all_inbounds(db: Session):
162 a47eb4c Dawsh 2y ago return db.query(Inbound).all()
163 a47eb4c Dawsh 2y ago
164 a47eb4c Dawsh 2y ago
165 2de9d72 Dawsh 2y ago def get_inbound(db: Session, inbound_id: int) -> Inbound | None:
166 2de9d72 Dawsh 2y ago return db.query(Inbound).filter(Inbound.id == inbound_id).first()
167 a47eb4c Dawsh 2y ago
168 a47eb4c Dawsh 2y ago
169 2de9d72 Dawsh 2y ago def get_host(db: Session, host_id) -> InboundHost:
170 2de9d72 Dawsh 2y ago return db.query(InboundHost).filter(InboundHost.id == host_id).first()
171 a47eb4c Dawsh 2y ago
172 a47eb4c Dawsh 2y ago
173 a47eb4c Dawsh 2y ago def add_host(db: Session, inbound: Inbound, host: InboundHostModify):
174 2de9d72 Dawsh 2y ago host = InboundHost(
175 2de9d72 Dawsh 2y ago remark=host.remark,
176 2de9d72 Dawsh 2y ago address=host.address,
177 2de9d72 Dawsh 2y ago port=host.port,
178 2de9d72 Dawsh 2y ago path=host.path,
179 2de9d72 Dawsh 2y ago sni=host.sni,
180 2de9d72 Dawsh 2y ago host=host.host,
181 2de9d72 Dawsh 2y ago inbound=inbound,
182 2de9d72 Dawsh 2y ago security=host.security,
183 244c9d8 Dawsh 2y ago alpn=host.alpn.value,
184 01cdf2e Dawsh 2y ago fingerprint=host.fingerprint,
185 0b152b5 Mrdn 2y ago fragment=host.fragment.model_dump() if host.fragment else None,
186 0b152b5 Mrdn 2y ago mux=host.mux,
187 dd32f70 Dawsh 2y ago allowinsecure=host.allowinsecure,
188 e422d72 Dawsh 2y ago weight=host.weight,
189 a47eb4c Dawsh 2y ago )
190 2de9d72 Dawsh 2y ago inbound.hosts.append(host)
191 a47eb4c Dawsh 2y ago db.commit()
192 2de9d72 Dawsh 2y ago db.refresh(host)
193 2de9d72 Dawsh 2y ago return host
194 2de9d72 Dawsh 2y ago
195 2de9d72 Dawsh 2y ago
196 2de9d72 Dawsh 2y ago def update_host(db: Session, db_host: InboundHost, host: InboundHostModify):
197 2de9d72 Dawsh 2y ago db_host.remark = host.remark
198 2de9d72 Dawsh 2y ago db_host.address = host.address
199 2de9d72 Dawsh 2y ago db_host.port = host.port
200 2de9d72 Dawsh 2y ago db_host.path = host.path
201 2de9d72 Dawsh 2y ago db_host.sni = host.sni
202 2de9d72 Dawsh 2y ago db_host.host = host.host
203 2de9d72 Dawsh 2y ago db_host.security = host.security
204 244c9d8 Dawsh 2y ago db_host.alpn = host.alpn.value
205 2de9d72 Dawsh 2y ago db_host.fingerprint = host.fingerprint
206 0b152b5 Mrdn 2y ago db_host.fragment = host.fragment.model_dump() if host.fragment else None
207 0b152b5 Mrdn 2y ago db_host.mux = host.mux
208 a88ed47 Mardin 2y ago db_host.is_disabled = host.is_disabled
209 dd32f70 Dawsh 2y ago db_host.allowinsecure = host.allowinsecure
210 e422d72 Dawsh 2y ago db_host.weight = host.weight
211 2de9d72 Dawsh 2y ago db.commit()
212 2de9d72 Dawsh 2y ago db.refresh(db_host)
213 2de9d72 Dawsh 2y ago return db_host
214 a47eb4c Dawsh 2y ago
215 a47eb4c Dawsh 2y ago
216 a47eb4c Dawsh 2y ago def get_user(db: Session, username: str):
217 a47eb4c Dawsh 2y ago return db.query(User).filter(User.username == username).first()
218 a47eb4c Dawsh 2y ago
219 a47eb4c Dawsh 2y ago
220 a47eb4c Dawsh 2y ago def get_user_by_id(db: Session, user_id: int):
221 a47eb4c Dawsh 2y ago return db.query(User).filter(User.id == user_id).first()
222 a47eb4c Dawsh 2y ago
223 a47eb4c Dawsh 2y ago
224 01cdf2e Dawsh 2y ago UsersSortingOptions = Enum(
225 01cdf2e Dawsh 2y ago "UsersSortingOptions",
226 01cdf2e Dawsh 2y ago {
227 01cdf2e Dawsh 2y ago "username": User.username.asc(),
228 01cdf2e Dawsh 2y ago "used_traffic": User.used_traffic.asc(),
229 01cdf2e Dawsh 2y ago "data_limit": User.data_limit.asc(),
230 d377c4c Dawsh 2y ago "expire": User.expire_date.asc(),
231 01cdf2e Dawsh 2y ago "created_at": User.created_at.asc(),
232 01cdf2e Dawsh 2y ago "-username": User.username.desc(),
233 01cdf2e Dawsh 2y ago "-used_traffic": User.used_traffic.desc(),
234 01cdf2e Dawsh 2y ago "-data_limit": User.data_limit.desc(),
235 d377c4c Dawsh 2y ago "-expire": User.expire_date.desc(),
236 01cdf2e Dawsh 2y ago "-created_at": User.created_at.desc(),
237 01cdf2e Dawsh 2y ago },
238 01cdf2e Dawsh 2y ago )
239 01cdf2e Dawsh 2y ago
240 01cdf2e Dawsh 2y ago
241 01cdf2e Dawsh 2y ago def get_users(
242 01cdf2e Dawsh 2y ago db: Session,
243 01cdf2e Dawsh 2y ago offset: Optional[int] = None,
244 01cdf2e Dawsh 2y ago limit: Optional[int] = None,
245 01cdf2e Dawsh 2y ago usernames: Optional[List[str]] = None,
246 01cdf2e Dawsh 2y ago sort: Optional[List[UsersSortingOptions]] = None,
247 01cdf2e Dawsh 2y ago admin: Optional[Admin] = None,
248 d377c4c Dawsh 2y ago reset_strategy: Optional[Union[UserDataUsageResetStrategy, list]] = None,
249 d377c4c Dawsh 2y ago expire_strategy: (
250 d377c4c Dawsh 2y ago UserExpireStrategy | list[UserExpireStrategy] | None
251 d377c4c Dawsh 2y ago ) = None,
252 d377c4c Dawsh 2y ago is_active: bool | None = None,
253 d377c4c Dawsh 2y ago activated: bool | None = None,
254 d377c4c Dawsh 2y ago expired: bool | None = None,
255 d377c4c Dawsh 2y ago data_limit_reached: bool | None = None,
256 4708fcb Dawsh 2y ago enabled: bool | None = None,
257 01cdf2e Dawsh 2y ago ) -> Union[List[User], Tuple[List[User], int]]:
258 24826a2 Dawsh 2y ago query = db.query(User).filter(User.removed == False)
259 a47eb4c Dawsh 2y ago
260 a47eb4c Dawsh 2y ago if usernames:
261 a47eb4c Dawsh 2y ago if len(usernames) == 1:
262 a47eb4c Dawsh 2y ago query = query.filter(User.username.ilike(f"%{usernames[0]}%"))
263 a47eb4c Dawsh 2y ago else:
264 a47eb4c Dawsh 2y ago query = query.filter(User.username.in_(usernames))
265 a47eb4c Dawsh 2y ago
266 a47eb4c Dawsh 2y ago if reset_strategy:
267 a47eb4c Dawsh 2y ago if isinstance(reset_strategy, list):
268 f1ab16c Dawsh 2y ago query = query.filter(
269 f1ab16c Dawsh 2y ago User.data_limit_reset_strategy.in_(reset_strategy)
270 f1ab16c Dawsh 2y ago )
271 a47eb4c Dawsh 2y ago else:
272 f1ab16c Dawsh 2y ago query = query.filter(
273 f1ab16c Dawsh 2y ago User.data_limit_reset_strategy == reset_strategy
274 f1ab16c Dawsh 2y ago )
275 a47eb4c Dawsh 2y ago
276 d377c4c Dawsh 2y ago if expire_strategy:
277 d377c4c Dawsh 2y ago if isinstance(expire_strategy, list):
278 d377c4c Dawsh 2y ago query = query.filter(User.expire_strategy.in_(expire_strategy))
279 d377c4c Dawsh 2y ago else:
280 d377c4c Dawsh 2y ago query = query.filter(User.expire_strategy == expire_strategy)
281 d377c4c Dawsh 2y ago
282 d377c4c Dawsh 2y ago if isinstance(is_active, bool):
283 d377c4c Dawsh 2y ago query = query.filter(User.is_active == is_active)
284 d377c4c Dawsh 2y ago
285 d377c4c Dawsh 2y ago if isinstance(activated, bool):
286 d377c4c Dawsh 2y ago query = query.filter(User.activated == activated)
287 d377c4c Dawsh 2y ago
288 d377c4c Dawsh 2y ago if isinstance(expired, bool):
289 d377c4c Dawsh 2y ago query = query.filter(User.expired == expired)
290 d377c4c Dawsh 2y ago
291 d377c4c Dawsh 2y ago if isinstance(data_limit_reached, bool):
292 d377c4c Dawsh 2y ago query = query.filter(User.data_limit_reached == data_limit_reached)
293 d377c4c Dawsh 2y ago
294 4708fcb Dawsh 2y ago if isinstance(enabled, bool):
295 4708fcb Dawsh 2y ago query = query.filter(User.enabled == enabled)
296 4708fcb Dawsh 2y ago
297 a47eb4c Dawsh 2y ago if admin:
298 a47eb4c Dawsh 2y ago query = query.filter(User.admin == admin)
299 a47eb4c Dawsh 2y ago
300 a47eb4c Dawsh 2y ago if sort:
301 a47eb4c Dawsh 2y ago query = query.order_by(*(opt.value for opt in sort))
302 a47eb4c Dawsh 2y ago
303 a47eb4c Dawsh 2y ago if offset:
304 a47eb4c Dawsh 2y ago query = query.offset(offset)
305 d377c4c Dawsh 2y ago
306 a47eb4c Dawsh 2y ago if limit:
307 a47eb4c Dawsh 2y ago query = query.limit(limit)
308 a47eb4c Dawsh 2y ago
309 a47eb4c Dawsh 2y ago return query.all()
310 a47eb4c Dawsh 2y ago
311 a47eb4c Dawsh 2y ago
312 01cdf2e Dawsh 2y ago def get_user_usages(
313 01cdf2e Dawsh 2y ago db: Session,
314 01cdf2e Dawsh 2y ago dbuser: User,
315 01cdf2e Dawsh 2y ago start: datetime,
316 01cdf2e Dawsh 2y ago end: datetime,
317 01cdf2e Dawsh 2y ago ) -> List[UserUsageResponse]:
318 eb74041 Dawsh 2y ago usages = dict()
319 a47eb4c Dawsh 2y ago
320 a47eb4c Dawsh 2y ago for node in db.query(Node).all():
321 a47eb4c Dawsh 2y ago usages[node.id] = UserUsageResponse(
322 01cdf2e Dawsh 2y ago node_id=node.id, node_name=node.name, used_traffic=0
323 a47eb4c Dawsh 2y ago )
324 a47eb4c Dawsh 2y ago
325 01cdf2e Dawsh 2y ago cond = and_(
326 01cdf2e Dawsh 2y ago NodeUserUsage.user_id == dbuser.id,
327 01cdf2e Dawsh 2y ago NodeUserUsage.created_at >= start,
328 01cdf2e Dawsh 2y ago NodeUserUsage.created_at <= end,
329 01cdf2e Dawsh 2y ago )
330 a47eb4c Dawsh 2y ago
331 a47eb4c Dawsh 2y ago for v in db.query(NodeUserUsage).filter(cond):
332 a47eb4c Dawsh 2y ago try:
333 a47eb4c Dawsh 2y ago usages[v.node_id or 0].used_traffic += v.used_traffic
334 a47eb4c Dawsh 2y ago except KeyError:
335 a47eb4c Dawsh 2y ago pass
336 a47eb4c Dawsh 2y ago
337 a47eb4c Dawsh 2y ago return list(usages.values())
338 a47eb4c Dawsh 2y ago
339 a47eb4c Dawsh 2y ago
340 e5dcd82 Dawsh 2y ago def get_users_count(
341 e5dcd82 Dawsh 2y ago db: Session,
342 e5dcd82 Dawsh 2y ago admin: Admin | None = None,
343 e5dcd82 Dawsh 2y ago enabled: bool | None = None,
344 e5dcd82 Dawsh 2y ago online: bool | None = None,
345 d377c4c Dawsh 2y ago expire_strategy: UserExpireStrategy | None = None,
346 d377c4c Dawsh 2y ago is_active: bool | None = None,
347 d377c4c Dawsh 2y ago expired: bool | None = None,
348 d377c4c Dawsh 2y ago data_limit_reached: bool | None = None,
349 e5dcd82 Dawsh 2y ago ):
350 45803f3 Dawsh 2y ago query = db.query(User.id).filter(User.removed == False)
351 a47eb4c Dawsh 2y ago if admin:
352 9ddfbc8 Mardin 2y ago query = query.filter(User.admin_id == admin.id)
353 d377c4c Dawsh 2y ago if is_active:
354 d377c4c Dawsh 2y ago query = query.filter(User.is_active == is_active)
355 d377c4c Dawsh 2y ago if expired:
356 d377c4c Dawsh 2y ago query = query.filter(User.expired == expired)
357 d377c4c Dawsh 2y ago if data_limit_reached:
358 d377c4c Dawsh 2y ago query = query.filter(User.data_limit_reached == data_limit_reached)
359 e5dcd82 Dawsh 2y ago if enabled:
360 e5dcd82 Dawsh 2y ago query = query.filter(User.enabled == enabled)
361 d377c4c Dawsh 2y ago if online is True:
362 e5dcd82 Dawsh 2y ago query = query.filter(
363 e5dcd82 Dawsh 2y ago User.online_at > (datetime.utcnow() - timedelta(seconds=30))
364 e5dcd82 Dawsh 2y ago )
365 d377c4c Dawsh 2y ago elif online is False:
366 d377c4c Dawsh 2y ago query = query.filter(
367 d377c4c Dawsh 2y ago User.online_at < (datetime.utcnow() - timedelta(seconds=30))
368 d377c4c Dawsh 2y ago )
369 d377c4c Dawsh 2y ago if expire_strategy:
370 d377c4c Dawsh 2y ago query = query.filter(User.expire_strategy == expire_strategy)
371 d377c4c Dawsh 2y ago
372 a47eb4c Dawsh 2y ago return query.count()
373 a47eb4c Dawsh 2y ago
374 a47eb4c Dawsh 2y ago
375 4708fcb Dawsh 2y ago def create_user(
376 4708fcb Dawsh 2y ago db: Session,
377 4708fcb Dawsh 2y ago user: UserCreate,
378 4708fcb Dawsh 2y ago admin: Admin = None,
379 4708fcb Dawsh 2y ago allowed_services: list | None = None,
380 4708fcb Dawsh 2y ago ):
381 4708fcb Dawsh 2y ago service_ids = (
382 4708fcb Dawsh 2y ago [sid for sid in user.service_ids if sid in allowed_services]
383 4708fcb Dawsh 2y ago if allowed_services is not None
384 4708fcb Dawsh 2y ago else user.service_ids
385 4708fcb Dawsh 2y ago )
386 a47eb4c Dawsh 2y ago dbuser = User(
387 a47eb4c Dawsh 2y ago username=user.username,
388 a47eb4c Dawsh 2y ago key=user.key,
389 d377c4c Dawsh 2y ago expire_strategy=user.expire_strategy,
390 d377c4c Dawsh 2y ago expire_date=user.expire_date,
391 d377c4c Dawsh 2y ago usage_duration=user.usage_duration,
392 d377c4c Dawsh 2y ago activation_deadline=user.activation_deadline,
393 01cdf2e Dawsh 2y ago services=db.query(Service)
394 4708fcb Dawsh 2y ago .filter(Service.id.in_(service_ids))
395 01cdf2e Dawsh 2y ago .all(), # user.services,
396 a47eb4c Dawsh 2y ago data_limit=(user.data_limit or None),
397 a47eb4c Dawsh 2y ago admin=admin,
398 a47eb4c Dawsh 2y ago data_limit_reset_strategy=user.data_limit_reset_strategy,
399 a47eb4c Dawsh 2y ago note=user.note,
400 a47eb4c Dawsh 2y ago )
401 a47eb4c Dawsh 2y ago db.add(dbuser)
402 a47eb4c Dawsh 2y ago db.commit()
403 a47eb4c Dawsh 2y ago db.refresh(dbuser)
404 a47eb4c Dawsh 2y ago return dbuser
405 a47eb4c Dawsh 2y ago
406 a47eb4c Dawsh 2y ago
407 a47eb4c Dawsh 2y ago def remove_user(db: Session, dbuser: User):
408 24826a2 Dawsh 2y ago dbuser.removed = True
409 24826a2 Dawsh 2y ago dbuser.activated = False
410 a47eb4c Dawsh 2y ago # db.query(User).filter_by(id=user_id).delete()
411 a47eb4c Dawsh 2y ago db.commit()
412 a47eb4c Dawsh 2y ago
413 a47eb4c Dawsh 2y ago
414 4708fcb Dawsh 2y ago def update_user(
415 4708fcb Dawsh 2y ago db: Session,
416 4708fcb Dawsh 2y ago dbuser: User,
417 4708fcb Dawsh 2y ago modify: UserModify,
418 4708fcb Dawsh 2y ago allowed_services: list | None = None,
419 4708fcb Dawsh 2y ago ):
420 a47eb4c Dawsh 2y ago if modify.data_limit is not None:
421 d377c4c Dawsh 2y ago dbuser.data_limit = modify.data_limit or None
422 d377c4c Dawsh 2y ago
423 d377c4c Dawsh 2y ago if modify.expire_strategy is not None:
424 d377c4c Dawsh 2y ago dbuser.expire_strategy = modify.expire_strategy or None
425 d377c4c Dawsh 2y ago
426 d377c4c Dawsh 2y ago if modify.expire_date is not None:
427 d377c4c Dawsh 2y ago dbuser.expire_date = modify.expire_date or None
428 a47eb4c Dawsh 2y ago
429 a47eb4c Dawsh 2y ago if modify.note is not None:
430 a47eb4c Dawsh 2y ago dbuser.note = modify.note or None
431 a47eb4c Dawsh 2y ago
432 a47eb4c Dawsh 2y ago if modify.data_limit_reset_strategy is not None:
433 d377c4c Dawsh 2y ago dbuser.data_limit_reset_strategy = modify.data_limit_reset_strategy
434 a47eb4c Dawsh 2y ago
435 d377c4c Dawsh 2y ago if modify.activation_deadline is not None:
436 d377c4c Dawsh 2y ago dbuser.activation_deadline = modify.activation_deadline
437 a47eb4c Dawsh 2y ago
438 d377c4c Dawsh 2y ago if modify.usage_duration is not None:
439 d377c4c Dawsh 2y ago dbuser.usage_duration = modify.usage_duration
440 2de9d72 Dawsh 2y ago
441 db68a58 Dawsh 2y ago if modify.service_ids is not None:
442 4708fcb Dawsh 2y ago if allowed_services is not None:
443 4708fcb Dawsh 2y ago service_ids = [
444 4708fcb Dawsh 2y ago sid for sid in modify.service_ids if sid in allowed_services
445 4708fcb Dawsh 2y ago ]
446 4708fcb Dawsh 2y ago else:
447 4708fcb Dawsh 2y ago service_ids = modify.service_ids
448 4708fcb Dawsh 2y ago
449 01cdf2e Dawsh 2y ago dbuser.services = (
450 4708fcb Dawsh 2y ago db.query(Service).filter(Service.id.in_(service_ids)).all()
451 01cdf2e Dawsh 2y ago )
452 a47eb4c Dawsh 2y ago dbuser.edit_at = datetime.utcnow()
453 a47eb4c Dawsh 2y ago
454 a47eb4c Dawsh 2y ago db.commit()
455 a47eb4c Dawsh 2y ago db.refresh(dbuser)
456 a47eb4c Dawsh 2y ago return dbuser
457 a47eb4c Dawsh 2y ago
458 a47eb4c Dawsh 2y ago
459 a47eb4c Dawsh 2y ago def reset_user_data_usage(db: Session, dbuser: User):
460 42a1073 Dawsh 2y ago dbuser.traffic_reset_at = datetime.utcnow()
461 a47eb4c Dawsh 2y ago
462 a47eb4c Dawsh 2y ago dbuser.used_traffic = 0
463 42a1073 Dawsh 2y ago
464 a47eb4c Dawsh 2y ago db.add(dbuser)
465 a47eb4c Dawsh 2y ago
466 a47eb4c Dawsh 2y ago db.commit()
467 a47eb4c Dawsh 2y ago db.refresh(dbuser)
468 a47eb4c Dawsh 2y ago return dbuser
469 a47eb4c Dawsh 2y ago
470 a47eb4c Dawsh 2y ago
471 a47eb4c Dawsh 2y ago def revoke_user_sub(db: Session, dbuser: User):
472 34a7612 Dawsh 2y ago dbuser.key = secrets.token_hex(16)
473 a47eb4c Dawsh 2y ago db.commit()
474 a47eb4c Dawsh 2y ago db.refresh(dbuser)
475 a47eb4c Dawsh 2y ago return dbuser
476 a47eb4c Dawsh 2y ago
477 a47eb4c Dawsh 2y ago
478 a47eb4c Dawsh 2y ago def update_user_sub(db: Session, dbuser: User, user_agent: str):
479 a47eb4c Dawsh 2y ago dbuser.sub_updated_at = datetime.utcnow()
480 a47eb4c Dawsh 2y ago dbuser.sub_last_user_agent = user_agent
481 a47eb4c Dawsh 2y ago
482 a47eb4c Dawsh 2y ago db.commit()
483 a47eb4c Dawsh 2y ago db.refresh(dbuser)
484 a47eb4c Dawsh 2y ago return dbuser
485 a47eb4c Dawsh 2y ago
486 a47eb4c Dawsh 2y ago
487 a47eb4c Dawsh 2y ago def reset_all_users_data_usage(db: Session, admin: Optional[Admin] = None):
488 a47eb4c Dawsh 2y ago query = db.query(User)
489 a47eb4c Dawsh 2y ago
490 a47eb4c Dawsh 2y ago if admin:
491 a47eb4c Dawsh 2y ago query = query.filter(User.admin == admin)
492 a47eb4c Dawsh 2y ago
493 d377c4c Dawsh 2y ago for db_user in query.all():
494 d377c4c Dawsh 2y ago db_user.used_traffic = 0
495 a47eb4c Dawsh 2y ago
496 a47eb4c Dawsh 2y ago db.commit()
497 a47eb4c Dawsh 2y ago
498 a47eb4c Dawsh 2y ago
499 a47eb4c Dawsh 2y ago def update_user_status(db: Session, dbuser: User, status: UserStatus):
500 a47eb4c Dawsh 2y ago dbuser.status = status
501 a47eb4c Dawsh 2y ago db.commit()
502 a47eb4c Dawsh 2y ago db.refresh(dbuser)
503 a47eb4c Dawsh 2y ago return dbuser
504 a47eb4c Dawsh 2y ago
505 a47eb4c Dawsh 2y ago
506 a47eb4c Dawsh 2y ago def set_owner(db: Session, dbuser: User, admin: Admin):
507 a47eb4c Dawsh 2y ago dbuser.admin = admin
508 a47eb4c Dawsh 2y ago db.commit()
509 a47eb4c Dawsh 2y ago db.refresh(dbuser)
510 a47eb4c Dawsh 2y ago return dbuser
511 a47eb4c Dawsh 2y ago
512 a47eb4c Dawsh 2y ago
513 a47eb4c Dawsh 2y ago def get_system_usage(db: Session):
514 a47eb4c Dawsh 2y ago return db.query(System).first()
515 a47eb4c Dawsh 2y ago
516 a47eb4c Dawsh 2y ago
517 a47eb4c Dawsh 2y ago def get_jwt_secret_key(db: Session):
518 a47eb4c Dawsh 2y ago return db.query(JWT).first().secret_key
519 a47eb4c Dawsh 2y ago
520 a47eb4c Dawsh 2y ago
521 a47eb4c Dawsh 2y ago def get_tls_certificate(db: Session):
522 a47eb4c Dawsh 2y ago return db.query(TLS).first()
523 a47eb4c Dawsh 2y ago
524 a47eb4c Dawsh 2y ago
525 4708fcb Dawsh 2y ago def get_admin(db: Session, username: str) -> Admin | None:
526 a47eb4c Dawsh 2y ago return db.query(Admin).filter(Admin.username == username).first()
527 a47eb4c Dawsh 2y ago
528 a47eb4c Dawsh 2y ago
529 a47eb4c Dawsh 2y ago def create_admin(db: Session, admin: AdminCreate):
530 a47eb4c Dawsh 2y ago dbadmin = Admin(
531 a47eb4c Dawsh 2y ago username=admin.username,
532 a47eb4c Dawsh 2y ago hashed_password=admin.hashed_password,
533 01cdf2e Dawsh 2y ago is_sudo=admin.is_sudo,
534 4708fcb Dawsh 2y ago enabled=admin.enabled,
535 4708fcb Dawsh 2y ago all_services_access=admin.all_services_access,
536 4708fcb Dawsh 2y ago modify_users_access=admin.modify_users_access,
537 4708fcb Dawsh 2y ago services=db.query(Service)
538 4708fcb Dawsh 2y ago .filter(Service.id.in_(admin.service_ids))
539 4708fcb Dawsh 2y ago .all(),
540 4708fcb Dawsh 2y ago subscription_url_prefix=admin.subscription_url_prefix,
541 a47eb4c Dawsh 2y ago )
542 a47eb4c Dawsh 2y ago db.add(dbadmin)
543 a47eb4c Dawsh 2y ago db.commit()
544 a47eb4c Dawsh 2y ago db.refresh(dbadmin)
545 a47eb4c Dawsh 2y ago return dbadmin
546 a47eb4c Dawsh 2y ago
547 a47eb4c Dawsh 2y ago
548 4708fcb Dawsh 2y ago def update_admin(
549 4708fcb Dawsh 2y ago db: Session, dbadmin: Admin, modifications: AdminPartialModify
550 4708fcb Dawsh 2y ago ):
551 4708fcb Dawsh 2y ago for attribute in [
552 4708fcb Dawsh 2y ago "is_sudo",
553 4708fcb Dawsh 2y ago "hashed_password",
554 4708fcb Dawsh 2y ago "enabled",
555 4708fcb Dawsh 2y ago "all_services_access",
556 4708fcb Dawsh 2y ago "modify_users_access",
557 4708fcb Dawsh 2y ago "subscription_url_prefix",
558 4708fcb Dawsh 2y ago ]:
559 4708fcb Dawsh 2y ago if not isinstance(getattr(modifications, attribute), NoneType):
560 4708fcb Dawsh 2y ago setattr(dbadmin, attribute, getattr(modifications, attribute))
561 4708fcb Dawsh 2y ago if attribute == "hashed_password":
562 4708fcb Dawsh 2y ago dbadmin.password_reset_at = datetime.utcnow()
563 4708fcb Dawsh 2y ago if isinstance(modifications.service_ids, list):
564 4708fcb Dawsh 2y ago dbadmin.services = (
565 4708fcb Dawsh 2y ago db.query(Service)
566 4708fcb Dawsh 2y ago .filter(Service.id.in_(modifications.service_ids))
567 4708fcb Dawsh 2y ago .all()
568 4708fcb Dawsh 2y ago )
569 a47eb4c Dawsh 2y ago db.commit()
570 a47eb4c Dawsh 2y ago db.refresh(dbadmin)
571 2de9d72 Dawsh 2y ago return dbadmin
572 a47eb4c Dawsh 2y ago
573 a47eb4c Dawsh 2y ago
574 01cdf2e Dawsh 2y ago def partial_update_admin(
575 01cdf2e Dawsh 2y ago db: Session, dbadmin: Admin, modified_admin: AdminPartialModify
576 01cdf2e Dawsh 2y ago ):
577 a47eb4c Dawsh 2y ago if modified_admin.is_sudo is not None:
578 a47eb4c Dawsh 2y ago dbadmin.is_sudo = modified_admin.is_sudo
579 01cdf2e Dawsh 2y ago if (
580 01cdf2e Dawsh 2y ago modified_admin.password is not None
581 01cdf2e Dawsh 2y ago and dbadmin.hashed_password != modified_admin.hashed_password
582 01cdf2e Dawsh 2y ago ):
583 a47eb4c Dawsh 2y ago dbadmin.hashed_password = modified_admin.hashed_password
584 a47eb4c Dawsh 2y ago dbadmin.password_reset_at = datetime.utcnow()
585 a47eb4c Dawsh 2y ago
586 a47eb4c Dawsh 2y ago db.commit()
587 a47eb4c Dawsh 2y ago db.refresh(dbadmin)
588 a47eb4c Dawsh 2y ago return dbadmin
589 a47eb4c Dawsh 2y ago
590 a47eb4c Dawsh 2y ago
591 a47eb4c Dawsh 2y ago def remove_admin(db: Session, dbadmin: Admin):
592 a47eb4c Dawsh 2y ago db.delete(dbadmin)
593 a47eb4c Dawsh 2y ago db.commit()
594 a47eb4c Dawsh 2y ago return dbadmin
595 a47eb4c Dawsh 2y ago
596 a47eb4c Dawsh 2y ago
597 01cdf2e Dawsh 2y ago def get_admins(
598 01cdf2e Dawsh 2y ago db: Session,
599 01cdf2e Dawsh 2y ago offset: Optional[int] = None,
600 01cdf2e Dawsh 2y ago limit: Optional[int] = None,
601 01cdf2e Dawsh 2y ago username: Optional[str] = None,
602 01cdf2e Dawsh 2y ago ):
603 a47eb4c Dawsh 2y ago query = db.query(Admin)
604 a47eb4c Dawsh 2y ago if username:
605 01cdf2e Dawsh 2y ago query = query.filter(Admin.username.ilike(f"%{username}%"))
606 a47eb4c Dawsh 2y ago if offset:
607 a47eb4c Dawsh 2y ago query = query.offset(offset)
608 a47eb4c Dawsh 2y ago if limit:
609 a47eb4c Dawsh 2y ago query = query.limit(limit)
610 a47eb4c Dawsh 2y ago return query.all()
611 a47eb4c Dawsh 2y ago
612 a47eb4c Dawsh 2y ago
613 ee1351a Dawsh 2y ago def create_service(db: Session, service: ServiceCreate) -> Service:
614 a47eb4c Dawsh 2y ago dbservice = Service(
615 a47eb4c Dawsh 2y ago name=service.name,
616 f1ab16c Dawsh 2y ago inbounds=db.query(Inbound)
617 f1ab16c Dawsh 2y ago .filter(Inbound.id.in_(service.inbound_ids))
618 f1ab16c Dawsh 2y ago .all(),
619 db68a58 Dawsh 2y ago users=[],
620 2de9d72 Dawsh 2y ago )
621 a47eb4c Dawsh 2y ago db.add(dbservice)
622 a47eb4c Dawsh 2y ago db.commit()
623 a47eb4c Dawsh 2y ago db.refresh(dbservice)
624 a47eb4c Dawsh 2y ago return dbservice
625 a47eb4c Dawsh 2y ago
626 a47eb4c Dawsh 2y ago
627 ee1351a Dawsh 2y ago def get_service(db: Session, service_id: id) -> Service:
628 ee1351a Dawsh 2y ago return db.query(Service).filter(Service.id == service_id).first()
629 a47eb4c Dawsh 2y ago
630 12b255b Dawsh 2y ago
631 a47eb4c Dawsh 2y ago def get_services(db: Session) -> List[Service]:
632 a47eb4c Dawsh 2y ago return db.query(Service).all()
633 a47eb4c Dawsh 2y ago
634 12b255b Dawsh 2y ago
635 f1ab16c Dawsh 2y ago def update_service(
636 f1ab16c Dawsh 2y ago db: Session, db_service: Service, modification: ServiceModify
637 f1ab16c Dawsh 2y ago ):
638 ee1351a Dawsh 2y ago if modification.name is not None:
639 ee1351a Dawsh 2y ago db_service.name = modification.name
640 ee1351a Dawsh 2y ago
641 db68a58 Dawsh 2y ago if modification.inbound_ids is not None:
642 01cdf2e Dawsh 2y ago db_service.inbounds = (
643 f1ab16c Dawsh 2y ago db.query(Inbound)
644 f1ab16c Dawsh 2y ago .filter(Inbound.id.in_(modification.inbound_ids))
645 f1ab16c Dawsh 2y ago .all()
646 01cdf2e Dawsh 2y ago )
647 ee1351a Dawsh 2y ago
648 ee1351a Dawsh 2y ago db.commit()
649 ee1351a Dawsh 2y ago db.refresh(db_service)
650 ee1351a Dawsh 2y ago return db_service
651 ee1351a Dawsh 2y ago
652 ee1351a Dawsh 2y ago
653 ee1351a Dawsh 2y ago def remove_service(db: Session, db_service: Service):
654 ee1351a Dawsh 2y ago db.delete(db_service)
655 ee1351a Dawsh 2y ago db.commit()
656 ee1351a Dawsh 2y ago return db_service
657 ee1351a Dawsh 2y ago
658 ee1351a Dawsh 2y ago
659 a47eb4c Dawsh 2y ago def get_node(db: Session, name: str):
660 a47eb4c Dawsh 2y ago return db.query(Node).filter(Node.name == name).first()
661 a47eb4c Dawsh 2y ago
662 a47eb4c Dawsh 2y ago
663 a47eb4c Dawsh 2y ago def get_node_by_id(db: Session, node_id: int):
664 a47eb4c Dawsh 2y ago return db.query(Node).filter(Node.id == node_id).first()
665 a47eb4c Dawsh 2y ago
666 a47eb4c Dawsh 2y ago
667 01cdf2e Dawsh 2y ago def get_nodes(
668 f1ab16c Dawsh 2y ago db: Session,
669 f1ab16c Dawsh 2y ago status: Optional[Union[NodeStatus, list]] = None,
670 f1ab16c Dawsh 2y ago enabled: bool = None,
671 01cdf2e Dawsh 2y ago ):
672 a47eb4c Dawsh 2y ago query = db.query(Node)
673 a47eb4c Dawsh 2y ago
674 a47eb4c Dawsh 2y ago if status:
675 a47eb4c Dawsh 2y ago if isinstance(status, list):
676 a47eb4c Dawsh 2y ago query = query.filter(Node.status.in_(status))
677 a47eb4c Dawsh 2y ago else:
678 a47eb4c Dawsh 2y ago query = query.filter(Node.status == status)
679 a47eb4c Dawsh 2y ago
680 a47eb4c Dawsh 2y ago if enabled:
681 a47eb4c Dawsh 2y ago query = query.filter(Node.status != NodeStatus.disabled)
682 a47eb4c Dawsh 2y ago
683 a47eb4c Dawsh 2y ago return query.all()
684 a47eb4c Dawsh 2y ago
685 a47eb4c Dawsh 2y ago
686 01cdf2e Dawsh 2y ago def get_nodes_usage(
687 01cdf2e Dawsh 2y ago db: Session, start: datetime, end: datetime
688 01cdf2e Dawsh 2y ago ) -> List[NodeUsageResponse]:
689 12b255b Dawsh 2y ago usages = dict()
690 a47eb4c Dawsh 2y ago
691 a47eb4c Dawsh 2y ago usages[0] = NodeUsageResponse( # Main Core
692 01cdf2e Dawsh 2y ago node_id=None, node_name="Master", uplink=0, downlink=0
693 a47eb4c Dawsh 2y ago )
694 a47eb4c Dawsh 2y ago for node in db.query(Node).all():
695 a47eb4c Dawsh 2y ago usages[node.id] = NodeUsageResponse(
696 01cdf2e Dawsh 2y ago node_id=node.id, node_name=node.name, uplink=0, downlink=0
697 a47eb4c Dawsh 2y ago )
698 a47eb4c Dawsh 2y ago
699 a47eb4c Dawsh 2y ago cond = and_(NodeUsage.created_at >= start, NodeUsage.created_at <= end)
700 a47eb4c Dawsh 2y ago
701 a47eb4c Dawsh 2y ago for v in db.query(NodeUsage).filter(cond):
702 a47eb4c Dawsh 2y ago try:
703 a47eb4c Dawsh 2y ago usages[v.node_id or 0].uplink += v.uplink
704 a47eb4c Dawsh 2y ago usages[v.node_id or 0].downlink += v.downlink
705 a47eb4c Dawsh 2y ago except KeyError:
706 a47eb4c Dawsh 2y ago pass
707 a47eb4c Dawsh 2y ago
708 a47eb4c Dawsh 2y ago return list(usages.values())
709 a47eb4c Dawsh 2y ago
710 a47eb4c Dawsh 2y ago
711 a47eb4c Dawsh 2y ago def create_node(db: Session, node: NodeCreate):
712 01cdf2e Dawsh 2y ago dbnode = Node(
713 01cdf2e Dawsh 2y ago name=node.name,
714 01cdf2e Dawsh 2y ago address=node.address,
715 01cdf2e Dawsh 2y ago port=node.port,
716 01cdf2e Dawsh 2y ago connection_backend=node.connection_backend,
717 01cdf2e Dawsh 2y ago )
718 a47eb4c Dawsh 2y ago
719 a47eb4c Dawsh 2y ago db.add(dbnode)
720 a47eb4c Dawsh 2y ago db.commit()
721 a47eb4c Dawsh 2y ago db.refresh(dbnode)
722 a47eb4c Dawsh 2y ago return dbnode
723 a47eb4c Dawsh 2y ago
724 a47eb4c Dawsh 2y ago
725 a47eb4c Dawsh 2y ago def remove_node(db: Session, dbnode: Node):
726 a47eb4c Dawsh 2y ago db.delete(dbnode)
727 a47eb4c Dawsh 2y ago db.commit()
728 a47eb4c Dawsh 2y ago return dbnode
729 a47eb4c Dawsh 2y ago
730 a47eb4c Dawsh 2y ago
731 a47eb4c Dawsh 2y ago def update_node(db: Session, dbnode: Node, modify: NodeModify):
732 a47eb4c Dawsh 2y ago if modify.name is not None:
733 a47eb4c Dawsh 2y ago dbnode.name = modify.name
734 a47eb4c Dawsh 2y ago
735 a47eb4c Dawsh 2y ago if modify.address is not None:
736 a47eb4c Dawsh 2y ago dbnode.address = modify.address
737 a47eb4c Dawsh 2y ago
738 a47eb4c Dawsh 2y ago if modify.port is not None:
739 a47eb4c Dawsh 2y ago dbnode.port = modify.port
740 a47eb4c Dawsh 2y ago
741 a47eb4c Dawsh 2y ago if modify.status is NodeStatus.disabled:
742 a47eb4c Dawsh 2y ago dbnode.status = modify.status
743 a47eb4c Dawsh 2y ago dbnode.xray_version = None
744 a47eb4c Dawsh 2y ago dbnode.message = None
745 a47eb4c Dawsh 2y ago else:
746 12b255b Dawsh 2y ago dbnode.status = NodeStatus.unhealthy
747 a47eb4c Dawsh 2y ago
748 a47eb4c Dawsh 2y ago if modify.usage_coefficient:
749 a47eb4c Dawsh 2y ago dbnode.usage_coefficient = modify.usage_coefficient
750 a47eb4c Dawsh 2y ago
751 5584d4a Mrdn 2y ago if modify.connection_backend:
752 5584d4a Mrdn 2y ago dbnode.connection_backend = modify.connection_backend
753 5584d4a Mrdn 2y ago
754 a47eb4c Dawsh 2y ago db.commit()
755 a47eb4c Dawsh 2y ago db.refresh(dbnode)
756 a47eb4c Dawsh 2y ago return dbnode
757 a47eb4c Dawsh 2y ago
758 a47eb4c Dawsh 2y ago
759 01cdf2e Dawsh 2y ago def update_node_status(
760 01cdf2e Dawsh 2y ago db: Session,
761 01cdf2e Dawsh 2y ago node_id: int,
762 01cdf2e Dawsh 2y ago status: NodeStatus,
763 01cdf2e Dawsh 2y ago message: str = None,
764 01cdf2e Dawsh 2y ago version: str = None,
765 01cdf2e Dawsh 2y ago ):
766 f6a32ea Dawsh 2y ago db_node = db.query(Node).where(Node.id == node_id).first()
767 f6a32ea Dawsh 2y ago db_node.status = status
768 f6a32ea Dawsh 2y ago if message:
769 f6a32ea Dawsh 2y ago db_node.message = message
770 f6a32ea Dawsh 2y ago db_node.last_status_change = datetime.utcnow()
771 a47eb4c Dawsh 2y ago db.commit()
772 a47eb4c Dawsh 2y ago
773 a47eb4c Dawsh 2y ago
774 a47eb4c Dawsh 2y ago def create_notification_reminder(
775 f1ab16c Dawsh 2y ago db: Session,
776 f1ab16c Dawsh 2y ago reminder_type: ReminderType,
777 f1ab16c Dawsh 2y ago expires_at: datetime,
778 f1ab16c Dawsh 2y ago user_id: int,
779 01cdf2e Dawsh 2y ago ) -> NotificationReminder:
780 01cdf2e Dawsh 2y ago reminder = NotificationReminder(
781 01cdf2e Dawsh 2y ago type=reminder_type, expires_at=expires_at, user_id=user_id
782 01cdf2e Dawsh 2y ago )
783 a47eb4c Dawsh 2y ago db.add(reminder)
784 a47eb4c Dawsh 2y ago db.commit()
785 a47eb4c Dawsh 2y ago db.refresh(reminder)
786 a47eb4c Dawsh 2y ago return reminder
787 a47eb4c Dawsh 2y ago
788 a47eb4c Dawsh 2y ago
789 a47eb4c Dawsh 2y ago def get_notification_reminder(
790 01cdf2e Dawsh 2y ago db: Session,
791 01cdf2e Dawsh 2y ago user_id: int,
792 01cdf2e Dawsh 2y ago reminder_type: ReminderType,
793 12b255b Dawsh 2y ago ) -> Optional[NotificationReminder]:
794 01cdf2e Dawsh 2y ago reminder = (
795 01cdf2e Dawsh 2y ago db.query(NotificationReminder)
796 01cdf2e Dawsh 2y ago .filter(NotificationReminder.user_id == user_id)
797 01cdf2e Dawsh 2y ago .filter(NotificationReminder.type == reminder_type)
798 01cdf2e Dawsh 2y ago .first()
799 01cdf2e Dawsh 2y ago )
800 a47eb4c Dawsh 2y ago if reminder is None:
801 a47eb4c Dawsh 2y ago return
802 a47eb4c Dawsh 2y ago if reminder.expires_at and reminder.expires_at < datetime.utcnow():
803 a47eb4c Dawsh 2y ago db.delete(reminder)
804 a47eb4c Dawsh 2y ago db.commit()
805 a47eb4c Dawsh 2y ago return
806 a47eb4c Dawsh 2y ago return reminder
807 a47eb4c Dawsh 2y ago
808 a47eb4c Dawsh 2y ago
809 01cdf2e Dawsh 2y ago def delete_notification_reminder_by_type(
810 01cdf2e Dawsh 2y ago db: Session, user_id: int, reminder_type: ReminderType
811 01cdf2e Dawsh 2y ago ) -> None:
812 a47eb4c Dawsh 2y ago """Deletes notification reminder filtered by user_id and type if exists"""
813 a47eb4c Dawsh 2y ago stmt = delete(NotificationReminder).where(
814 a47eb4c Dawsh 2y ago NotificationReminder.user_id == user_id,
815 a47eb4c Dawsh 2y ago NotificationReminder.type == reminder_type,
816 a47eb4c Dawsh 2y ago )
817 a47eb4c Dawsh 2y ago db.execute(stmt)
818 a47eb4c Dawsh 2y ago db.commit()
819 a47eb4c Dawsh 2y ago return
820 a47eb4c Dawsh 2y ago
821 a47eb4c Dawsh 2y ago
822 f1ab16c Dawsh 2y ago def delete_notification_reminder(
823 f1ab16c Dawsh 2y ago db: Session, dbreminder: NotificationReminder
824 f1ab16c Dawsh 2y ago ) -> None:
825 a47eb4c Dawsh 2y ago db.delete(dbreminder)
826 a47eb4c Dawsh 2y ago db.commit()
827 a47eb4c Dawsh 2y ago return