162 lines
5.5 KiB
Python
162 lines
5.5 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
from tempfile import NamedTemporaryFile
|
|
from threading import RLock
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
def _normalize_command(value: str) -> str:
|
|
token = value.strip().lower()
|
|
token = token.lstrip(".")
|
|
if "@" in token:
|
|
token = token.split("@", 1)[0]
|
|
return token
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class RuntimeAclStore:
|
|
path: Path
|
|
data: dict[str, dict[str, set[str]]] = field(init=False)
|
|
_lock: RLock = field(init=False, repr=False)
|
|
|
|
def __post_init__(self) -> None:
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._lock = RLock()
|
|
self.data = {"allow": {}, "deny": {}}
|
|
self.load()
|
|
|
|
def load(self) -> None:
|
|
with self._lock:
|
|
if not self.path.exists():
|
|
return
|
|
|
|
try:
|
|
raw = json.loads(self.path.read_text(encoding="utf-8"))
|
|
except json.JSONDecodeError:
|
|
log.warning("Ignoring malformed ACL file: %s", self.path)
|
|
return
|
|
|
|
if not isinstance(raw, dict):
|
|
log.warning("Ignoring ACL file with unexpected root type: %s", self.path)
|
|
return
|
|
|
|
allow_raw = raw.get("allow", {})
|
|
deny_raw = raw.get("deny", {})
|
|
self.data = {
|
|
"allow": {
|
|
str(user_id): {str(command) for command in commands}
|
|
for user_id, commands in allow_raw.items()
|
|
},
|
|
"deny": {
|
|
str(user_id): {str(command) for command in commands}
|
|
for user_id, commands in deny_raw.items()
|
|
},
|
|
}
|
|
|
|
def save(self) -> None:
|
|
with self._lock:
|
|
payload = {
|
|
"version": 1,
|
|
"allow": {
|
|
user_id: sorted(commands)
|
|
for user_id, commands in sorted(self.data["allow"].items())
|
|
if commands
|
|
},
|
|
"deny": {
|
|
user_id: sorted(commands)
|
|
for user_id, commands in sorted(self.data["deny"].items())
|
|
if commands
|
|
},
|
|
}
|
|
|
|
serialized = json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True)
|
|
with NamedTemporaryFile(
|
|
"w",
|
|
encoding="utf-8",
|
|
delete=False,
|
|
dir=self.path.parent,
|
|
prefix=self.path.name,
|
|
suffix=".tmp",
|
|
) as handle:
|
|
handle.write(serialized)
|
|
temp_path = Path(handle.name)
|
|
temp_path.replace(self.path)
|
|
|
|
def _allow_set(self, user_id: int) -> set[str]:
|
|
return self.data["allow"].setdefault(str(user_id), set())
|
|
|
|
def _deny_set(self, user_id: int) -> set[str]:
|
|
return self.data["deny"].setdefault(str(user_id), set())
|
|
|
|
def allow(self, user_id: int, command: str) -> bool:
|
|
with self._lock:
|
|
command_name = _normalize_command(command)
|
|
if not command_name:
|
|
return False
|
|
|
|
allow_set = self._allow_set(user_id)
|
|
deny_set = self._deny_set(user_id)
|
|
changed = command_name not in allow_set or command_name in deny_set
|
|
allow_set.add(command_name)
|
|
deny_set.discard(command_name)
|
|
if changed:
|
|
self.save()
|
|
return changed
|
|
|
|
def deny(self, user_id: int, command: str) -> bool:
|
|
with self._lock:
|
|
command_name = _normalize_command(command)
|
|
if not command_name:
|
|
return False
|
|
|
|
allow_set = self._allow_set(user_id)
|
|
deny_set = self._deny_set(user_id)
|
|
changed = command_name not in deny_set or command_name in allow_set
|
|
deny_set.add(command_name)
|
|
allow_set.discard(command_name)
|
|
if changed:
|
|
self.save()
|
|
return changed
|
|
|
|
def is_allowed(self, user_id: int, command: str) -> bool:
|
|
with self._lock:
|
|
command_name = _normalize_command(command)
|
|
user_key = str(user_id)
|
|
deny_set = self.data["deny"].get(user_key, set())
|
|
if command_name in deny_set:
|
|
return False
|
|
allow_set = self.data["allow"].get(user_key, set())
|
|
return command_name in allow_set
|
|
|
|
def list_allowed(self, user_id: int) -> list[str]:
|
|
with self._lock:
|
|
return sorted(self.data["allow"].get(str(user_id), set()))
|
|
|
|
def list_denied(self, user_id: int) -> list[str]:
|
|
with self._lock:
|
|
return sorted(self.data["deny"].get(str(user_id), set()))
|
|
|
|
def count_allowed_users(self, command: str) -> int:
|
|
with self._lock:
|
|
command_name = _normalize_command(command)
|
|
return sum(1 for commands in self.data["allow"].values() if command_name in commands)
|
|
|
|
def count_denied_users(self, command: str) -> int:
|
|
with self._lock:
|
|
command_name = _normalize_command(command)
|
|
return sum(1 for commands in self.data["deny"].values() if command_name in commands)
|
|
|
|
def stats(self) -> dict[str, int]:
|
|
with self._lock:
|
|
return {
|
|
"allowed_users": len(self.data["allow"]),
|
|
"denied_users": len(self.data["deny"]),
|
|
"allow_rules": sum(len(commands) for commands in self.data["allow"].values()),
|
|
"deny_rules": sum(len(commands) for commands in self.data["deny"].values()),
|
|
}
|