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