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