diff --git a/app/db/crud.py b/app/db/crud.py
@@ -352,7 +352,7 @@ def reset_user_data_usage(db: Session, dbuser: User):
dbuser.used_traffic = 0
dbuser.node_usages.clear()
- if dbuser.status not in (UserStatus.expired or UserStatus.disabled):
+ if dbuser.status is UserStatus.limited:
dbuser.status = UserStatus.active.value
db.add(dbuser)
diff --git a/app/jobs/record_usages.py b/app/jobs/record_usages.py
@@ -0,0 +1,122 @@
+import asyncio
+from collections import defaultdict
+from datetime import datetime
+
+from grpclib import GRPCError
+from sqlalchemy import and_, select, insert, update, bindparam
+
+from app import marznode, scheduler
+from app.db import GetDB
+from app.db.models import NodeUsage, NodeUserUsage, User
+from app.marznode import MarzNodeBase
+
+
+def record_user_stats(params: list, node_id: int,
+ consumption_factor: int = 1):
+ if not params:
+ return
+
+ created_at = datetime.fromisoformat(datetime.utcnow().strftime('%Y-%m-%dT%H:00:00'))
+
+ with GetDB() as db:
+ # make user usage row if doesn't exist
+ select_stmt = select(NodeUserUsage.user_id) \
+ .where(and_(NodeUserUsage.node_id == node_id, NodeUserUsage.created_at == created_at))
+ existings = [r[0] for r in db.execute(select_stmt).fetchall()]
+ uids_to_insert = set()
+
+ for p in params:
+ uid = p['uid']
+ if uid in existings:
+ continue
+ uids_to_insert.add(uid)
+
+ if uids_to_insert:
+ stmt = insert(NodeUserUsage).values(
+ user_id=bindparam('uid'),
+ created_at=created_at,
+ node_id=node_id,
+ used_traffic=0
+ )
+ db.execute(stmt, [{'uid': uid} for uid in uids_to_insert])
+
+ # record
+ stmt = update(NodeUserUsage) \
+ .values(used_traffic=NodeUserUsage.used_traffic + bindparam('value') * consumption_factor) \
+ .where(and_(NodeUserUsage.user_id == bindparam('uid'),
+ NodeUserUsage.node_id == node_id,
+ NodeUserUsage.created_at == created_at))
+ db.connection().execute(stmt, params, execution_options={"synchronize_session": None})
+ db.commit()
+
+
+def record_node_stats(params: dict, node_id: int):
+ if not params:
+ return
+
+ created_at = datetime.fromisoformat(datetime.utcnow().strftime('%Y-%m-%dT%H:00:00'))
+
+ with GetDB() as db:
+
+ # make node usage row if doesn't exist
+ select_stmt = select(NodeUsage.node_id). \
+ where(and_(NodeUsage.node_id == node_id, NodeUsage.created_at == created_at))
+ notfound = db.execute(select_stmt).first() is None
+ if notfound:
+ stmt = insert(NodeUsage).values(created_at=created_at, node_id=node_id, uplink=0, downlink=0)
+ db.execute(stmt)
+ db.commit()
+
+ # record
+ stmt = update(NodeUsage). \
+ values(uplink=NodeUsage.uplink + bindparam('up'), downlink=NodeUsage.downlink + bindparam('down')). \
+ where(and_(NodeUsage.node_id == node_id, NodeUsage.created_at == created_at))
+
+ db.execute(stmt, params)
+
+
+async def get_users_stats(node_id: int, node: MarzNodeBase) -> tuple[int, list[dict]]:
+ try:
+ params = list()
+ for stat in await node.fetch_users_stats():
+ if stat.usage:
+ params.append({"uid": stat.uid, "value": stat.usage})
+ return node_id, params
+ except (GRPCError, ConnectionError, OSError):
+ return node_id, []
+
+
+async def record_user_usages():
+ # usage_coefficient = {None: 1} # default usage coefficient for the main api instance
+
+ results = await asyncio.gather(
+ *[get_users_stats(node_id, node) for node_id, node in marznode.nodes.items()]
+ )
+ api_params = {node_id: params for node_id, params in list(results)}
+
+ users_usage = defaultdict(int)
+ for node_id, params in api_params.items():
+ coefficient = node.usage_coefficient if (node := marznode.nodes.get(node_id)) else 1
+ for param in params:
+ users_usage[param['uid']] += param['value'] * coefficient # apply the usage coefficient
+
+ users_usage = list({"id": uid, "value": value} for uid, value in users_usage.items())
+ if not users_usage:
+ return
+
+ # record users usage
+ with GetDB() as db:
+ stmt = update(User)\
+ .values(
+ used_traffic=User.used_traffic + bindparam('value'),
+ online_at=datetime.utcnow()
+ )
+
+ db.execute(stmt, users_usage, execution_options={"synchronize_session": None})
+ db.commit()
+
+ for node_id, params in api_params.items():
+ record_user_stats(params, node_id,
+ node.usage_coefficient if (node := marznode.nodes.get(node_id)) else 1)
+
+scheduler.add_job(record_user_usages, 'interval', coalesce=True, seconds=30)
diff --git a/app/jobs/reset_user_data_usage.py b/app/jobs/reset_user_data_usage.py
@@ -1,6 +1,6 @@
from datetime import datetime
-from app import logger, scheduler
+from app import logger, scheduler, marznode
from app.db import crud, GetDB, get_users
from app.models.user import UserDataLimitResetStrategy, UserStatus
@@ -12,7 +12,7 @@ reset_strategy_to_days = {
}
-def reset_user_data_usage():
+async def reset_user_data_usage():
now = datetime.utcnow()
with GetDB() as db:
for user in get_users(db, reset_strategy=[
@@ -27,10 +27,11 @@ def reset_user_data_usage():
if not (now - last_reset_time).days >= num_days_to_reset:
continue
+ old_status = user.status
crud.reset_user_data_usage(db, user)
# make user active if limited on usage reset
- if user.status == UserStatus.limited:
- xray.operations.add_user(user)
+ if user.status == UserStatus.active and old_status == UserStatus.limited:
+ await marznode.operations.add_user(user)
logger.info(f"User data usage reset for User \"{user.username}\"")
diff --git a/app/jobs/review_users.py b/app/jobs/review_users.py
@@ -1,14 +1,14 @@
+import asyncio
from datetime import datetime
from typing import TYPE_CHECKING
from sqlalchemy.orm import Session
-from app import logger, scheduler
+from app import logger, scheduler, marznode
from app.db import (GetDB, get_notification_reminder, get_users,
start_user_expire, update_user_status)
from app.models.user import ReminderType, UserResponse, UserStatus
from app.utils import report
-# from app.utils.concurrency import GetBG
from app.utils.helpers import (calculate_expiration_days,
calculate_usage_percent)
from config import (NOTIFY_DAYS_LEFT, NOTIFY_REACHED_USAGE_PERCENT,
@@ -37,7 +37,7 @@ def add_notification_reminders(db: Session, user: "User", now: datetime = dateti
async def review():
now = datetime.utcnow()
- with GetDB() as db: #, GetBG() as bg:
+ with GetDB() as db:
for user in get_users(db, status=UserStatus.active):
limited = user.data_limit and user.used_traffic >= user.data_limit
@@ -51,12 +51,13 @@ async def review():
add_notification_reminders(db, user, now)
continue
- # await xray.operations.remove_user(user)
+ await marznode.operations.remove_user(user)
update_user_status(db, user, status)
- #bg.add_task(
- # report.status_change, user.username, status, UserResponse.model_validate(user)
- #)
+ asyncio.create_task(report.status_change(
+ user.username,
+ user.status,
+ UserResponse.model_validate(user)))
logger.info(f"User \"{user.username}\" status changed to {status}")
@@ -80,22 +81,11 @@ async def review():
update_user_status(db, user, status)
start_user_expire(db, user)
- #bg.add_task(report.status_change, user.username, status,
- # UserResponse.model_validate(user))
-
- logger.info(f"User \"{user.username}\" status changed to {status}")
-
-
- active = user.expire and user.expire >= now.timestamp()
- if active:
- status = UserStatus.active
- else:
- continue
-
- update_user_status(db, user, status)
- # await xray.operations.add_user(user)
-
- logger.info(f"User \"{user.username}\" status fixed.")
+ asyncio.create_task(report.status_change(
+ user.username,
+ user.status,
+ UserResponse.model_validate(user)))
+ logger.info(f"on hold user `{user.username}` has been activated")
-scheduler.add_job(review, 'interval', seconds=10, coalesce=True, max_instances=1)
+scheduler.add_job(review, 'interval', seconds=600, coalesce=True, max_instances=1)
diff --git a/app/marznode/grpc.py b/app/marznode/grpc.py
@@ -21,7 +21,7 @@ def string_to_temp_file(content: str):
class MarzNodeGRPC(MarzNodeBase):
- def __init__(self, address: str, port: int, ssl_key: str, ssl_cert: str): # TODO: implement ssl
+ def __init__(self, address: str, port: int, ssl_key: str, ssl_cert: str, usage_coefficient: int = 1):
self._address = address
self._port = port
@@ -36,6 +36,8 @@ class MarzNodeGRPC(MarzNodeBase):
self._stub = MarzServiceStub(self._channel)
self._health = HealthStub(self._channel)
self.synced = False
+
+ self.usage_coefficient = usage_coefficient
atexit.register(self._channel.close)
@staticmethod
@@ -89,7 +91,10 @@ class MarzNodeGRPC(MarzNodeBase):
await stream.send_message(UsersUpdate(users_updates=updates))
async def fetch_users_stats(self):
- raise NotImplementedError
+ async with self._stub.FetchUsersStats.open() as stream:
+ await stream.send_message(Empty())
+ stats = await stream.recv_message()
+ return stats.users_stats
async def fetch_inbounds(self) -> list:
async with self._stub.FetchInbounds.open() as stm:
diff --git a/app/marznode/marznode.proto b/app/marznode/marznode.proto
@@ -14,7 +14,7 @@ message InboundsResponse {
}
message User {
- int32 id = 1;
+ uint32 id = 1;
string username = 2;
string key = 3;
}
@@ -31,9 +31,8 @@ message UsersUpdate {
message UsersStats {
message UserStats {
- User user = 1;
- uint64 uplink = 2;
- uint64 downlink = 3;
+ uint32 uid = 1;
+ uint64 usage = 2;
}
repeated UserStats users_stats = 1;
}
diff --git a/app/marznode/marznode_pb2.py b/app/marznode/marznode_pb2.py
@@ -14,7 +14,7 @@ _sym_db = _symbol_database.Default()
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1b\x61pp/marznode/marznode.proto\x12\x08marznode\"\x07\n\x05\x45mpty\"6\n\x07Inbound\x12\x0b\n\x03tag\x18\x01 \x01(\t\x12\x13\n\x06\x63onfig\x18\x02 \x01(\tH\x00\x88\x01\x01\x42\t\n\x07_config\"7\n\x10InboundsResponse\x12#\n\x08inbounds\x18\x01 \x03(\x0b\x32\x11.marznode.Inbound\"1\n\x04User\x12\n\n\x02id\x18\x01 \x01(\x05\x12\x10\n\x08username\x18\x02 \x01(\t\x12\x0b\n\x03key\x18\x03 \x01(\t\"\x87\x01\n\nUserUpdate\x12\x1c\n\x04user\x18\x01 \x01(\x0b\x32\x0e.marznode.User\x12,\n\x11inbound_additions\x18\x02 \x03(\x0b\x32\x11.marznode.Inbound\x12-\n\x12inbound_reductions\x18\x03 \x03(\x0b\x32\x11.marznode.Inbound\":\n\x0bUsersUpdate\x12+\n\rusers_updates\x18\x01 \x03(\x0b\x32\x14.marznode.UserUpdate\"\x8e\x01\n\nUsersStats\x12\x33\n\x0busers_stats\x18\x01 \x03(\x0b\x32\x1e.marznode.UsersStats.UserStats\x1aK\n\tUserStats\x12\x1c\n\x04user\x18\x01 \x01(\x0b\x32\x0e.marznode.User\x12\x0e\n\x06uplink\x18\x02 \x01(\x04\x12\x10\n\x08\x64ownlink\x18\x03 \x01(\x04\x32\xe4\x02\n\x0bMarzService\x12\x30\n\x07\x41\x64\x64User\x12\x14.marznode.UserUpdate\x1a\x0f.marznode.Empty\x12\x33\n\nRemoveUser\x12\x14.marznode.UserUpdate\x1a\x0f.marznode.Empty\x12;\n\x12UpdateUserInbounds\x12\x14.marznode.UserUpdate\x1a\x0f.marznode.Empty\x12\x39\n\x0fRepopulateUsers\x12\x15.marznode.UsersUpdate\x1a\x0f.marznode.Empty\x12<\n\rFetchInbounds\x12\x0f.marznode.Empty\x1a\x1a.marznode.InboundsResponse\x12\x38\n\x0f\x46\x65tchUsersStats\x12\x0f.marznode.Empty\x1a\x14.marznode.UsersStatsb\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1b\x61pp/marznode/marznode.proto\x12\x08marznode\"\x07\n\x05\x45mpty\"6\n\x07Inbound\x12\x0b\n\x03tag\x18\x01 \x01(\t\x12\x13\n\x06\x63onfig\x18\x02 \x01(\tH\x00\x88\x01\x01\x42\t\n\x07_config\"7\n\x10InboundsResponse\x12#\n\x08inbounds\x18\x01 \x03(\x0b\x32\x11.marznode.Inbound\"1\n\x04User\x12\n\n\x02id\x18\x01 \x01(\r\x12\x10\n\x08username\x18\x02 \x01(\t\x12\x0b\n\x03key\x18\x03 \x01(\t\"\x87\x01\n\nUserUpdate\x12\x1c\n\x04user\x18\x01 \x01(\x0b\x32\x0e.marznode.User\x12,\n\x11inbound_additions\x18\x02 \x03(\x0b\x32\x11.marznode.Inbound\x12-\n\x12inbound_reductions\x18\x03 \x03(\x0b\x32\x11.marznode.Inbound\":\n\x0bUsersUpdate\x12+\n\rusers_updates\x18\x01 \x03(\x0b\x32\x14.marznode.UserUpdate\"j\n\nUsersStats\x12\x33\n\x0busers_stats\x18\x01 \x03(\x0b\x32\x1e.marznode.UsersStats.UserStats\x1a\'\n\tUserStats\x12\x0b\n\x03uid\x18\x01 \x01(\r\x12\r\n\x05usage\x18\x02 \x01(\x04\x32\xe4\x02\n\x0bMarzService\x12\x30\n\x07\x41\x64\x64User\x12\x14.marznode.UserUpdate\x1a\x0f.marznode.Empty\x12\x33\n\nRemoveUser\x12\x14.marznode.UserUpdate\x1a\x0f.marznode.Empty\x12;\n\x12UpdateUserInbounds\x12\x14.marznode.UserUpdate\x1a\x0f.marznode.Empty\x12\x39\n\x0fRepopulateUsers\x12\x15.marznode.UsersUpdate\x1a\x0f.marznode.Empty\x12<\n\rFetchInbounds\x12\x0f.marznode.Empty\x1a\x1a.marznode.InboundsResponse\x12\x38\n\x0f\x46\x65tchUsersStats\x12\x0f.marznode.Empty\x1a\x14.marznode.UsersStatsb\x06proto3')
_globals = globals()
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -33,10 +33,10 @@ if _descriptor._USE_C_DESCRIPTORS == False:
_globals['_USERUPDATE']._serialized_end=350
_globals['_USERSUPDATE']._serialized_start=352
_globals['_USERSUPDATE']._serialized_end=410
- _globals['_USERSSTATS']._serialized_start=413
- _globals['_USERSSTATS']._serialized_end=555
- _globals['_USERSSTATS_USERSTATS']._serialized_start=480
- _globals['_USERSSTATS_USERSTATS']._serialized_end=555
- _globals['_MARZSERVICE']._serialized_start=558
- _globals['_MARZSERVICE']._serialized_end=914
+ _globals['_USERSSTATS']._serialized_start=412
+ _globals['_USERSSTATS']._serialized_end=518
+ _globals['_USERSSTATS_USERSTATS']._serialized_start=479
+ _globals['_USERSSTATS_USERSTATS']._serialized_end=518
+ _globals['_MARZSERVICE']._serialized_start=521
+ _globals['_MARZSERVICE']._serialized_end=877
# @@protoc_insertion_point(module_scope)
diff --git a/app/utils/notification.py b/app/utils/notification.py
@@ -100,6 +100,6 @@ class UserSubscriptionRevoked(UserNotification):
user: UserResponse
-async def notify(message: Type[Notification]) -> None:
+async def notify(message: Notification) -> None:
if config.WEBHOOK_ADDRESS:
queue.append(message)
diff --git a/app/views/user.py b/app/views/user.py
@@ -149,25 +149,27 @@ async def reset_user_data_usage(username: str,
"""
Reset user data usage
"""
- if not (admin.is_sudo or (dbuser.admin and dbuser.admin.username == admin.username)):
+ db_user = crud.get_user(db, username)
+ if not (admin.is_sudo or (db_user.admin and db_user.admin.username == admin.username)):
raise HTTPException(status_code=403, detail="You're not allowed")
- dbuser = crud.get_user(db, username)
- if not dbuser:
+ if not db_user:
raise HTTPException(status_code=404, detail="User not found")
- dbuser = crud.reset_user_data_usage(db=db, dbuser=dbuser)
- if dbuser.status in [UserStatus.active, UserStatus.on_hold]:
- asyncio.create_task(marznode.operations.add_user(dbuser=dbuser))
+ previous_status = db_user.status
+ db_user = crud.reset_user_data_usage(db=db, dbuser=db_user)
- user = UserResponse.model_validate(dbuser)
+ if db_user.status == UserStatus.active and previous_status == UserStatus.limited:
+ await marznode.operations.add_user(db_user)
+
+ user = UserResponse.model_validate(db_user)
asyncio.create_task(report.user_data_usage_reset(
user=user,
by=admin))
logger.info(f"User \"{username}\"'s usage was reset")
- return dbuser
+ return user
@app.post("/api/user/{username}/revoke_sub", tags=['User'], response_model=UserResponse)