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