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