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(
22            db,
23            reset_strategy=[
24                UserDataLimitResetStrategy.day.value,
25                UserDataLimitResetStrategy.week.value,
26                UserDataLimitResetStrategy.month.value,
27                UserDataLimitResetStrategy.year.value,
28            ],
29        ):
30            last_reset_time = user.last_traffic_reset_time
31            num_days_to_reset = reset_strategy_to_days[user.data_limit_reset_strategy]
32
33            if not (now - last_reset_time).days >= num_days_to_reset:
34                continue
35
36            old_status = user.status
37            crud.reset_user_data_usage(db, user)
38            # make user active if limited on usage reset
39            if user.status == UserStatus.active and old_status == UserStatus.limited:
40                await marznode.operations.update_user(user)
41
42            logger.info(f'User data usage reset for User "{user.username}"')