gitscope
← marzneshin / v0.1.2 / app/tasks/reset_user_data_usage.py
https://dev.kronoz.win/r/marzneshin.git
app/tasks/reset_user_data_usage.py · Python
Blame Raw
 1import logging
 2from datetime import datetime
 3
 4from app import marznode
 5from app.db import crud, GetDB, get_users
 6from app.models.user import UserDataLimitResetStrategy, UserStatus
 7
 8logger = logging.getLogger(__name__)
 9
10reset_strategy_to_days = {
11    UserDataLimitResetStrategy.day.value: 1,
12    UserDataLimitResetStrategy.week.value: 7,
13    UserDataLimitResetStrategy.month.value: 30,
14    UserDataLimitResetStrategy.year.value: 365,
15}
16
17
18async def reset_user_data_usage():
19    now = datetime.utcnow()
20    with GetDB() as db:
21        for user in get_users(db, reset_strategy=[
22            UserDataLimitResetStrategy.day.value,
23            UserDataLimitResetStrategy.week.value,
24            UserDataLimitResetStrategy.month.value,
25            UserDataLimitResetStrategy.year.value,
26        ]):
27            last_reset_time = user.last_traffic_reset_time
28            num_days_to_reset = reset_strategy_to_days[user.data_limit_reset_strategy]
29
30            if not (now - last_reset_time).days >= num_days_to_reset:
31                continue
32
33            old_status = user.status
34            crud.reset_user_data_usage(db, user)
35            # make user active if limited on usage reset
36            if user.status == UserStatus.active and old_status == UserStatus.limited:
37                await marznode.operations.update_user(user)
38
39            logger.info(f"User data usage reset for User \"{user.username}\"")