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.lifetime_used_traffic += dbuser.used_traffic
461 dbuser.traffic_reset_at = datetime.utcnow()
462
463 dbuser.used_traffic = 0
464
465 db.add(dbuser)
466
467 db.commit()
468 db.refresh(dbuser)
469 return dbuser
470
471
472def revoke_user_sub(db: Session, dbuser: User):
473 dbuser.key = secrets.token_hex(16)
474 db.commit()
475 db.refresh(dbuser)
476 return dbuser
477
478
479def update_user_sub(db: Session, dbuser: User, user_agent: str):
480 dbuser.sub_updated_at = datetime.utcnow()
481 dbuser.sub_last_user_agent = user_agent
482
483 db.commit()
484 db.refresh(dbuser)
485 return dbuser
486
487
488def reset_all_users_data_usage(db: Session, admin: Optional[Admin] = None):
489 query = db.query(User)
490
491 if admin:
492 query = query.filter(User.admin == admin)
493
494 for db_user in query.all():
495 db_user.used_traffic = 0
496
497 db.commit()
498
499
500def update_user_status(db: Session, dbuser: User, status: UserStatus):
501 dbuser.status = status
502 db.commit()
503 db.refresh(dbuser)
504 return dbuser
505
506
507def set_owner(db: Session, dbuser: User, admin: Admin):
508 dbuser.admin = admin
509 db.commit()
510 db.refresh(dbuser)
511 return dbuser
512
513
514def get_system_usage(db: Session):
515 return db.query(System).first()
516
517
518def get_jwt_secret_key(db: Session):
519 return db.query(JWT).first().secret_key
520
521
522def get_tls_certificate(db: Session):
523 return db.query(TLS).first()
524
525
526def get_admin(db: Session, username: str) -> Admin | None:
527 return db.query(Admin).filter(Admin.username == username).first()
528
529
530def create_admin(db: Session, admin: AdminCreate):
531 dbadmin = Admin(
532 username=admin.username,
533 hashed_password=admin.hashed_password,
534 is_sudo=admin.is_sudo,
535 enabled=admin.enabled,
536 all_services_access=admin.all_services_access,
537 modify_users_access=admin.modify_users_access,
538 services=db.query(Service)
539 .filter(Service.id.in_(admin.service_ids))
540 .all(),
541 subscription_url_prefix=admin.subscription_url_prefix,
542 )
543 db.add(dbadmin)
544 db.commit()
545 db.refresh(dbadmin)
546 return dbadmin
547
548
549def update_admin(
550 db: Session, dbadmin: Admin, modifications: AdminPartialModify
551):
552 for attribute in [
553 "is_sudo",
554 "hashed_password",
555 "enabled",
556 "all_services_access",
557 "modify_users_access",
558 "subscription_url_prefix",
559 ]:
560 if not isinstance(getattr(modifications, attribute), NoneType):
561 setattr(dbadmin, attribute, getattr(modifications, attribute))
562 if attribute == "hashed_password":
563 dbadmin.password_reset_at = datetime.utcnow()
564 if isinstance(modifications.service_ids, list):
565 dbadmin.services = (
566 db.query(Service)
567 .filter(Service.id.in_(modifications.service_ids))
568 .all()
569 )
570 db.commit()
571 db.refresh(dbadmin)
572 return dbadmin
573
574
575def partial_update_admin(
576 db: Session, dbadmin: Admin, modified_admin: AdminPartialModify
577):
578 if modified_admin.is_sudo is not None:
579 dbadmin.is_sudo = modified_admin.is_sudo
580 if (
581 modified_admin.password is not None
582 and dbadmin.hashed_password != modified_admin.hashed_password
583 ):
584 dbadmin.hashed_password = modified_admin.hashed_password
585 dbadmin.password_reset_at = datetime.utcnow()
586
587 db.commit()
588 db.refresh(dbadmin)
589 return dbadmin
590
591
592def remove_admin(db: Session, dbadmin: Admin):
593 db.delete(dbadmin)
594 db.commit()
595 return dbadmin
596
597
598def get_admins(
599 db: Session,
600 offset: Optional[int] = None,
601 limit: Optional[int] = None,
602 username: Optional[str] = None,
603):
604 query = db.query(Admin)
605 if username:
606 query = query.filter(Admin.username.ilike(f"%{username}%"))
607 if offset:
608 query = query.offset(offset)
609 if limit:
610 query = query.limit(limit)
611 return query.all()
612
613
614def create_service(db: Session, service: ServiceCreate) -> Service:
615 dbservice = Service(
616 name=service.name,
617 inbounds=db.query(Inbound)
618 .filter(Inbound.id.in_(service.inbound_ids))
619 .all(),
620 users=[],
621 )
622 db.add(dbservice)
623 db.commit()
624 db.refresh(dbservice)
625 return dbservice
626
627
628def get_service(db: Session, service_id: id) -> Service:
629 return db.query(Service).filter(Service.id == service_id).first()
630
631
632def get_services(db: Session) -> List[Service]:
633 return db.query(Service).all()
634
635
636def update_service(
637 db: Session, db_service: Service, modification: ServiceModify
638):
639 if modification.name is not None:
640 db_service.name = modification.name
641
642 if modification.inbound_ids is not None:
643 db_service.inbounds = (
644 db.query(Inbound)
645 .filter(Inbound.id.in_(modification.inbound_ids))
646 .all()
647 )
648
649 db.commit()
650 db.refresh(db_service)
651 return db_service
652
653
654def remove_service(db: Session, db_service: Service):
655 db.delete(db_service)
656 db.commit()
657 return db_service
658
659
660def get_node(db: Session, name: str):
661 return db.query(Node).filter(Node.name == name).first()
662
663
664def get_node_by_id(db: Session, node_id: int):
665 return db.query(Node).filter(Node.id == node_id).first()
666
667
668def get_nodes(
669 db: Session,
670 status: Optional[Union[NodeStatus, list]] = None,
671 enabled: bool = None,
672):
673 query = db.query(Node)
674
675 if status:
676 if isinstance(status, list):
677 query = query.filter(Node.status.in_(status))
678 else:
679 query = query.filter(Node.status == status)
680
681 if enabled:
682 query = query.filter(Node.status != NodeStatus.disabled)
683
684 return query.all()
685
686
687def get_nodes_usage(
688 db: Session, start: datetime, end: datetime
689) -> List[NodeUsageResponse]:
690 usages = dict()
691
692 usages[0] = NodeUsageResponse( # Main Core
693 node_id=None, node_name="Master", uplink=0, downlink=0
694 )
695 for node in db.query(Node).all():
696 usages[node.id] = NodeUsageResponse(
697 node_id=node.id, node_name=node.name, uplink=0, downlink=0
698 )
699
700 cond = and_(NodeUsage.created_at >= start, NodeUsage.created_at <= end)
701
702 for v in db.query(NodeUsage).filter(cond):
703 try:
704 usages[v.node_id or 0].uplink += v.uplink
705 usages[v.node_id or 0].downlink += v.downlink
706 except KeyError:
707 pass
708
709 return list(usages.values())
710
711
712def create_node(db: Session, node: NodeCreate):
713 dbnode = Node(
714 name=node.name,
715 address=node.address,
716 port=node.port,
717 connection_backend=node.connection_backend,
718 )
719
720 db.add(dbnode)
721 db.commit()
722 db.refresh(dbnode)
723 return dbnode
724
725
726def remove_node(db: Session, dbnode: Node):
727 db.delete(dbnode)
728 db.commit()
729 return dbnode
730
731
732def update_node(db: Session, dbnode: Node, modify: NodeModify):
733 if modify.name is not None:
734 dbnode.name = modify.name
735
736 if modify.address is not None:
737 dbnode.address = modify.address
738
739 if modify.port is not None:
740 dbnode.port = modify.port
741
742 if modify.status is NodeStatus.disabled:
743 dbnode.status = modify.status
744 dbnode.xray_version = None
745 dbnode.message = None
746 else:
747 dbnode.status = NodeStatus.unhealthy
748
749 if modify.usage_coefficient:
750 dbnode.usage_coefficient = modify.usage_coefficient
751
752 if modify.connection_backend:
753 dbnode.connection_backend = modify.connection_backend
754
755 db.commit()
756 db.refresh(dbnode)
757 return dbnode
758
759
760def update_node_status(
761 db: Session,
762 node_id: int,
763 status: NodeStatus,
764 message: str = None,
765 version: str = None,
766):
767 db_node = db.query(Node).where(Node.id == node_id).first()
768 db_node.status = status
769 if message:
770 db_node.message = message
771 db_node.last_status_change = datetime.utcnow()
772 db.commit()
773
774
775def create_notification_reminder(
776 db: Session,
777 reminder_type: ReminderType,
778 expires_at: datetime,
779 user_id: int,
780) -> NotificationReminder:
781 reminder = NotificationReminder(
782 type=reminder_type, expires_at=expires_at, user_id=user_id
783 )
784 db.add(reminder)
785 db.commit()
786 db.refresh(reminder)
787 return reminder
788
789
790def get_notification_reminder(
791 db: Session,
792 user_id: int,
793 reminder_type: ReminderType,
794) -> Optional[NotificationReminder]:
795 reminder = (
796 db.query(NotificationReminder)
797 .filter(NotificationReminder.user_id == user_id)
798 .filter(NotificationReminder.type == reminder_type)
799 .first()
800 )
801 if reminder is None:
802 return
803 if reminder.expires_at and reminder.expires_at < datetime.utcnow():
804 db.delete(reminder)
805 db.commit()
806 return
807 return reminder
808
809
810def delete_notification_reminder_by_type(
811 db: Session, user_id: int, reminder_type: ReminderType
812) -> None:
813 """Deletes notification reminder filtered by user_id and type if exists"""
814 stmt = delete(NotificationReminder).where(
815 NotificationReminder.user_id == user_id,
816 NotificationReminder.type == reminder_type,
817 )
818 db.execute(stmt)
819 db.commit()
820 return
821
822
823def delete_notification_reminder(
824 db: Session, dbreminder: NotificationReminder
825) -> None:
826 db.delete(dbreminder)
827 db.commit()
828 return