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