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