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