1import json
2import secrets
3from collections import defaultdict
4from datetime import datetime, timedelta, timezone
5from enum import Enum
6from types import NoneType
7from typing import List, Optional, Tuple, Union
8
9from sqlalchemy import and_, delete, update, select, func, cast, Date
10from sqlalchemy.orm import Session
11
12from app.db.models import (
13 JWT,
14 TLS,
15 Admin,
16 Node,
17 NodeUserUsage,
18 InboundHost,
19 Service,
20 Inbound,
21 System,
22 User,
23 Backend,
24)
25from app.models.admin import AdminCreate, AdminPartialModify
26from app.models.node import (
27 NodeCreate,
28 NodeModify,
29 NodeStatus,
30)
31from app.models.proxy import InboundHost as InboundHostModify
32from app.models.service import Service as ServiceModify, ServiceCreate
33from app.models.system import TrafficUsageSeries
34from app.models.user import (
35 UserCreate,
36 UserDataUsageResetStrategy,
37 UserModify,
38 UserStatus,
39 UserExpireStrategy,
40 UserNodeUsageSeries,
41 UserUsageSeriesResponse,
42)
43
44
45def add_default_hosts(db: Session, inbounds: List[Inbound]):
46 hosts = [
47 InboundHost(
48 remark="🚀 Marz ({USERNAME}) [{PROTOCOL} - {TRANSPORT}]",
49 address="{SERVER_IP}",
50 inbound=i,
51 )
52 for i in inbounds
53 ]
54 db.add_all(hosts)
55 db.commit()
56
57
58def ensure_node_backends(db: Session, backends, node_id: int):
59 old_backends = db.query(Backend).where(Backend.node_id == node_id)
60 for backend in old_backends:
61 db.delete(backend)
62 backends = [
63 Backend(
64 name=backend.name,
65 backend_type=backend.type,
66 version=backend.version,
67 node_id=node_id,
68 )
69 for backend in backends
70 ]
71 db.add_all(backends)
72 db.flush()
73
74
75def ensure_node_inbounds(db: Session, inbounds: List[Inbound], node_id: int):
76 current_tags = [
77 i[0]
78 for i in db.execute(
79 select(Inbound.tag).filter(Inbound.node_id == node_id)
80 ).all()
81 ]
82 updated_tags = set(i.tag for i in list(inbounds))
83 inbound_additions, tag_deletions = list(), set()
84 for tag in current_tags:
85 if tag not in updated_tags:
86 tag_deletions.add(tag)
87 removals = db.query(Inbound).where(
88 and_(Inbound.node_id == node_id, Inbound.tag.in_(tag_deletions))
89 )
90 for i in removals:
91 db.delete(i)
92
93 for inb in inbounds:
94 if inb.tag in current_tags:
95 stmt = (
96 update(Inbound)
97 .where(
98 and_(Inbound.node_id == node_id, Inbound.tag == inb.tag)
99 )
100 .values(
101 protocol=json.loads(inb.config)["protocol"],
102 config=inb.config,
103 )
104 )
105 db.execute(stmt)
106 else:
107 inbound_additions.append(inb)
108 new_inbounds = [
109 Inbound(
110 tag=inb.tag,
111 protocol=json.loads(inb.config)["protocol"],
112 config=inb.config,
113 node_id=node_id,
114 )
115 for inb in inbound_additions
116 ]
117 db.add_all(new_inbounds)
118 db.flush()
119 for i in new_inbounds:
120 db.refresh(i)
121 add_default_hosts(db, new_inbounds)
122 db.commit()
123
124
125def get_node_users(
126 db: Session,
127 node_id: int,
128):
129 query = (
130 db.query(User.id, User.username, User.key, Inbound)
131 .distinct()
132 .join(Inbound.services)
133 .join(Service.users)
134 .filter(Inbound.node_id == node_id)
135 .filter(User.activated == True)
136 )
137 return query.all()
138
139
140def get_user_hosts(db: Session, user_id: int):
141 return (
142 db.query(InboundHost)
143 .distinct()
144 .join(User.services)
145 .join(Service.inbounds)
146 .join(Inbound.hosts)
147 .filter(User.id == user_id)
148 .all()
149 )
150
151
152def get_inbound_hosts(db: Session, inbound_id: int) -> List[InboundHost]:
153 return (
154 db.query(InboundHost)
155 .filter(InboundHost.inbound_id == inbound_id)
156 .all()
157 )
158
159
160def get_all_inbounds(db: Session):
161 return db.query(Inbound).all()
162
163
164def get_inbound(db: Session, inbound_id: int) -> Inbound | None:
165 return db.query(Inbound).filter(Inbound.id == inbound_id).first()
166
167
168def get_host(db: Session, host_id) -> InboundHost:
169 return db.query(InboundHost).filter(InboundHost.id == host_id).first()
170
171
172def add_host(db: Session, inbound: Inbound, host: InboundHostModify):
173 host = InboundHost(
174 remark=host.remark,
175 address=host.address,
176 port=host.port,
177 path=host.path,
178 sni=host.sni,
179 host=host.host,
180 inbound=inbound,
181 security=host.security,
182 alpn=host.alpn.value,
183 fingerprint=host.fingerprint,
184 fragment=host.fragment.model_dump() if host.fragment else None,
185 udp_noises=host.udp_noises,
186 http_headers=host.http_headers,
187 mtu=host.mtu,
188 dns_servers=host.dns_servers,
189 allowed_ips=host.allowed_ips,
190 mux=host.mux,
191 allowinsecure=host.allowinsecure,
192 weight=host.weight,
193 )
194 inbound.hosts.append(host)
195 db.commit()
196 db.refresh(host)
197 return host
198
199
200def update_host(db: Session, db_host: InboundHost, host: InboundHostModify):
201 db_host.remark = host.remark
202 db_host.address = host.address
203 db_host.port = host.port
204 db_host.path = host.path
205 db_host.sni = host.sni
206 db_host.host = host.host
207 db_host.security = host.security
208 db_host.alpn = host.alpn.value
209 db_host.fingerprint = host.fingerprint
210 db_host.fragment = host.fragment.model_dump() if host.fragment else None
211 db_host.mux = host.mux
212 db_host.is_disabled = host.is_disabled
213 db_host.allowinsecure = host.allowinsecure
214 db_host.udp_noises = host.udp_noises
215 db_host.http_headers = host.http_headers
216 db_host.mtu = host.mtu
217 db_host.dns_servers = host.dns_servers
218 db_host.allowed_ips = host.allowed_ips
219 db_host.weight = host.weight
220 db.commit()
221 db.refresh(db_host)
222 return db_host
223
224
225def get_user(db: Session, username: str):
226 return db.query(User).filter(User.username == username).first()
227
228
229def get_user_by_id(db: Session, user_id: int):
230 return db.query(User).filter(User.id == user_id).first()
231
232
233UsersSortingOptions = Enum(
234 "UsersSortingOptions",
235 {
236 "username": User.username.asc(),
237 "used_traffic": User.used_traffic.asc(),
238 "data_limit": User.data_limit.asc(),
239 "expire": User.expire_date.asc(),
240 "created_at": User.created_at.asc(),
241 "-username": User.username.desc(),
242 "-used_traffic": User.used_traffic.desc(),
243 "-data_limit": User.data_limit.desc(),
244 "-expire": User.expire_date.desc(),
245 "-created_at": User.created_at.desc(),
246 },
247)
248
249
250def get_users(
251 db: Session,
252 offset: Optional[int] = None,
253 limit: Optional[int] = None,
254 usernames: Optional[List[str]] = None,
255 sort: Optional[List[UsersSortingOptions]] = None,
256 admin: Optional[Admin] = None,
257 reset_strategy: Optional[Union[UserDataUsageResetStrategy, list]] = None,
258 expire_strategy: (
259 UserExpireStrategy | list[UserExpireStrategy] | None
260 ) = None,
261 is_active: bool | None = None,
262 activated: bool | None = None,
263 expired: bool | None = None,
264 data_limit_reached: bool | None = None,
265 enabled: bool | None = None,
266) -> Union[List[User], Tuple[List[User], int]]:
267 query = db.query(User).filter(User.removed == False)
268
269 if usernames:
270 if len(usernames) == 1:
271 query = query.filter(User.username.ilike(f"%{usernames[0]}%"))
272 else:
273 query = query.filter(User.username.in_(usernames))
274
275 if reset_strategy:
276 if isinstance(reset_strategy, list):
277 query = query.filter(
278 User.data_limit_reset_strategy.in_(reset_strategy)
279 )
280 else:
281 query = query.filter(
282 User.data_limit_reset_strategy == reset_strategy
283 )
284
285 if expire_strategy:
286 if isinstance(expire_strategy, list):
287 query = query.filter(User.expire_strategy.in_(expire_strategy))
288 else:
289 query = query.filter(User.expire_strategy == expire_strategy)
290
291 if isinstance(is_active, bool):
292 query = query.filter(User.is_active == is_active)
293
294 if isinstance(activated, bool):
295 query = query.filter(User.activated == activated)
296
297 if isinstance(expired, bool):
298 query = query.filter(User.expired == expired)
299
300 if isinstance(data_limit_reached, bool):
301 query = query.filter(User.data_limit_reached == data_limit_reached)
302
303 if isinstance(enabled, bool):
304 query = query.filter(User.enabled == enabled)
305
306 if admin:
307 query = query.filter(User.admin == admin)
308
309 if sort:
310 query = query.order_by(*(opt.value for opt in sort))
311
312 if offset:
313 query = query.offset(offset)
314
315 if limit:
316 query = query.limit(limit)
317
318 return query.all()
319
320
321def get_user_total_usage(
322 db: Session, user: User, start: datetime, end: datetime, per_day=False
323):
324 usages = defaultdict(int)
325
326 query = db.query(
327 (
328 cast(NodeUserUsage.created_at, Date).label("day")
329 if per_day
330 else NodeUserUsage.created_at
331 ),
332 func.sum(NodeUserUsage.used_traffic),
333 ).filter(
334 and_(
335 NodeUserUsage.user_id == user.id,
336 NodeUserUsage.created_at >= start,
337 NodeUserUsage.created_at <= end,
338 )
339 )
340 if per_day:
341 query = query.group_by(cast(NodeUserUsage.created_at, Date))
342 else:
343 query = query.group_by(NodeUserUsage.created_at)
344
345 for date, used_traffic in query:
346 if per_day:
347 timestamp = datetime(
348 date.year, date.month, date.day, tzinfo=timezone.utc
349 ).timestamp()
350 usages[timestamp] += int(used_traffic)
351 else:
352 usages[date.replace(tzinfo=timezone.utc).timestamp()] += int(
353 used_traffic
354 )
355
356 result = TrafficUsageSeries(usages=[])
357 current = start.astimezone(timezone.utc).replace(
358 minute=0, second=0, microsecond=0
359 )
360 if per_day:
361 current = current.replace(hour=0)
362
363 step = timedelta(days=1) if per_day else timedelta(hours=1)
364 while current <= end.replace(tzinfo=timezone.utc):
365 current_usage = usages.get(current.timestamp()) or 0
366 result.usages.append((int(current.timestamp()), current_usage))
367 result.total += current_usage
368 current += step
369
370 result.step = int(step.total_seconds())
371 return result
372
373
374def get_total_usages(
375 db: Session, admin: Admin, start: datetime, end: datetime
376) -> TrafficUsageSeries:
377 usages = defaultdict(int)
378
379 query = (
380 db.query(
381 NodeUserUsage.created_at, func.sum(NodeUserUsage.used_traffic)
382 )
383 .group_by(NodeUserUsage.created_at)
384 .filter(
385 and_(
386 NodeUserUsage.created_at >= start,
387 NodeUserUsage.created_at <= end,
388 )
389 )
390 )
391
392 if not admin.is_sudo:
393 query = (
394 query.filter(
395 Admin.id == admin.id,
396 )
397 .join(User, NodeUserUsage.user_id == User.id)
398 .join(Admin, User.admin_id == Admin.id)
399 )
400
401 for created_at, used_traffic in query.all():
402 timestamp = created_at.replace(tzinfo=timezone.utc).timestamp()
403 usages[timestamp] += int(used_traffic)
404
405 result = TrafficUsageSeries(usages=[], total=0)
406 current = start.astimezone(timezone.utc).replace(
407 minute=0, second=0, microsecond=0
408 )
409
410 while current <= end.replace(tzinfo=timezone.utc):
411 usage = usages.get(current.timestamp()) or 0
412 result.usages.append((int(current.timestamp()), usage))
413 result.total += usage
414 current += timedelta(hours=1)
415
416 return result
417
418
419def get_user_usages(
420 db: Session,
421 db_user: User,
422 start: datetime,
423 end: datetime,
424) -> UserUsageSeriesResponse:
425 usages = defaultdict(dict)
426
427 cond = and_(
428 NodeUserUsage.user_id == db_user.id,
429 NodeUserUsage.created_at >= start,
430 NodeUserUsage.created_at <= end,
431 )
432
433 for v in db.query(NodeUserUsage).filter(cond):
434 timestamp = v.created_at.replace(tzinfo=timezone.utc).timestamp()
435 usages[v.node_id][timestamp] = v.used_traffic
436
437 node_ids = list(usages.keys())
438 nodes = db.query(Node).where(Node.id.in_(node_ids))
439 node_id_names = {node.id: node.name for node in nodes}
440
441 result = UserUsageSeriesResponse(
442 username=db_user.username, node_usages=[], total=0
443 )
444
445 for node_id, rows in usages.items():
446 node_usages = UserNodeUsageSeries(
447 node_id=node_id, node_name=node_id_names[node_id], usages=[]
448 )
449 current = start.astimezone(timezone.utc).replace(
450 minute=0, second=0, microsecond=0
451 )
452
453 while current <= end:
454 usage = rows.get(current.timestamp()) or 0
455 node_usages.usages.append((int(current.timestamp()), usage))
456 current += timedelta(hours=1)
457 result.total += usage
458 result.node_usages.append(node_usages)
459
460 return result
461
462
463def get_users_count(
464 db: Session,
465 admin: Admin | None = None,
466 enabled: bool | None = None,
467 online: bool | None = None,
468 expire_strategy: UserExpireStrategy | None = None,
469 is_active: bool | None = None,
470 expired: bool | None = None,
471 data_limit_reached: bool | None = None,
472):
473 query = db.query(User.id).filter(User.removed == False)
474 if admin:
475 query = query.filter(User.admin_id == admin.id)
476 if is_active:
477 query = query.filter(User.is_active == is_active)
478 if expired:
479 query = query.filter(User.expired == expired)
480 if data_limit_reached:
481 query = query.filter(User.data_limit_reached == data_limit_reached)
482 if enabled:
483 query = query.filter(User.enabled == enabled)
484 if online is True:
485 query = query.filter(
486 User.online_at > (datetime.utcnow() - timedelta(seconds=30))
487 )
488 elif online is False:
489 query = query.filter(
490 User.online_at < (datetime.utcnow() - timedelta(seconds=30))
491 )
492 if expire_strategy:
493 query = query.filter(User.expire_strategy == expire_strategy)
494
495 return query.count()
496
497
498def create_user(
499 db: Session,
500 user: UserCreate,
501 admin: Admin = None,
502 allowed_services: list | None = None,
503):
504 service_ids = (
505 [sid for sid in user.service_ids if sid in allowed_services]
506 if allowed_services is not None
507 else user.service_ids
508 )
509 dbuser = User(
510 username=user.username,
511 key=user.key,
512 expire_strategy=user.expire_strategy,
513 expire_date=user.expire_date,
514 usage_duration=user.usage_duration,
515 activation_deadline=user.activation_deadline,
516 services=db.query(Service)
517 .filter(Service.id.in_(service_ids))
518 .all(), # user.services,
519 data_limit=(user.data_limit or None),
520 admin=admin,
521 data_limit_reset_strategy=user.data_limit_reset_strategy,
522 note=user.note,
523 )
524 db.add(dbuser)
525 db.commit()
526 db.refresh(dbuser)
527 return dbuser
528
529
530def remove_user(db: Session, dbuser: User):
531 dbuser.username = None
532 dbuser.removed = True
533 dbuser.activated = False
534 # db.query(User).filter_by(id=user_id).delete()
535 db.commit()
536
537
538def update_user(
539 db: Session,
540 dbuser: User,
541 modify: UserModify,
542 allowed_services: list | None = None,
543):
544 if modify.data_limit is not None:
545 dbuser.data_limit = modify.data_limit or None
546
547 if modify.expire_strategy is not None:
548 dbuser.expire_strategy = modify.expire_strategy or None
549
550 if modify.expire_strategy == UserExpireStrategy.FIXED_DATE:
551 dbuser.usage_duration = None
552 dbuser.activation_deadline = None
553 elif modify.expire_strategy == UserExpireStrategy.START_ON_FIRST_USE:
554 dbuser.expire_date = None
555 elif modify.expire_strategy == UserExpireStrategy.NEVER:
556 dbuser.expire_date = None
557 dbuser.usage_duration = None
558 dbuser.activation_deadline = None
559
560 if modify.expire_date is not None:
561 dbuser.expire_date = modify.expire_date or None
562
563 if modify.note is not None:
564 dbuser.note = modify.note or None
565
566 if modify.data_limit_reset_strategy is not None:
567 dbuser.data_limit_reset_strategy = modify.data_limit_reset_strategy
568
569 if modify.activation_deadline is not None:
570 dbuser.activation_deadline = modify.activation_deadline
571
572 if modify.usage_duration is not None:
573 dbuser.usage_duration = modify.usage_duration
574
575 if modify.service_ids is not None:
576 if allowed_services is not None:
577 service_ids = [
578 sid for sid in modify.service_ids if sid in allowed_services
579 ]
580 else:
581 service_ids = modify.service_ids
582
583 dbuser.services = (
584 db.query(Service).filter(Service.id.in_(service_ids)).all()
585 )
586 dbuser.edit_at = datetime.utcnow()
587
588 db.commit()
589 db.refresh(dbuser)
590 return dbuser
591
592
593def reset_user_data_usage(db: Session, dbuser: User):
594 dbuser.traffic_reset_at = datetime.utcnow()
595
596 dbuser.used_traffic = 0
597
598 db.add(dbuser)
599
600 db.commit()
601 db.refresh(dbuser)
602 return dbuser
603
604
605def revoke_user_sub(db: Session, dbuser: User):
606 dbuser.key = secrets.token_hex(16)
607 dbuser.sub_revoked_at = datetime.utcnow()
608 db.commit()
609 db.refresh(dbuser)
610 return dbuser
611
612
613def update_user_sub(db: Session, dbuser: User, user_agent: str):
614 dbuser.sub_updated_at = datetime.utcnow()
615 dbuser.sub_last_user_agent = user_agent
616
617 db.commit()
618 db.refresh(dbuser)
619 return dbuser
620
621
622def reset_all_users_data_usage(db: Session, admin: Optional[Admin] = None):
623 query = db.query(User)
624
625 if admin:
626 query = query.filter(User.admin == admin)
627
628 for db_user in query.all():
629 db_user.used_traffic = 0
630
631 db.commit()
632
633
634def update_user_status(db: Session, dbuser: User, status: UserStatus):
635 dbuser.status = status
636 db.commit()
637 db.refresh(dbuser)
638 return dbuser
639
640
641def set_owner(db: Session, dbuser: User, admin: Admin):
642 dbuser.admin = admin
643 db.commit()
644 db.refresh(dbuser)
645 return dbuser
646
647
648def get_system_usage(db: Session):
649 return db.query(System).first()
650
651
652def get_jwt_secret_key(db: Session):
653 return db.query(JWT).first().secret_key
654
655
656def get_tls_certificate(db: Session):
657 return db.query(TLS).first()
658
659
660def get_admin(db: Session, username: str) -> Admin | None:
661 return db.query(Admin).filter(Admin.username == username).first()
662
663
664def create_admin(db: Session, admin: AdminCreate):
665 dbadmin = Admin(
666 username=admin.username,
667 hashed_password=admin.hashed_password,
668 is_sudo=admin.is_sudo,
669 enabled=admin.enabled,
670 all_services_access=admin.all_services_access,
671 modify_users_access=admin.modify_users_access,
672 services=db.query(Service)
673 .filter(Service.id.in_(admin.service_ids))
674 .all(),
675 subscription_url_prefix=admin.subscription_url_prefix,
676 )
677 db.add(dbadmin)
678 db.commit()
679 db.refresh(dbadmin)
680 return dbadmin
681
682
683def update_admin(
684 db: Session, dbadmin: Admin, modifications: AdminPartialModify
685):
686 for attribute in [
687 "is_sudo",
688 "hashed_password",
689 "enabled",
690 "all_services_access",
691 "modify_users_access",
692 "subscription_url_prefix",
693 ]:
694 if not isinstance(getattr(modifications, attribute), NoneType):
695 setattr(dbadmin, attribute, getattr(modifications, attribute))
696 if attribute == "hashed_password":
697 dbadmin.password_reset_at = datetime.utcnow()
698 if isinstance(modifications.service_ids, list):
699 dbadmin.services = (
700 db.query(Service)
701 .filter(Service.id.in_(modifications.service_ids))
702 .all()
703 )
704 db.commit()
705 db.refresh(dbadmin)
706 return dbadmin
707
708
709def partial_update_admin(
710 db: Session, dbadmin: Admin, modified_admin: AdminPartialModify
711):
712 if modified_admin.is_sudo is not None:
713 dbadmin.is_sudo = modified_admin.is_sudo
714 if (
715 modified_admin.password is not None
716 and dbadmin.hashed_password != modified_admin.hashed_password
717 ):
718 dbadmin.hashed_password = modified_admin.hashed_password
719 dbadmin.password_reset_at = datetime.utcnow()
720
721 db.commit()
722 db.refresh(dbadmin)
723 return dbadmin
724
725
726def remove_admin(db: Session, dbadmin: Admin):
727 db.delete(dbadmin)
728 db.commit()
729 return dbadmin
730
731
732def get_admins(
733 db: Session,
734 offset: Optional[int] = None,
735 limit: Optional[int] = None,
736 username: Optional[str] = None,
737):
738 query = db.query(Admin)
739 if username:
740 query = query.filter(Admin.username.ilike(f"%{username}%"))
741 if offset:
742 query = query.offset(offset)
743 if limit:
744 query = query.limit(limit)
745 return query.all()
746
747
748def create_service(db: Session, service: ServiceCreate) -> Service:
749 dbservice = Service(
750 name=service.name,
751 inbounds=db.query(Inbound)
752 .filter(Inbound.id.in_(service.inbound_ids))
753 .all(),
754 users=[],
755 )
756 db.add(dbservice)
757 db.commit()
758 db.refresh(dbservice)
759 return dbservice
760
761
762def get_service(db: Session, service_id: id) -> Service:
763 return db.query(Service).filter(Service.id == service_id).first()
764
765
766def get_services(db: Session) -> List[Service]:
767 return db.query(Service).all()
768
769
770def update_service(
771 db: Session, db_service: Service, modification: ServiceModify
772):
773 if modification.name is not None:
774 db_service.name = modification.name
775
776 if modification.inbound_ids is not None:
777 db_service.inbounds = (
778 db.query(Inbound)
779 .filter(Inbound.id.in_(modification.inbound_ids))
780 .all()
781 )
782
783 db.commit()
784 db.refresh(db_service)
785 return db_service
786
787
788def remove_service(db: Session, db_service: Service):
789 db.delete(db_service)
790 db.commit()
791 return db_service
792
793
794def get_node(db: Session, name: str):
795 return db.query(Node).filter(Node.name == name).first()
796
797
798def get_node_by_id(db: Session, node_id: int):
799 return db.query(Node).filter(Node.id == node_id).first()
800
801
802def get_nodes(
803 db: Session,
804 status: Optional[Union[NodeStatus, list]] = None,
805 enabled: bool = None,
806):
807 query = db.query(Node)
808
809 if status:
810 if isinstance(status, list):
811 query = query.filter(Node.status.in_(status))
812 else:
813 query = query.filter(Node.status == status)
814
815 if enabled:
816 query = query.filter(Node.status != NodeStatus.disabled)
817
818 return query.all()
819
820
821def get_node_usage(
822 db: Session, start: datetime, end: datetime, node: Node
823) -> TrafficUsageSeries:
824 usages = defaultdict(int)
825
826 query = (
827 db.query(
828 NodeUserUsage.created_at, func.sum(NodeUserUsage.used_traffic)
829 )
830 .group_by(NodeUserUsage.created_at)
831 .filter(
832 and_(
833 NodeUserUsage.node_id == node.id,
834 NodeUserUsage.created_at >= start,
835 NodeUserUsage.created_at <= end,
836 )
837 )
838 )
839
840 for created_at, used_traffic in query.all():
841 usages[created_at.replace(tzinfo=timezone.utc).timestamp()] += int(
842 used_traffic
843 )
844
845 result = TrafficUsageSeries(usages=[], total=0)
846 current = start.astimezone(timezone.utc).replace(
847 minute=0, second=0, microsecond=0
848 )
849
850 while current <= end.replace(tzinfo=timezone.utc):
851 usage = usages.get(current.timestamp()) or 0
852 result.usages.append((int(current.timestamp()), usage))
853 result.total += usage
854 current += timedelta(hours=1)
855
856 return result
857
858
859def create_node(db: Session, node: NodeCreate):
860 dbnode = Node(
861 name=node.name,
862 address=node.address,
863 port=node.port,
864 connection_backend=node.connection_backend,
865 )
866
867 db.add(dbnode)
868 db.commit()
869 db.refresh(dbnode)
870 return dbnode
871
872
873def remove_node(db: Session, dbnode: Node):
874 db.delete(dbnode)
875 db.commit()
876 return dbnode
877
878
879def update_node(db: Session, dbnode: Node, modify: NodeModify):
880 if modify.name is not None:
881 dbnode.name = modify.name
882
883 if modify.address is not None:
884 dbnode.address = modify.address
885
886 if modify.port is not None:
887 dbnode.port = modify.port
888
889 if modify.status is NodeStatus.disabled:
890 dbnode.status = modify.status
891 dbnode.xray_version = None
892 dbnode.message = None
893 else:
894 dbnode.status = NodeStatus.unhealthy
895
896 if modify.usage_coefficient:
897 dbnode.usage_coefficient = modify.usage_coefficient
898
899 if modify.connection_backend:
900 dbnode.connection_backend = modify.connection_backend
901
902 db.commit()
903 db.refresh(dbnode)
904 return dbnode
905
906
907def update_node_status(
908 db: Session,
909 node_id: int,
910 status: NodeStatus,
911 message: str = None,
912 version: str = None,
913):
914 db_node = db.query(Node).where(Node.id == node_id).first()
915 db_node.status = status
916 if message:
917 db_node.message = message
918 db_node.last_status_change = datetime.utcnow()
919 db.commit()