diff --git a/app/db/crud.py b/app/db/crud.py index c321c62..bfdd276 100644 --- 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 new file mode 100644 index 0000000..7420493 --- /dev/null +++ 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 index 9bfe7b6..3ce2c25 100644 --- 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 index 77fd45e..2748997 100644 --- 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 index 393de69..0324a8d 100644 --- 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 index 6c31727..815133b 100644 --- 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 index 30843be..f67c7a0 100644 --- 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 index 48e973c..327e2ea 100644 --- 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 index de9cea8..993ec89 100644 --- 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)