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