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