gitscope
← marzneshin / master / 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    with db.get_bind().begin() as conn:
 729        conn.execute(
 730            update(User)
 731            .where(User.id == dbuser.id)
 732            .values(
 733                sub_updated_at=datetime.utcnow(),
 734                sub_last_user_agent=user_agent,
 735            )
 736        )
 737
 738
 739def reset_all_users_data_usage(db: Session, admin: Optional[Admin] = None):
 740    query = db.query(User)
 741
 742    if admin:
 743        query = query.filter(User.admin == admin)
 744
 745    for db_user in query.all():
 746        db_user.used_traffic = 0
 747
 748    db.commit()
 749
 750
 751def update_user_status(db: Session, dbuser: User, status: UserStatus):
 752    dbuser.status = status
 753    db.commit()
 754    db.refresh(dbuser)
 755    return dbuser
 756
 757
 758def set_owner(db: Session, dbuser: User, admin: Admin):
 759    dbuser.admin = admin
 760    db.commit()
 761    db.refresh(dbuser)
 762    return dbuser
 763
 764
 765def get_system_usage(db: Session):
 766    return db.query(System).first()
 767
 768
 769def get_jwt_secret_key(db: Session):
 770    return db.query(JWT).first().secret_key
 771
 772
 773def get_tls_certificate(db: Session):
 774    return db.query(TLS).first()
 775
 776
 777def get_admin(db: Session, username: str) -> Admin | None:
 778    return db.query(Admin).filter(Admin.username == username).first()
 779
 780
 781def create_admin(db: Session, admin: AdminCreate):
 782    dbadmin = Admin(
 783        username=admin.username,
 784        hashed_password=admin.hashed_password,
 785        is_sudo=admin.is_sudo,
 786        enabled=admin.enabled,
 787        all_services_access=admin.all_services_access,
 788        modify_users_access=admin.modify_users_access,
 789        services=db.query(Service)
 790        .filter(Service.id.in_(admin.service_ids))
 791        .all(),
 792        subscription_url_prefix=admin.subscription_url_prefix,
 793    )
 794    db.add(dbadmin)
 795    db.commit()
 796    db.refresh(dbadmin)
 797    return dbadmin
 798
 799
 800def update_admin(
 801    db: Session, dbadmin: Admin, modifications: AdminPartialModify
 802):
 803    for attribute in [
 804        "is_sudo",
 805        "hashed_password",
 806        "enabled",
 807        "all_services_access",
 808        "modify_users_access",
 809        "subscription_url_prefix",
 810    ]:
 811        if not isinstance(getattr(modifications, attribute), NoneType):
 812            setattr(dbadmin, attribute, getattr(modifications, attribute))
 813            if attribute == "hashed_password":
 814                dbadmin.password_reset_at = datetime.utcnow()
 815    if isinstance(modifications.service_ids, list):
 816        dbadmin.services = (
 817            db.query(Service)
 818            .filter(Service.id.in_(modifications.service_ids))
 819            .all()
 820        )
 821    db.commit()
 822    db.refresh(dbadmin)
 823    return dbadmin
 824
 825
 826def partial_update_admin(
 827    db: Session, dbadmin: Admin, modified_admin: AdminPartialModify
 828):
 829    if modified_admin.is_sudo is not None:
 830        dbadmin.is_sudo = modified_admin.is_sudo
 831    if (
 832        modified_admin.password is not None
 833        and dbadmin.hashed_password != modified_admin.hashed_password
 834    ):
 835        dbadmin.hashed_password = modified_admin.hashed_password
 836        dbadmin.password_reset_at = datetime.utcnow()
 837
 838    db.commit()
 839    db.refresh(dbadmin)
 840    return dbadmin
 841
 842
 843def remove_admin(db: Session, dbadmin: Admin):
 844    db.delete(dbadmin)
 845    db.commit()
 846    return dbadmin
 847
 848
 849def get_admins(
 850    db: Session,
 851    offset: Optional[int] = None,
 852    limit: Optional[int] = None,
 853    username: Optional[str] = None,
 854):
 855    query = db.query(Admin)
 856    if username:
 857        query = query.filter(Admin.username.ilike(f"%{username}%"))
 858    if offset:
 859        query = query.offset(offset)
 860    if limit:
 861        query = query.limit(limit)
 862    return query.all()
 863
 864
 865def create_service(db: Session, service: ServiceCreate) -> Service:
 866    dbservice = Service(
 867        name=service.name,
 868        inbounds=db.query(Inbound)
 869        .filter(Inbound.id.in_(service.inbound_ids))
 870        .all(),
 871        users=[],
 872    )
 873    db.add(dbservice)
 874    db.commit()
 875    db.refresh(dbservice)
 876    return dbservice
 877
 878
 879def get_service(db: Session, service_id: id) -> Service:
 880    return db.query(Service).filter(Service.id == service_id).first()
 881
 882
 883def get_services(db: Session) -> List[Service]:
 884    return db.query(Service).all()
 885
 886
 887def update_service(
 888    db: Session, db_service: Service, modification: ServiceModify
 889):
 890    if modification.name is not None:
 891        db_service.name = modification.name
 892
 893    if modification.inbound_ids is not None:
 894        db_service.inbounds = (
 895            db.query(Inbound)
 896            .filter(Inbound.id.in_(modification.inbound_ids))
 897            .all()
 898        )
 899
 900    db.commit()
 901    db.refresh(db_service)
 902    return db_service
 903
 904
 905def remove_service(db: Session, db_service: Service):
 906    db.delete(db_service)
 907    db.commit()
 908    return db_service
 909
 910
 911def get_node(db: Session, name: str):
 912    return db.query(Node).filter(Node.name == name).first()
 913
 914
 915def get_node_by_id(db: Session, node_id: int):
 916    return db.query(Node).filter(Node.id == node_id).first()
 917
 918
 919def get_nodes(
 920    db: Session,
 921    status: Optional[Union[NodeStatus, list]] = None,
 922    enabled: bool = None,
 923):
 924    query = db.query(Node)
 925
 926    if status:
 927        if isinstance(status, list):
 928            query = query.filter(Node.status.in_(status))
 929        else:
 930            query = query.filter(Node.status == status)
 931
 932    if enabled:
 933        query = query.filter(Node.status != NodeStatus.disabled)
 934
 935    return query.all()
 936
 937
 938def get_node_usage(
 939    db: Session, start: datetime, end: datetime, node: Node
 940) -> TrafficUsageSeries:
 941    usages = defaultdict(int)
 942
 943    query = (
 944        db.query(
 945            NodeUserUsage.created_at, func.sum(NodeUserUsage.used_traffic)
 946        )
 947        .group_by(NodeUserUsage.created_at)
 948        .filter(
 949            and_(
 950                NodeUserUsage.node_id == node.id,
 951                NodeUserUsage.created_at >= start,
 952                NodeUserUsage.created_at <= end,
 953            )
 954        )
 955    )
 956
 957    for created_at, used_traffic in query.all():
 958        usages[created_at.replace(tzinfo=timezone.utc).timestamp()] += int(
 959            used_traffic
 960        )
 961
 962    result = TrafficUsageSeries(usages=[], total=0)
 963    current = start.astimezone(timezone.utc).replace(
 964        minute=0, second=0, microsecond=0
 965    )
 966
 967    while current <= end.replace(tzinfo=timezone.utc):
 968        usage = usages.get(current.timestamp()) or 0
 969        result.usages.append((int(current.timestamp()), usage))
 970        result.total += usage
 971        current += timedelta(hours=1)
 972
 973    return result
 974
 975
 976def create_node(db: Session, node: NodeCreate):
 977    dbnode = Node(
 978        name=node.name,
 979        address=node.address,
 980        port=node.port,
 981        connection_backend=node.connection_backend,
 982    )
 983
 984    db.add(dbnode)
 985    db.commit()
 986    db.refresh(dbnode)
 987    return dbnode
 988
 989
 990def remove_node(db: Session, dbnode: Node):
 991    db.delete(dbnode)
 992    db.commit()
 993    return dbnode
 994
 995
 996def update_node(db: Session, dbnode: Node, modify: NodeModify):
 997    if modify.name is not None:
 998        dbnode.name = modify.name
 999
1000    if modify.address is not None:
1001        dbnode.address = modify.address
1002
1003    if modify.port is not None:
1004        dbnode.port = modify.port
1005
1006    if modify.status is NodeStatus.disabled:
1007        dbnode.status = modify.status
1008        dbnode.xray_version = None
1009        dbnode.message = None
1010    else:
1011        dbnode.status = NodeStatus.unhealthy
1012
1013    if modify.usage_coefficient:
1014        dbnode.usage_coefficient = modify.usage_coefficient
1015
1016    if modify.connection_backend:
1017        dbnode.connection_backend = modify.connection_backend
1018
1019    db.commit()
1020    db.refresh(dbnode)
1021    return dbnode
1022
1023
1024def update_node_status(
1025    db: Session,
1026    node_id: int,
1027    status: NodeStatus,
1028    message: str = None,
1029    version: str = None,
1030):
1031    db_node = db.query(Node).where(Node.id == node_id).first()
1032    db_node.status = status
1033    if message:
1034        db_node.message = message
1035    db_node.last_status_change = datetime.utcnow()
1036    db.commit()