981 lines
19 KiB
Python
981 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
import subprocess
|
|
import threading
|
|
|
|
from .config import (
|
|
DAY_NAMES,
|
|
find_user,
|
|
get_day_name,
|
|
load_state,
|
|
load_users_config,
|
|
normalize_access_windows,
|
|
normalize_allowances,
|
|
parse_time,
|
|
python_weekday_to_name,
|
|
save_state,
|
|
)
|
|
from .users import (
|
|
is_locked,
|
|
lock_user,
|
|
terminate_user,
|
|
unlock_user,
|
|
)
|
|
|
|
|
|
_state_lock = threading.RLock()
|
|
|
|
|
|
def user_has_session(
|
|
username: str,
|
|
) -> bool:
|
|
result = subprocess.run(
|
|
[
|
|
"loginctl",
|
|
"list-users",
|
|
"--no-legend",
|
|
],
|
|
capture_output=True,
|
|
text=True,
|
|
check=False,
|
|
)
|
|
|
|
if result.returncode != 0:
|
|
return False
|
|
|
|
for line in result.stdout.splitlines():
|
|
parts = line.split()
|
|
|
|
if len(parts) >= 2:
|
|
if parts[1] == username:
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
def current_time():
|
|
now = datetime.now()
|
|
|
|
weekday = python_weekday_to_name(
|
|
now.weekday()
|
|
)
|
|
|
|
minute = (
|
|
now.hour * 60
|
|
+ now.minute
|
|
)
|
|
|
|
return (
|
|
now,
|
|
weekday,
|
|
minute,
|
|
)
|
|
|
|
|
|
def _today() -> str:
|
|
return (
|
|
datetime.now()
|
|
.date()
|
|
.isoformat()
|
|
)
|
|
|
|
|
|
def _get_usage(
|
|
state: dict,
|
|
username: str,
|
|
date: str,
|
|
) -> int:
|
|
usage = state.setdefault(
|
|
"usage",
|
|
{},
|
|
)
|
|
|
|
user_usage = usage.get(
|
|
username,
|
|
{},
|
|
)
|
|
|
|
if not isinstance(
|
|
user_usage,
|
|
dict,
|
|
):
|
|
return 0
|
|
|
|
try:
|
|
return max(
|
|
0,
|
|
int(
|
|
user_usage.get(
|
|
date,
|
|
0,
|
|
)
|
|
),
|
|
)
|
|
except (
|
|
TypeError,
|
|
ValueError,
|
|
):
|
|
return 0
|
|
|
|
|
|
def _set_usage(
|
|
state: dict,
|
|
username: str,
|
|
date: str,
|
|
seconds: int,
|
|
) -> None:
|
|
usage = state.setdefault(
|
|
"usage",
|
|
{},
|
|
)
|
|
|
|
user_usage = usage.setdefault(
|
|
username,
|
|
{},
|
|
)
|
|
|
|
user_usage[date] = max(
|
|
0,
|
|
int(seconds),
|
|
)
|
|
|
|
|
|
def _get_grants(
|
|
state: dict,
|
|
) -> list[dict]:
|
|
grants = state.setdefault(
|
|
"temporary_grants",
|
|
[],
|
|
)
|
|
|
|
if not isinstance(
|
|
grants,
|
|
list,
|
|
):
|
|
grants = []
|
|
state[
|
|
"temporary_grants"
|
|
] = grants
|
|
|
|
return grants
|
|
|
|
|
|
def _grant_is_active(
|
|
grant: dict,
|
|
username: str,
|
|
now: datetime,
|
|
) -> bool:
|
|
if str(
|
|
grant.get("username", "")
|
|
) != username:
|
|
return False
|
|
|
|
try:
|
|
remaining = int(
|
|
grant.get(
|
|
"remaining_seconds",
|
|
0,
|
|
)
|
|
)
|
|
except (
|
|
TypeError,
|
|
ValueError,
|
|
):
|
|
return False
|
|
|
|
if remaining <= 0:
|
|
return False
|
|
|
|
expires_at = grant.get(
|
|
"expires_at"
|
|
)
|
|
|
|
if not expires_at:
|
|
return True
|
|
|
|
try:
|
|
expiry = datetime.fromisoformat(
|
|
str(expires_at)
|
|
)
|
|
except ValueError:
|
|
return False
|
|
|
|
return expiry > now
|
|
|
|
|
|
def get_remaining_grant_seconds(
|
|
user_id: int,
|
|
) -> int:
|
|
user = find_user(
|
|
user_id
|
|
)
|
|
|
|
if user is None:
|
|
return 0
|
|
|
|
username = str(
|
|
user.get(
|
|
"username",
|
|
"",
|
|
)
|
|
)
|
|
|
|
now = datetime.now()
|
|
|
|
with _state_lock:
|
|
state = load_state()
|
|
|
|
total = 0
|
|
|
|
for grant in _get_grants(
|
|
state
|
|
):
|
|
if _grant_is_active(
|
|
grant,
|
|
username,
|
|
now,
|
|
):
|
|
try:
|
|
total += int(
|
|
grant.get(
|
|
"remaining_seconds",
|
|
0,
|
|
)
|
|
)
|
|
except (
|
|
TypeError,
|
|
ValueError,
|
|
):
|
|
pass
|
|
|
|
return max(
|
|
0,
|
|
total,
|
|
)
|
|
|
|
|
|
def consume_grant_seconds(
|
|
user_id: int,
|
|
seconds: int,
|
|
) -> None:
|
|
if seconds <= 0:
|
|
return
|
|
|
|
user = find_user(
|
|
user_id
|
|
)
|
|
|
|
if user is None:
|
|
return
|
|
|
|
username = str(
|
|
user.get(
|
|
"username",
|
|
"",
|
|
)
|
|
)
|
|
|
|
now = datetime.now()
|
|
|
|
with _state_lock:
|
|
state = load_state()
|
|
|
|
remaining_to_consume = (
|
|
int(seconds)
|
|
)
|
|
|
|
for grant in _get_grants(
|
|
state
|
|
):
|
|
if remaining_to_consume <= 0:
|
|
break
|
|
|
|
if not _grant_is_active(
|
|
grant,
|
|
username,
|
|
now,
|
|
):
|
|
continue
|
|
|
|
try:
|
|
available = int(
|
|
grant.get(
|
|
"remaining_seconds",
|
|
0,
|
|
)
|
|
)
|
|
except (
|
|
TypeError,
|
|
ValueError,
|
|
):
|
|
continue
|
|
|
|
consumed = min(
|
|
available,
|
|
remaining_to_consume,
|
|
)
|
|
|
|
grant[
|
|
"remaining_seconds"
|
|
] = available - consumed
|
|
|
|
remaining_to_consume -= (
|
|
consumed
|
|
)
|
|
|
|
if (
|
|
grant[
|
|
"remaining_seconds"
|
|
] <= 0
|
|
):
|
|
grant[
|
|
"remaining_seconds"
|
|
] = 0
|
|
|
|
save_state(
|
|
state
|
|
)
|
|
|
|
|
|
def record_usage(
|
|
user_id: int,
|
|
seconds: int,
|
|
) -> None:
|
|
if seconds <= 0:
|
|
return
|
|
|
|
user = find_user(
|
|
user_id
|
|
)
|
|
|
|
if user is None:
|
|
return
|
|
|
|
username = str(
|
|
user.get(
|
|
"username",
|
|
"",
|
|
)
|
|
)
|
|
|
|
today = _today()
|
|
|
|
with _state_lock:
|
|
state = load_state()
|
|
|
|
current = _get_usage(
|
|
state,
|
|
username,
|
|
today,
|
|
)
|
|
|
|
_set_usage(
|
|
state,
|
|
username,
|
|
today,
|
|
current + int(seconds),
|
|
)
|
|
|
|
save_state(
|
|
state
|
|
)
|
|
|
|
|
|
def get_user_policy(
|
|
user_id: int,
|
|
weekday: str,
|
|
):
|
|
user = find_user(
|
|
user_id
|
|
)
|
|
|
|
if user is None:
|
|
return (
|
|
0,
|
|
0,
|
|
[],
|
|
0,
|
|
)
|
|
|
|
allowances = normalize_allowances(
|
|
user.get(
|
|
"daily_allowance",
|
|
{},
|
|
)
|
|
)
|
|
|
|
allowance_seconds = (
|
|
allowances.get(
|
|
weekday,
|
|
0,
|
|
)
|
|
)
|
|
|
|
today = _today()
|
|
|
|
with _state_lock:
|
|
state = load_state()
|
|
|
|
username = str(
|
|
user.get(
|
|
"username",
|
|
"",
|
|
)
|
|
)
|
|
|
|
usage_seconds = _get_usage(
|
|
state,
|
|
username,
|
|
today,
|
|
)
|
|
|
|
grants = _get_grants(
|
|
state
|
|
)
|
|
|
|
now = datetime.now()
|
|
|
|
grant_seconds = 0
|
|
|
|
for grant in grants:
|
|
if _grant_is_active(
|
|
grant,
|
|
username,
|
|
now,
|
|
):
|
|
try:
|
|
grant_seconds += int(
|
|
grant.get(
|
|
"remaining_seconds",
|
|
0,
|
|
)
|
|
)
|
|
except (
|
|
TypeError,
|
|
ValueError,
|
|
):
|
|
pass
|
|
|
|
access_windows = normalize_access_windows(
|
|
user.get(
|
|
"access_windows",
|
|
{},
|
|
)
|
|
)
|
|
|
|
windows = access_windows.get(
|
|
weekday,
|
|
[],
|
|
)
|
|
|
|
return (
|
|
allowance_seconds,
|
|
usage_seconds,
|
|
windows,
|
|
max(
|
|
0,
|
|
grant_seconds,
|
|
),
|
|
)
|
|
|
|
|
|
def is_inside_window(
|
|
windows,
|
|
minute: int,
|
|
) -> bool:
|
|
if not windows:
|
|
return True
|
|
|
|
for window in windows:
|
|
try:
|
|
start = parse_time(
|
|
window["start"]
|
|
)
|
|
end = parse_time(
|
|
window["end"]
|
|
)
|
|
except (
|
|
KeyError,
|
|
TypeError,
|
|
ValueError,
|
|
):
|
|
continue
|
|
|
|
# Same time means all day.
|
|
if start == end:
|
|
return True
|
|
|
|
# Normal same-day window.
|
|
if start < end:
|
|
if start <= minute < end:
|
|
return True
|
|
|
|
# Overnight window.
|
|
else:
|
|
if (
|
|
minute >= start
|
|
or minute < end
|
|
):
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
def _clean_old_state() -> None:
|
|
today = _today()
|
|
|
|
with _state_lock:
|
|
state = load_state()
|
|
|
|
usage = state.get(
|
|
"usage",
|
|
{},
|
|
)
|
|
|
|
if isinstance(
|
|
usage,
|
|
dict,
|
|
):
|
|
for username in list(
|
|
usage.keys()
|
|
):
|
|
user_usage = usage[
|
|
username
|
|
]
|
|
|
|
if not isinstance(
|
|
user_usage,
|
|
dict,
|
|
):
|
|
del usage[
|
|
username
|
|
]
|
|
continue
|
|
|
|
dates = sorted(
|
|
user_usage.keys()
|
|
)
|
|
|
|
if len(dates) > 31:
|
|
for old_date in dates[
|
|
:-31
|
|
]:
|
|
del user_usage[
|
|
old_date
|
|
]
|
|
|
|
grants = state.get(
|
|
"temporary_grants",
|
|
[],
|
|
)
|
|
|
|
if isinstance(
|
|
grants,
|
|
list,
|
|
):
|
|
now = datetime.now()
|
|
|
|
cleaned = []
|
|
|
|
for grant in grants:
|
|
try:
|
|
remaining = int(
|
|
grant.get(
|
|
"remaining_seconds",
|
|
0,
|
|
)
|
|
)
|
|
except (
|
|
TypeError,
|
|
ValueError,
|
|
):
|
|
continue
|
|
|
|
if remaining <= 0:
|
|
continue
|
|
|
|
expires_at = grant.get(
|
|
"expires_at"
|
|
)
|
|
|
|
if expires_at:
|
|
try:
|
|
expiry = datetime.fromisoformat(
|
|
str(expires_at)
|
|
)
|
|
except ValueError:
|
|
continue
|
|
|
|
if expiry <= now:
|
|
continue
|
|
|
|
cleaned.append(
|
|
grant
|
|
)
|
|
|
|
state[
|
|
"temporary_grants"
|
|
] = cleaned
|
|
|
|
save_state(
|
|
state
|
|
)
|
|
|
|
|
|
def evaluate_user(
|
|
user_id: int,
|
|
username: str,
|
|
):
|
|
now, weekday, minute = (
|
|
current_time()
|
|
)
|
|
|
|
user = find_user(
|
|
user_id
|
|
)
|
|
|
|
if user is None:
|
|
return {
|
|
"user_id": user_id,
|
|
"username": username,
|
|
"timestamp": now.isoformat(),
|
|
"weekday": weekday,
|
|
"minute": minute,
|
|
"inside_window": False,
|
|
"logged_in": False,
|
|
"allowance_seconds": 0,
|
|
"usage_seconds": 0,
|
|
"allowance_remaining": 0,
|
|
"grant_seconds": 0,
|
|
"total_remaining": 0,
|
|
"allowed_by_schedule": False,
|
|
"allowed_by_grant": False,
|
|
"allowed": False,
|
|
}
|
|
|
|
if str(
|
|
user.get(
|
|
"username",
|
|
"",
|
|
)
|
|
) != username:
|
|
return {
|
|
"user_id": user_id,
|
|
"username": username,
|
|
"timestamp": now.isoformat(),
|
|
"weekday": weekday,
|
|
"minute": minute,
|
|
"inside_window": False,
|
|
"logged_in": False,
|
|
"allowance_seconds": 0,
|
|
"usage_seconds": 0,
|
|
"allowance_remaining": 0,
|
|
"grant_seconds": 0,
|
|
"total_remaining": 0,
|
|
"allowed_by_schedule": False,
|
|
"allowed_by_grant": False,
|
|
"allowed": False,
|
|
}
|
|
|
|
if not bool(
|
|
user.get(
|
|
"enabled",
|
|
True,
|
|
)
|
|
):
|
|
logged_in = user_has_session(
|
|
username
|
|
)
|
|
|
|
if logged_in:
|
|
try:
|
|
terminate_user(
|
|
username
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
if not is_locked(
|
|
username
|
|
):
|
|
lock_user(
|
|
username
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
return {
|
|
"user_id": user_id,
|
|
"username": username,
|
|
"timestamp": now.isoformat(),
|
|
"weekday": weekday,
|
|
"minute": minute,
|
|
"inside_window": False,
|
|
"logged_in": logged_in,
|
|
"allowance_seconds": 0,
|
|
"usage_seconds": 0,
|
|
"allowance_remaining": 0,
|
|
"grant_seconds": 0,
|
|
"total_remaining": 0,
|
|
"allowed_by_schedule": False,
|
|
"allowed_by_grant": False,
|
|
"allowed": False,
|
|
}
|
|
|
|
(
|
|
allowance_seconds,
|
|
usage_seconds,
|
|
windows,
|
|
grant_seconds,
|
|
) = get_user_policy(
|
|
user_id,
|
|
weekday,
|
|
)
|
|
|
|
inside_window = is_inside_window(
|
|
windows,
|
|
minute,
|
|
)
|
|
|
|
allowance_remaining = max(
|
|
0,
|
|
allowance_seconds
|
|
- usage_seconds,
|
|
)
|
|
|
|
total_remaining = (
|
|
allowance_remaining
|
|
+ grant_seconds
|
|
)
|
|
|
|
logged_in = user_has_session(
|
|
username
|
|
)
|
|
|
|
allowed_by_schedule = (
|
|
inside_window
|
|
and allowance_remaining > 0
|
|
)
|
|
|
|
allowed_by_grant = (
|
|
grant_seconds > 0
|
|
)
|
|
|
|
should_allow = (
|
|
allowed_by_schedule
|
|
or allowed_by_grant
|
|
)
|
|
|
|
# ---------------------------------------------------------
|
|
# DAILY USAGE ACCOUNTING
|
|
#
|
|
# The scheduler runs once every second.
|
|
#
|
|
# Only count usage when:
|
|
# 1. The user is actually logged in.
|
|
# 2. The user is inside an allowed access window.
|
|
# 3. The user still has daily allowance remaining.
|
|
#
|
|
# Temporary grants are consumed separately below.
|
|
# ---------------------------------------------------------
|
|
if (
|
|
logged_in
|
|
and inside_window
|
|
and allowance_remaining > 0
|
|
):
|
|
record_usage(
|
|
user_id,
|
|
1,
|
|
)
|
|
|
|
# Refresh the usage value immediately so
|
|
# the returned policy information reflects
|
|
# the second we just consumed.
|
|
usage_seconds += 1
|
|
|
|
allowance_remaining = max(
|
|
0,
|
|
allowance_seconds
|
|
- usage_seconds,
|
|
)
|
|
|
|
total_remaining = (
|
|
allowance_remaining
|
|
+ grant_seconds
|
|
)
|
|
|
|
# If this second consumed the final
|
|
# allowance second, the user must be
|
|
# locked during this same scheduler pass.
|
|
if allowance_remaining <= 0:
|
|
should_allow = (
|
|
grant_seconds > 0
|
|
)
|
|
|
|
# ---------------------------------------------------------
|
|
# TEMPORARY GRANT ACCOUNTING
|
|
#
|
|
# Temporary time is consumed only while the
|
|
# user is actually logged in and is being
|
|
# allowed by the temporary grant.
|
|
#
|
|
# The grant also has its own expires_at timestamp,
|
|
# so it cannot remain valid beyond its expiry.
|
|
# ---------------------------------------------------------
|
|
if (
|
|
logged_in
|
|
and grant_seconds > 0
|
|
):
|
|
consume_grant_seconds(
|
|
user_id,
|
|
1,
|
|
)
|
|
|
|
grant_seconds = max(
|
|
0,
|
|
grant_seconds - 1,
|
|
)
|
|
|
|
total_remaining = (
|
|
allowance_remaining
|
|
+ grant_seconds
|
|
)
|
|
|
|
should_allow = (
|
|
allowed_by_schedule
|
|
or grant_seconds > 0
|
|
)
|
|
|
|
try:
|
|
locked = is_locked(
|
|
username
|
|
)
|
|
except Exception:
|
|
locked = False
|
|
|
|
if should_allow:
|
|
if locked:
|
|
try:
|
|
unlock_user(
|
|
username
|
|
)
|
|
|
|
if not is_locked(
|
|
username
|
|
):
|
|
pass
|
|
|
|
except Exception as exc:
|
|
with _state_lock:
|
|
state = load_state()
|
|
|
|
_ = state
|
|
_ = exc
|
|
|
|
else:
|
|
if logged_in:
|
|
try:
|
|
terminate_user(
|
|
username
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
if not locked:
|
|
try:
|
|
lock_user(
|
|
username
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
return {
|
|
"user_id": user_id,
|
|
"username": username,
|
|
"timestamp": now.isoformat(),
|
|
"weekday": weekday,
|
|
"minute": minute,
|
|
"inside_window": inside_window,
|
|
"logged_in": logged_in,
|
|
"allowance_seconds": allowance_seconds,
|
|
"usage_seconds": usage_seconds,
|
|
"allowance_remaining": allowance_remaining,
|
|
"grant_seconds": grant_seconds,
|
|
"total_remaining": total_remaining,
|
|
"allowed_by_schedule": (
|
|
inside_window
|
|
and allowance_remaining > 0
|
|
),
|
|
"allowed_by_grant": (
|
|
grant_seconds > 0
|
|
),
|
|
"allowed": should_allow,
|
|
}
|
|
|
|
|
|
def enforce_all_users():
|
|
_clean_old_state()
|
|
|
|
config = load_users_config()
|
|
|
|
users = config.get(
|
|
"users",
|
|
[],
|
|
)
|
|
|
|
results = []
|
|
|
|
for user in users:
|
|
if not isinstance(
|
|
user,
|
|
dict,
|
|
):
|
|
continue
|
|
|
|
if not bool(
|
|
user.get(
|
|
"enabled",
|
|
True,
|
|
)
|
|
):
|
|
continue
|
|
|
|
try:
|
|
user_id = int(
|
|
user["id"]
|
|
)
|
|
|
|
username = str(
|
|
user["username"]
|
|
)
|
|
|
|
except (
|
|
KeyError,
|
|
TypeError,
|
|
ValueError,
|
|
):
|
|
continue
|
|
|
|
try:
|
|
result = evaluate_user(
|
|
user_id,
|
|
username,
|
|
)
|
|
|
|
results.append(
|
|
result
|
|
)
|
|
|
|
except Exception as exc:
|
|
results.append(
|
|
{
|
|
"user_id": user_id,
|
|
"username": username,
|
|
"error": str(exc),
|
|
}
|
|
)
|
|
|
|
return results
|