HomeAssistantVS/custom_components/ingress/config.py

import base64
from homeassistant.components.frontend import EVENT_PANELS_UPDATED  # type: ignore
import os
import re
import time
from typing import TYPE_CHECKING

from .const import LOGGER as _LOGGER

if TYPE_CHECKING:
    from homeassistant.core import HomeAssistant
    from typing import TypedDict
    from yarl import URL

    from .const import WorkMode, RewriteMode

    class Token(TypedDict):
        value: str
        expire: int

    class UserInfo(TypedDict):
        id: str
        name: str | None
        username: str | None


class RewriteCfg:
    mode: "RewriteMode"
    path = ""
    name = ""
    match: str
    replace: str

    def __init__(self, **kwargs):
        self.__dict__.update(
            (k, v) for k, v in kwargs.items() if v is not None and v != getattr(self, k, None)
        )


class IngressCfg:
    mode: "WorkMode"
    name: str
    entry: str
    origin: "URL"
    sub_path = ""
    headers = {}
    cookie_name = "ingress_token"
    _cookie_name_re = None
    expire_time = 3600
    static_token = ""
    rewrites: list[RewriteCfg] = []
    sub_apps: list["IngressCfg"] = []

    def __init__(self, **kwargs):
        self.__dict__.update(
            (k, v) for k, v in kwargs.items() if v is not None and v != getattr(self, k, None)
        )
        self.token: "Token" = {"value": "", "expire": 0}

    def remove_token_from_cookie(self, cookie):
        if (cookie_re := self._cookie_name_re) is None:
            cookie_re = _init_cookie_name_re(self.cookie_name)
            self._cookie_name_re = cookie_re
        cookie = cookie_re.sub("", cookie).strip("; ")
        return USER_TOKEN_COOKIE_RE.sub("", cookie).strip("; ")


class IngressStore:
    """ingress configs."""

    def __init__(self):
        self._configs: dict[str, IngressCfg] = {}
        self._tokens: dict[str, IngressCfg] = {}
        self._user_tokens: dict[str, dict] = {}

    def clear(self):
        self._configs.clear()
        self._tokens.clear()
        self._user_tokens.clear()

    def get(self, name: str) -> IngressCfg | None:
        return self._configs.get(name)

    def del_ingress(self, name: str) -> IngressCfg | None:
        if cfg := self._configs.pop(name, None):
            self._tokens.pop(cfg.token["value"], None)
        return cfg

    def add_ingress(self, cfg: IngressCfg, now: int):
        self.del_ingress(cfg.name)
        cfg.token["value"] = ""
        self.new_token(cfg, now)
        self._configs[cfg.name] = cfg

    def new_token(self, cfg: IngressCfg, now: int) -> str:
        if cfg.static_token:
            token = f"t-{cfg.static_token}"
            old = self._tokens.get(token)
            if old and old is not cfg:
                _LOGGER.error(
                    "static_token conflict with %s, use dynamic for %s!", old.name, cfg.name
                )
                del cfg.static_token
        while not cfg.static_token:
            token = base64.urlsafe_b64encode(os.urandom(33)).decode()
            if token not in self._tokens:
                break
        tkcfg = cfg.token
        self._tokens.pop(tkcfg.get("value"), None)
        self._tokens[token] = cfg
        tkcfg["value"] = token
        tkcfg["expire"] = now + cfg.expire_time
        return token

    def check_token(
        self, hass: "HomeAssistant", token: str, refresh=True
    ) -> tuple[IngressCfg | None, str]:
        cfg = self._tokens.get(token)
        if cfg and refresh:
            # token valid, check refresh
            now = int(time.time())
            if now >= cfg.token["expire"]:
                token = self.new_token(cfg, now)
                hass.bus.async_fire(EVENT_PANELS_UPDATED)
        return cfg, token

    def generate_user_token(self, user_info: dict) -> str:
        """Generate a token associated with user info."""
        now = int(time.time())
        for token in [k for k, v in self._user_tokens.items() if now >= v["expire"]]:
            self._user_tokens.pop(token, None)
        while True:
            token = base64.urlsafe_b64encode(os.urandom(33)).decode()
            if token not in self._user_tokens:
                break
        self._user_tokens[token] = {"user_info": user_info, "expire": now + USER_TOKEN_VALIDITY}
        return token

    def check_user_token(self, token: str) -> "UserInfo | None":
        """Check if a token is valid and return associated user info."""
        entry = self._user_tokens.get(token)
        if not entry:
            return None
        now = int(time.time())
        if now >= entry["expire"]:
            # Token expired, remove it
            self._user_tokens.pop(token, None)
            return None
        # Refresh expiration
        entry["expire"] = now + USER_TOKEN_VALIDITY
        return entry["user_info"]

    def user_cookie_name(self) -> str:
        return USER_TOKEN_COOKIE_NAME

    def cookie_name(self, name: str) -> str:
        cfg = self._configs.get(name)
        return cfg.cookie_name if cfg else IngressCfg.cookie_name


def _init_cookie_name_re(cookie_name):
    return re.compile(rf"(?:^|;\s*){re.escape(cookie_name)}=[^;]*(?=;|$)")


USER_TOKEN_VALIDITY = 300
USER_TOKEN_COOKIE_NAME = "ha_ingress_session"
USER_TOKEN_COOKIE_RE = _init_cookie_name_re(USER_TOKEN_COOKIE_NAME)