|
# retoor <retoor@molodetz.nl>
|
|
|
|
import json
|
|
import logging
|
|
import threading
|
|
from typing import Optional
|
|
from devplacepy.services.tunnel.config import TUNNEL_SUBDOMAIN_REGISTRY_PATH
|
|
|
|
logger = logging.getLogger(__name__)
|
|
_lock = threading.Lock()
|
|
|
|
|
|
class SubdomainRegistry:
|
|
def __init__(self) -> None:
|
|
self._sub_to_session: dict[str, str] = {}
|
|
self._session_to_sub: dict[str, str] = {}
|
|
self._load()
|
|
|
|
def _load(self) -> None:
|
|
path = TUNNEL_SUBDOMAIN_REGISTRY_PATH
|
|
if not path.exists():
|
|
return
|
|
try:
|
|
raw = path.read_text()
|
|
data = json.loads(raw)
|
|
with _lock:
|
|
self._sub_to_session = data.get("sub_to_session", {})
|
|
self._session_to_sub = data.get("session_to_sub", {})
|
|
except (json.JSONDecodeError, OSError):
|
|
logger.exception("Failed to load subdomain registry")
|
|
|
|
def _save(self) -> None:
|
|
path = TUNNEL_SUBDOMAIN_REGISTRY_PATH
|
|
try:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
data = {
|
|
"sub_to_session": self._sub_to_session,
|
|
"session_to_sub": self._session_to_sub,
|
|
}
|
|
path.write_text(json.dumps(data, indent=2))
|
|
except OSError:
|
|
logger.exception("Failed to save subdomain registry")
|
|
|
|
def claim(self, subdomain: str, session_id: str) -> bool:
|
|
with _lock:
|
|
if subdomain in self._sub_to_session:
|
|
return False
|
|
self._sub_to_session[subdomain] = session_id
|
|
self._session_to_sub[session_id] = subdomain
|
|
self._save()
|
|
return True
|
|
|
|
def release(self, session_id: str) -> Optional[str]:
|
|
with _lock:
|
|
sub = self._session_to_sub.pop(session_id, None)
|
|
if sub:
|
|
self._sub_to_session.pop(sub, None)
|
|
self._save()
|
|
return sub
|
|
|
|
def resolve(self, subdomain: str) -> Optional[str]:
|
|
with _lock:
|
|
return self._sub_to_session.get(subdomain)
|
|
|
|
def session_subdomain(self, session_id: str) -> Optional[str]:
|
|
with _lock:
|
|
return self._session_to_sub.get(session_id)
|
|
|
|
def all_subdomains(self) -> dict[str, str]:
|
|
with _lock:
|
|
return dict(self._sub_to_session)
|
|
|
|
def clear(self) -> None:
|
|
with _lock:
|
|
self._sub_to_session.clear()
|
|
self._session_to_sub.clear()
|
|
self._save()
|
|
|
|
|
|
_registry: Optional[SubdomainRegistry] = None
|
|
|
|
|
|
def get_registry() -> SubdomainRegistry:
|
|
global _registry
|
|
if _registry is None:
|
|
_registry = SubdomainRegistry()
|
|
return _registry
|
|
|
|
|
|
def reset_registry() -> None:
|
|
global _registry
|
|
_registry = None
|