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