added pam auth
This commit is contained in:
-105
@@ -1,105 +0,0 @@
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
from contextlib import contextmanager
|
||||
|
||||
BASE_DIR = Path(__file__).resolve().parent.parent
|
||||
|
||||
DATABASE_DIR = BASE_DIR / "data"
|
||||
DATABASE_PATH = DATABASE_DIR / "parental-control.db"
|
||||
|
||||
|
||||
def initialize_database():
|
||||
DATABASE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with sqlite3.connect(DATABASE_PATH) as db:
|
||||
db.execute("PRAGMA foreign_keys = ON")
|
||||
|
||||
db.executescript(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
username TEXT NOT NULL UNIQUE,
|
||||
enabled INTEGER NOT NULL DEFAULT 1,
|
||||
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS daily_allowances (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL,
|
||||
weekday INTEGER NOT NULL,
|
||||
allowance_seconds INTEGER NOT NULL DEFAULT 0,
|
||||
|
||||
UNIQUE(user_id, weekday),
|
||||
|
||||
FOREIGN KEY(user_id)
|
||||
REFERENCES users(id)
|
||||
ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS access_windows (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL,
|
||||
weekday INTEGER NOT NULL,
|
||||
start_minute INTEGER NOT NULL,
|
||||
end_minute INTEGER NOT NULL,
|
||||
|
||||
FOREIGN KEY(user_id)
|
||||
REFERENCES users(id)
|
||||
ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS usage (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL,
|
||||
date TEXT NOT NULL,
|
||||
used_seconds INTEGER NOT NULL DEFAULT 0,
|
||||
|
||||
UNIQUE(user_id, date),
|
||||
|
||||
FOREIGN KEY(user_id)
|
||||
REFERENCES users(id)
|
||||
ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS temporary_grants (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL,
|
||||
seconds INTEGER NOT NULL,
|
||||
remaining_seconds INTEGER NOT NULL,
|
||||
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
expires_at TEXT,
|
||||
consumed INTEGER NOT NULL DEFAULT 0,
|
||||
FOREIGN KEY(user_id) REFERENCES users(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS events (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER,
|
||||
event_type TEXT NOT NULL,
|
||||
details TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
|
||||
FOREIGN KEY(user_id)
|
||||
REFERENCES users(id)
|
||||
ON DELETE SET NULL
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
db.commit()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def get_db():
|
||||
db = sqlite3.connect(DATABASE_PATH)
|
||||
db.row_factory = sqlite3.Row
|
||||
db.execute("PRAGMA foreign_keys = ON")
|
||||
|
||||
try:
|
||||
yield db
|
||||
db.commit()
|
||||
except Exception:
|
||||
db.rollback()
|
||||
raise
|
||||
finally:
|
||||
db.close()
|
||||
+33
-361
@@ -1,7 +1,14 @@
|
||||
from datetime import datetime
|
||||
import subprocess
|
||||
|
||||
from .database import get_db
|
||||
from .storage import (
|
||||
get_user_policy,
|
||||
get_remaining_grant_seconds,
|
||||
consume_grant_seconds,
|
||||
record_usage,
|
||||
record_event,
|
||||
list_users,
|
||||
)
|
||||
from .users import (
|
||||
lock_user,
|
||||
unlock_user,
|
||||
@@ -12,11 +19,7 @@ from .users import (
|
||||
|
||||
def user_has_session(username: str) -> bool:
|
||||
result = subprocess.run(
|
||||
[
|
||||
"loginctl",
|
||||
"list-users",
|
||||
"--no-legend",
|
||||
],
|
||||
["loginctl", "list-users", "--no-legend"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
@@ -25,297 +28,28 @@ def user_has_session(username: str) -> bool:
|
||||
if result.returncode != 0:
|
||||
return False
|
||||
|
||||
for line in result.stdout.splitlines():
|
||||
parts = line.split()
|
||||
|
||||
if len(parts) >= 2 and parts[1] == username:
|
||||
return True
|
||||
|
||||
return False
|
||||
return any(
|
||||
len(parts := line.split()) >= 2 and parts[1] == username
|
||||
for line in result.stdout.splitlines()
|
||||
)
|
||||
|
||||
|
||||
def current_time():
|
||||
now = datetime.now()
|
||||
weekday = now.weekday()
|
||||
minute = now.hour * 60 + now.minute
|
||||
|
||||
return now, weekday, minute
|
||||
|
||||
|
||||
def get_user_policy(user_id: int, weekday: int):
|
||||
with get_db() as db:
|
||||
allowance_row = db.execute(
|
||||
"""
|
||||
SELECT allowance_seconds
|
||||
FROM daily_allowances
|
||||
WHERE user_id = ?
|
||||
AND weekday = ?
|
||||
""",
|
||||
(
|
||||
user_id,
|
||||
weekday,
|
||||
),
|
||||
).fetchone()
|
||||
|
||||
allowance_seconds = (
|
||||
allowance_row["allowance_seconds"]
|
||||
if allowance_row
|
||||
else 0
|
||||
)
|
||||
|
||||
today = datetime.now().date().isoformat()
|
||||
|
||||
usage_row = db.execute(
|
||||
"""
|
||||
SELECT used_seconds
|
||||
FROM usage
|
||||
WHERE user_id = ?
|
||||
AND date = ?
|
||||
""",
|
||||
(
|
||||
user_id,
|
||||
today,
|
||||
),
|
||||
).fetchone()
|
||||
|
||||
usage_seconds = (
|
||||
usage_row["used_seconds"]
|
||||
if usage_row
|
||||
else 0
|
||||
)
|
||||
|
||||
windows = db.execute(
|
||||
"""
|
||||
SELECT id, start_minute, end_minute
|
||||
FROM access_windows
|
||||
WHERE user_id = ?
|
||||
AND weekday = ?
|
||||
ORDER BY start_minute
|
||||
""",
|
||||
(
|
||||
user_id,
|
||||
weekday,
|
||||
),
|
||||
).fetchall()
|
||||
|
||||
grant_row = db.execute(
|
||||
"""
|
||||
SELECT COALESCE(
|
||||
SUM(remaining_seconds),
|
||||
0
|
||||
) AS total
|
||||
FROM temporary_grants
|
||||
WHERE user_id = ?
|
||||
AND consumed = 0
|
||||
AND (
|
||||
expires_at IS NULL
|
||||
OR expires_at > ?
|
||||
)
|
||||
""",
|
||||
(
|
||||
user_id,
|
||||
datetime.now().isoformat(),
|
||||
),
|
||||
).fetchone()
|
||||
|
||||
grant_seconds = grant_row["total"]
|
||||
|
||||
return (
|
||||
allowance_seconds,
|
||||
usage_seconds,
|
||||
windows,
|
||||
grant_seconds,
|
||||
)
|
||||
return now, now.weekday(), now.hour * 60 + now.minute
|
||||
|
||||
|
||||
def is_inside_window(windows, minute: int) -> bool:
|
||||
if not windows:
|
||||
return True
|
||||
|
||||
for window in windows:
|
||||
if (
|
||||
window["start_minute"]
|
||||
<= minute
|
||||
< window["end_minute"]
|
||||
):
|
||||
return True
|
||||
|
||||
return False
|
||||
return any(
|
||||
int(window["start_minute"]) <= minute < int(window["end_minute"])
|
||||
for window in windows
|
||||
)
|
||||
|
||||
|
||||
def get_remaining_grant_seconds(user_id: int) -> int:
|
||||
now = datetime.now().isoformat()
|
||||
|
||||
with get_db() as db:
|
||||
row = db.execute(
|
||||
"""
|
||||
SELECT COALESCE(
|
||||
SUM(remaining_seconds),
|
||||
0
|
||||
) AS total
|
||||
FROM temporary_grants
|
||||
WHERE user_id = ?
|
||||
AND consumed = 0
|
||||
AND (
|
||||
expires_at IS NULL
|
||||
OR expires_at > ?
|
||||
)
|
||||
""",
|
||||
(
|
||||
user_id,
|
||||
now,
|
||||
),
|
||||
).fetchone()
|
||||
|
||||
return row["total"]
|
||||
|
||||
|
||||
def consume_grant_seconds(
|
||||
user_id: int,
|
||||
seconds: int,
|
||||
):
|
||||
if seconds <= 0:
|
||||
return
|
||||
|
||||
now = datetime.now().isoformat()
|
||||
|
||||
with get_db() as db:
|
||||
grants = db.execute(
|
||||
"""
|
||||
SELECT id, remaining_seconds
|
||||
FROM temporary_grants
|
||||
WHERE user_id = ?
|
||||
AND consumed = 0
|
||||
AND remaining_seconds > 0
|
||||
AND (
|
||||
expires_at IS NULL
|
||||
OR expires_at > ?
|
||||
)
|
||||
ORDER BY id ASC
|
||||
""",
|
||||
(
|
||||
user_id,
|
||||
now,
|
||||
),
|
||||
).fetchall()
|
||||
|
||||
remaining = seconds
|
||||
|
||||
for grant in grants:
|
||||
if remaining <= 0:
|
||||
break
|
||||
|
||||
available = grant["remaining_seconds"]
|
||||
|
||||
consumed = min(
|
||||
available,
|
||||
remaining,
|
||||
)
|
||||
|
||||
new_remaining = (
|
||||
available - consumed
|
||||
)
|
||||
|
||||
db.execute(
|
||||
"""
|
||||
UPDATE temporary_grants
|
||||
SET remaining_seconds = ?,
|
||||
consumed = ?
|
||||
WHERE id = ?
|
||||
""",
|
||||
(
|
||||
new_remaining,
|
||||
1 if new_remaining <= 0 else 0,
|
||||
grant["id"],
|
||||
),
|
||||
)
|
||||
|
||||
remaining -= consumed
|
||||
|
||||
|
||||
def record_usage(
|
||||
user_id: int,
|
||||
seconds: int,
|
||||
):
|
||||
if seconds <= 0:
|
||||
return
|
||||
|
||||
today = datetime.now().date().isoformat()
|
||||
|
||||
with get_db() as db:
|
||||
row = db.execute(
|
||||
"""
|
||||
SELECT used_seconds
|
||||
FROM usage
|
||||
WHERE user_id = ?
|
||||
AND date = ?
|
||||
""",
|
||||
(
|
||||
user_id,
|
||||
today,
|
||||
),
|
||||
).fetchone()
|
||||
|
||||
if row is None:
|
||||
db.execute(
|
||||
"""
|
||||
INSERT INTO usage (
|
||||
user_id,
|
||||
date,
|
||||
used_seconds
|
||||
)
|
||||
VALUES (?, ?, ?)
|
||||
""",
|
||||
(
|
||||
user_id,
|
||||
today,
|
||||
seconds,
|
||||
),
|
||||
)
|
||||
else:
|
||||
db.execute(
|
||||
"""
|
||||
UPDATE usage
|
||||
SET used_seconds =
|
||||
used_seconds + ?
|
||||
WHERE user_id = ?
|
||||
AND date = ?
|
||||
""",
|
||||
(
|
||||
seconds,
|
||||
user_id,
|
||||
today,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def record_event(
|
||||
user_id: int,
|
||||
event_type: str,
|
||||
details: str = "",
|
||||
):
|
||||
with get_db() as db:
|
||||
db.execute(
|
||||
"""
|
||||
INSERT INTO events (
|
||||
user_id,
|
||||
event_type,
|
||||
details
|
||||
)
|
||||
VALUES (?, ?, ?)
|
||||
""",
|
||||
(
|
||||
user_id,
|
||||
event_type,
|
||||
details,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def evaluate_user(
|
||||
user_id: int,
|
||||
username: str,
|
||||
):
|
||||
def evaluate_user(user_id: int, username: str):
|
||||
now, weekday, minute = current_time()
|
||||
|
||||
(
|
||||
@@ -323,43 +57,16 @@ def evaluate_user(
|
||||
usage_seconds,
|
||||
windows,
|
||||
grant_seconds,
|
||||
) = get_user_policy(
|
||||
user_id,
|
||||
weekday,
|
||||
)
|
||||
) = get_user_policy(user_id, weekday)
|
||||
|
||||
inside_window = is_inside_window(
|
||||
windows,
|
||||
minute,
|
||||
)
|
||||
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)
|
||||
|
||||
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
|
||||
)
|
||||
allowed_by_schedule = inside_window and allowance_remaining > 0
|
||||
allowed_by_grant = grant_seconds > 0
|
||||
should_allow = allowed_by_schedule or allowed_by_grant
|
||||
|
||||
locked = is_locked(username)
|
||||
|
||||
@@ -367,23 +74,12 @@ def evaluate_user(
|
||||
if locked:
|
||||
try:
|
||||
unlock_user(username)
|
||||
|
||||
record_event(
|
||||
user_id,
|
||||
"auto_unlock",
|
||||
"Access became available",
|
||||
)
|
||||
record_event(user_id, "auto_unlock", "Access became available")
|
||||
except Exception as exc:
|
||||
record_event(
|
||||
user_id,
|
||||
"unlock_error",
|
||||
str(exc),
|
||||
)
|
||||
|
||||
record_event(user_id, "unlock_error", str(exc))
|
||||
else:
|
||||
if logged_in:
|
||||
terminate_user(username)
|
||||
|
||||
record_event(
|
||||
user_id,
|
||||
"session_terminated",
|
||||
@@ -393,18 +89,13 @@ def evaluate_user(
|
||||
if not locked:
|
||||
try:
|
||||
lock_user(username)
|
||||
|
||||
record_event(
|
||||
user_id,
|
||||
"auto_lock",
|
||||
"Access is not currently permitted",
|
||||
)
|
||||
except Exception as exc:
|
||||
record_event(
|
||||
user_id,
|
||||
"lock_error",
|
||||
str(exc),
|
||||
)
|
||||
record_event(user_id, "lock_error", str(exc))
|
||||
|
||||
return {
|
||||
"user_id": user_id,
|
||||
@@ -424,32 +115,13 @@ def evaluate_user(
|
||||
|
||||
|
||||
def enforce_all_users():
|
||||
with get_db() as db:
|
||||
users = db.execute(
|
||||
"""
|
||||
SELECT id, username, enabled
|
||||
FROM users
|
||||
WHERE enabled = 1
|
||||
ORDER BY id
|
||||
"""
|
||||
).fetchall()
|
||||
|
||||
users = list_users()
|
||||
results = []
|
||||
|
||||
for user in users:
|
||||
try:
|
||||
result = evaluate_user(
|
||||
user["id"],
|
||||
user["username"],
|
||||
)
|
||||
|
||||
results.append(result)
|
||||
|
||||
results.append(evaluate_user(user["id"], user["username"]))
|
||||
except Exception as exc:
|
||||
record_event(
|
||||
user["id"],
|
||||
"enforcement_error",
|
||||
str(exc),
|
||||
)
|
||||
record_event(user["id"], "enforcement_error", str(exc))
|
||||
|
||||
return results
|
||||
|
||||
+335
-614
File diff suppressed because it is too large
Load Diff
+2
-6
@@ -2,12 +2,8 @@ import threading
|
||||
import time
|
||||
from datetime import datetime
|
||||
|
||||
from .enforcement import (
|
||||
enforce_all_users,
|
||||
record_usage,
|
||||
get_user_policy,
|
||||
consume_grant_seconds,
|
||||
)
|
||||
from .enforcement import enforce_all_users
|
||||
from .storage import record_usage, get_user_policy, consume_grant_seconds
|
||||
|
||||
|
||||
CHECK_INTERVAL = 5
|
||||
|
||||
+445
@@ -0,0 +1,445 @@
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from threading import RLock
|
||||
|
||||
import yaml
|
||||
|
||||
BASE_DIR = Path(__file__).resolve().parent.parent
|
||||
DATA_DIR = BASE_DIR / "data"
|
||||
CONFIG_PATH = DATA_DIR / "config.yaml"
|
||||
STATE_PATH = DATA_DIR / "state.json"
|
||||
|
||||
_lock = RLock()
|
||||
|
||||
DEFAULT_CONFIG = {
|
||||
"version": 1,
|
||||
"auth": {
|
||||
"pam_service": "login",
|
||||
"admin_users": [],
|
||||
},
|
||||
"users": [],
|
||||
}
|
||||
|
||||
DEFAULT_STATE = {
|
||||
"version": 1,
|
||||
"next_ids": {
|
||||
"user": 1,
|
||||
"window": 1,
|
||||
"grant": 1,
|
||||
},
|
||||
"usage": {},
|
||||
"temporary_grants": [],
|
||||
"events": [],
|
||||
}
|
||||
|
||||
|
||||
def _atomic_write(path: Path, content: str, mode: int = 0o600):
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
fd, tmp_name = tempfile.mkstemp(
|
||||
prefix=f".{path.name}.",
|
||||
dir=str(path.parent),
|
||||
text=True,
|
||||
)
|
||||
try:
|
||||
os.fchmod(fd, mode)
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as handle:
|
||||
handle.write(content)
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
os.replace(tmp_name, path)
|
||||
finally:
|
||||
if os.path.exists(tmp_name):
|
||||
os.unlink(tmp_name)
|
||||
|
||||
|
||||
def _load_yaml():
|
||||
if not CONFIG_PATH.exists():
|
||||
return json.loads(json.dumps(DEFAULT_CONFIG))
|
||||
|
||||
with CONFIG_PATH.open("r", encoding="utf-8") as handle:
|
||||
data = yaml.safe_load(handle) or {}
|
||||
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("config.yaml must contain a YAML object")
|
||||
|
||||
data.setdefault("version", 1)
|
||||
data.setdefault("auth", {})
|
||||
data["auth"].setdefault("pam_service", "login")
|
||||
data["auth"].setdefault("admin_users", [])
|
||||
data.setdefault("users", [])
|
||||
return data
|
||||
|
||||
|
||||
def _load_json():
|
||||
if not STATE_PATH.exists():
|
||||
return json.loads(json.dumps(DEFAULT_STATE))
|
||||
|
||||
with STATE_PATH.open("r", encoding="utf-8") as handle:
|
||||
data = json.load(handle)
|
||||
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("state.json must contain a JSON object")
|
||||
|
||||
data.setdefault("version", 1)
|
||||
data.setdefault("next_ids", {})
|
||||
data["next_ids"].setdefault("user", 1)
|
||||
data["next_ids"].setdefault("window", 1)
|
||||
data["next_ids"].setdefault("grant", 1)
|
||||
data.setdefault("usage", {})
|
||||
data.setdefault("temporary_grants", [])
|
||||
data.setdefault("events", [])
|
||||
return data
|
||||
|
||||
|
||||
def _save_yaml(data):
|
||||
_atomic_write(
|
||||
CONFIG_PATH,
|
||||
yaml.safe_dump(
|
||||
data,
|
||||
sort_keys=False,
|
||||
default_flow_style=False,
|
||||
allow_unicode=True,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _save_json(data):
|
||||
_atomic_write(
|
||||
STATE_PATH,
|
||||
json.dumps(data, indent=2, ensure_ascii=False) + "\n",
|
||||
)
|
||||
|
||||
|
||||
def initialize_storage():
|
||||
DATA_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with _lock:
|
||||
if not CONFIG_PATH.exists():
|
||||
_save_yaml(DEFAULT_CONFIG)
|
||||
|
||||
if not STATE_PATH.exists():
|
||||
_save_json(DEFAULT_STATE)
|
||||
|
||||
# Keep files usable after manual edits while avoiding destructive
|
||||
# initialization or recreation of any database.
|
||||
config = _load_yaml()
|
||||
state = _load_json()
|
||||
_save_yaml(config)
|
||||
_save_json(state)
|
||||
|
||||
|
||||
def get_config():
|
||||
with _lock:
|
||||
return _load_yaml()
|
||||
|
||||
|
||||
def get_pam_service():
|
||||
return get_config()["auth"].get("pam_service", "login")
|
||||
|
||||
|
||||
def admin_users():
|
||||
value = get_config()["auth"].get("admin_users", [])
|
||||
return {str(item) for item in value}
|
||||
|
||||
|
||||
def is_admin_allowed(username: str) -> bool:
|
||||
allowed = admin_users()
|
||||
return not allowed or username in allowed
|
||||
|
||||
|
||||
def _find_user(config, user_id):
|
||||
for user in config["users"]:
|
||||
if int(user["id"]) == int(user_id):
|
||||
return user
|
||||
return None
|
||||
|
||||
|
||||
def _find_window(config, window_id):
|
||||
for user in config["users"]:
|
||||
for window in user.get("windows", []):
|
||||
if int(window["id"]) == int(window_id):
|
||||
return user, window
|
||||
return None, None
|
||||
|
||||
|
||||
def list_users():
|
||||
with _lock:
|
||||
config = _load_yaml()
|
||||
return sorted(
|
||||
[dict(user) for user in config["users"] if user.get("enabled", True)],
|
||||
key=lambda item: item["username"],
|
||||
)
|
||||
|
||||
|
||||
def get_user(user_id: int):
|
||||
with _lock:
|
||||
config = _load_yaml()
|
||||
user = _find_user(config, user_id)
|
||||
return dict(user) if user else None
|
||||
|
||||
|
||||
def get_user_by_username(username: str):
|
||||
with _lock:
|
||||
config = _load_yaml()
|
||||
for user in config["users"]:
|
||||
if user["username"] == username:
|
||||
return dict(user)
|
||||
return None
|
||||
|
||||
|
||||
def add_user(username: str):
|
||||
with _lock:
|
||||
config = _load_yaml()
|
||||
if any(u["username"] == username for u in config["users"]):
|
||||
raise ValueError("User is already configured")
|
||||
|
||||
state = _load_json()
|
||||
user_id = int(state["next_ids"]["user"])
|
||||
state["next_ids"]["user"] = user_id + 1
|
||||
|
||||
user = {
|
||||
"id": user_id,
|
||||
"username": username,
|
||||
"enabled": True,
|
||||
"allowances": {str(day): 0 for day in range(7)},
|
||||
"windows": [],
|
||||
}
|
||||
config["users"].append(user)
|
||||
_save_yaml(config)
|
||||
_save_json(state)
|
||||
return dict(user)
|
||||
|
||||
|
||||
def delete_user(user_id: int):
|
||||
with _lock:
|
||||
config = _load_yaml()
|
||||
user = _find_user(config, user_id)
|
||||
if user is None:
|
||||
return False
|
||||
|
||||
username = user["username"]
|
||||
config["users"] = [
|
||||
item for item in config["users"]
|
||||
if int(item["id"]) != int(user_id)
|
||||
]
|
||||
|
||||
state = _load_json()
|
||||
state["usage"] = {
|
||||
key: value
|
||||
for key, value in state["usage"].items()
|
||||
if not key.startswith(f"{int(user_id)}:")
|
||||
}
|
||||
state["temporary_grants"] = [
|
||||
grant for grant in state["temporary_grants"]
|
||||
if int(grant["user_id"]) != int(user_id)
|
||||
]
|
||||
state["events"] = [
|
||||
event for event in state["events"]
|
||||
if event.get("user_id") is None
|
||||
or int(event["user_id"]) != int(user_id)
|
||||
]
|
||||
|
||||
_save_yaml(config)
|
||||
_save_json(state)
|
||||
return username
|
||||
|
||||
|
||||
def set_allowance(user_id: int, weekday: int, seconds: int):
|
||||
with _lock:
|
||||
config = _load_yaml()
|
||||
user = _find_user(config, user_id)
|
||||
if user is None:
|
||||
raise KeyError("User not found")
|
||||
|
||||
user.setdefault("allowances", {})
|
||||
user["allowances"][str(weekday)] = int(seconds)
|
||||
_save_yaml(config)
|
||||
|
||||
|
||||
def add_window(user_id: int, weekday: int, start_minute: int, end_minute: int):
|
||||
with _lock:
|
||||
config = _load_yaml()
|
||||
user = _find_user(config, user_id)
|
||||
if user is None:
|
||||
raise KeyError("User not found")
|
||||
|
||||
state = _load_json()
|
||||
window_id = int(state["next_ids"]["window"])
|
||||
state["next_ids"]["window"] = window_id + 1
|
||||
|
||||
user.setdefault("windows", []).append({
|
||||
"id": window_id,
|
||||
"weekday": int(weekday),
|
||||
"start_minute": int(start_minute),
|
||||
"end_minute": int(end_minute),
|
||||
})
|
||||
|
||||
_save_yaml(config)
|
||||
_save_json(state)
|
||||
return window_id
|
||||
|
||||
|
||||
def delete_window(window_id: int):
|
||||
with _lock:
|
||||
config = _load_yaml()
|
||||
owner, window = _find_window(config, window_id)
|
||||
if owner is None:
|
||||
return None
|
||||
|
||||
owner["windows"] = [
|
||||
item for item in owner.get("windows", [])
|
||||
if int(item["id"]) != int(window_id)
|
||||
]
|
||||
_save_yaml(config)
|
||||
return int(owner["id"])
|
||||
|
||||
|
||||
def get_user_policy(user_id: int, weekday: int):
|
||||
with _lock:
|
||||
config = _load_yaml()
|
||||
user = _find_user(config, user_id)
|
||||
if user is None:
|
||||
return 0, 0, [], 0
|
||||
|
||||
allowance_seconds = int(
|
||||
user.get("allowances", {}).get(str(weekday), 0)
|
||||
)
|
||||
windows = sorted(
|
||||
[
|
||||
dict(window)
|
||||
for window in user.get("windows", [])
|
||||
if int(window["weekday"]) == int(weekday)
|
||||
],
|
||||
key=lambda item: int(item["start_minute"]),
|
||||
)
|
||||
|
||||
state = _load_json()
|
||||
today = datetime.now().date().isoformat()
|
||||
usage_seconds = int(
|
||||
state["usage"].get(f"{int(user_id)}:{today}", 0)
|
||||
)
|
||||
|
||||
now = datetime.now().isoformat()
|
||||
grant_seconds = sum(
|
||||
int(grant["remaining_seconds"])
|
||||
for grant in state["temporary_grants"]
|
||||
if int(grant["user_id"]) == int(user_id)
|
||||
and not grant.get("consumed", False)
|
||||
and (
|
||||
grant.get("expires_at") is None
|
||||
or grant["expires_at"] > now
|
||||
)
|
||||
)
|
||||
|
||||
return allowance_seconds, usage_seconds, windows, grant_seconds
|
||||
|
||||
|
||||
def get_remaining_grant_seconds(user_id: int) -> int:
|
||||
with _lock:
|
||||
state = _load_json()
|
||||
now = datetime.now().isoformat()
|
||||
return sum(
|
||||
int(grant["remaining_seconds"])
|
||||
for grant in state["temporary_grants"]
|
||||
if int(grant["user_id"]) == int(user_id)
|
||||
and not grant.get("consumed", False)
|
||||
and int(grant["remaining_seconds"]) > 0
|
||||
and (
|
||||
grant.get("expires_at") is None
|
||||
or grant["expires_at"] > now
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def add_grant(user_id: int, seconds: int):
|
||||
with _lock:
|
||||
state = _load_json()
|
||||
grant_id = int(state["next_ids"]["grant"])
|
||||
state["next_ids"]["grant"] = grant_id + 1
|
||||
state["temporary_grants"].append({
|
||||
"id": grant_id,
|
||||
"user_id": int(user_id),
|
||||
"seconds": int(seconds),
|
||||
"remaining_seconds": int(seconds),
|
||||
"created_at": datetime.now().isoformat(timespec="seconds"),
|
||||
"expires_at": None,
|
||||
"consumed": False,
|
||||
})
|
||||
_save_json(state)
|
||||
return grant_id
|
||||
|
||||
|
||||
def consume_grant_seconds(user_id: int, seconds: int):
|
||||
if seconds <= 0:
|
||||
return
|
||||
|
||||
with _lock:
|
||||
state = _load_json()
|
||||
now = datetime.now().isoformat()
|
||||
remaining = int(seconds)
|
||||
|
||||
for grant in state["temporary_grants"]:
|
||||
if remaining <= 0:
|
||||
break
|
||||
if int(grant["user_id"]) != int(user_id):
|
||||
continue
|
||||
if grant.get("consumed", False):
|
||||
continue
|
||||
if int(grant["remaining_seconds"]) <= 0:
|
||||
continue
|
||||
if (
|
||||
grant.get("expires_at") is not None
|
||||
and grant["expires_at"] <= now
|
||||
):
|
||||
continue
|
||||
|
||||
available = int(grant["remaining_seconds"])
|
||||
consumed = min(available, remaining)
|
||||
new_remaining = available - consumed
|
||||
grant["remaining_seconds"] = new_remaining
|
||||
grant["consumed"] = new_remaining <= 0
|
||||
remaining -= consumed
|
||||
|
||||
_save_json(state)
|
||||
|
||||
|
||||
def record_usage(user_id: int, seconds: int):
|
||||
if seconds <= 0:
|
||||
return
|
||||
|
||||
with _lock:
|
||||
state = _load_json()
|
||||
today = datetime.now().date().isoformat()
|
||||
key = f"{int(user_id)}:{today}"
|
||||
state["usage"][key] = int(state["usage"].get(key, 0)) + int(seconds)
|
||||
_save_json(state)
|
||||
|
||||
|
||||
def list_grants(user_id: int, limit: int = 20):
|
||||
with _lock:
|
||||
state = _load_json()
|
||||
grants = [
|
||||
dict(grant)
|
||||
for grant in state["temporary_grants"]
|
||||
if int(grant["user_id"]) == int(user_id)
|
||||
]
|
||||
grants.sort(key=lambda item: int(item["id"]), reverse=True)
|
||||
return grants[:limit]
|
||||
|
||||
|
||||
def record_event(user_id, event_type: str, details: str = ""):
|
||||
with _lock:
|
||||
state = _load_json()
|
||||
state["events"].append({
|
||||
"user_id": int(user_id) if user_id is not None else None,
|
||||
"event_type": event_type,
|
||||
"details": details,
|
||||
"created_at": datetime.now().isoformat(timespec="seconds"),
|
||||
})
|
||||
# Keep the state file bounded.
|
||||
state["events"] = state["events"][-2000:]
|
||||
_save_json(state)
|
||||
Reference in New Issue
Block a user