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