Files
stapler/stapler/registry.py
T
2026-04-21 22:51:35 +02:00

110 lines
4.0 KiB
Python

import logging
import typing
from .data_dir import DataDir
from .page import Page
if typing.TYPE_CHECKING:
from .params import Parameters
class Registry:
__slots__ = [
"data_dir",
"logger",
"pages",
]
HOST_FILE = ".host"
HOST_ONLY_FILE = ".host_only"
TOKEN_FILE = ".token" # noqa: S105
REDIRECT_FILE = ".redirect"
PROXY_FILE = ".proxy"
def __init__(self, params: Parameters) -> None:
self.logger: logging.Logger = logging.getLogger(self.__class__.__name__)
self.pages: dict[str, Page] = {}
self.data_dir = DataDir(params.data_dir)
def load_pages(self) -> None:
self.pages = {}
for path in self.data_dir.list_paths():
self.add(path)
def get_hosts(self) -> list[str]:
return [p.host for p in self.pages.values() if p.host is not None]
def add(self, path: str) -> None:
host = self.data_dir.get_file(path, self.HOST_FILE)
host_only = self.data_dir.get_file(path, self.HOST_ONLY_FILE)
self.pages[path] = Page(
path=path,
with_index=self.data_dir.has_index(path),
host=host if host is not None else host_only,
host_only=host_only is not None,
token_hash=self.data_dir.get_file(path, self.TOKEN_FILE),
redirect=self.data_dir.get_file(path, self.REDIRECT_FILE),
proxy=self.data_dir.get_file(path, self.PROXY_FILE),
)
self.logger.info("Updated %s", self.pages[path])
def set_host(self, path: str, host: str) -> None:
if path in self.pages and (
self.pages[path].host != host or self.pages[path].host_only
):
self.data_dir.set_file(path, self.HOST_FILE, host)
self.data_dir.remove_file(path, self.HOST_ONLY_FILE)
self.pages[path].host = host
self.logger.debug("Updated %s", self.pages[path])
def set_host_only(self, path: str, host: str) -> None:
if path in self.pages and (
self.pages[path].host != host or not self.pages[path].host_only
):
self.data_dir.set_file(path, self.HOST_ONLY_FILE, host)
self.data_dir.remove_file(path, self.HOST_FILE)
self.pages[path].host = host
self.pages[path].host_only = True
self.logger.debug("Updated %s", self.pages[path])
def set_token_hash(self, path: str, token_hash: str) -> None:
if path in self.pages and self.pages[path].token_hash != token_hash:
self.data_dir.set_file(path, self.TOKEN_FILE, token_hash, 0o600)
self.pages[path].token_hash = token_hash
self.logger.debug("Updated %s (token)", self.pages[path])
def set_redirect(self, path: str, redirect: str) -> None:
if path not in self.pages or self.pages[path].redirect != redirect:
self.data_dir.empty(path)
self.data_dir.set_file(path, self.REDIRECT_FILE, redirect)
if path not in self.pages:
self.pages[path] = Page(path)
self.pages[path].redirect = redirect
self.logger.debug("Updated %s", self.pages[path])
def set_proxy(self, path: str, proxy: str) -> None:
if path not in self.pages or self.pages[path].proxy != proxy:
self.data_dir.empty(path)
self.data_dir.set_file(path, self.PROXY_FILE, proxy)
if path not in self.pages:
self.pages[path] = Page(path)
self.pages[path].proxy = proxy
self.logger.debug("Updated %s", self.pages[path])
def remove(self, path: str) -> None:
if path in self.pages:
page = self.pages[path]
del self.pages[path]
self.logger.info("Removed %s", page)
def get_from_path(self, path: str) -> Page | None:
if path in self.pages:
return self.pages[path]
return None
def get_from_host(self, host: str) -> Page | None:
for p in self.pages.values():
if p.host == host:
return p
return None