forust
bc2e93353a
- Continuous endpoints reject non-list channels (400), null no longer wipes stored list (F-5) - '@'-prefixed legacy channel values normalized on load/save in manager and store (F-6) - _run_loop guarded: unexpected exceptions logged, backoff retry, no silent thread death (F-7) - +8 tests: continuous validation, normalization, loop survival, join-timeout/tombstone (F-8) - 58 tests passing; REVIEW.md updated
430 lines
16 KiB
Python
430 lines
16 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
from copy import deepcopy
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Dict, List, Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ── defaults ──────────────────────────────────────────────────────────
|
|
|
|
GLOBAL_DEFAULTS: Dict[str, Any] = {
|
|
"accounts": [],
|
|
"version": 2,
|
|
}
|
|
|
|
ACCOUNT_DEFAULTS: Dict[str, Any] = {
|
|
"label": "",
|
|
"api_id": None,
|
|
"api_hash": None,
|
|
"channels": {},
|
|
"channel_names": {},
|
|
"scrape_media": True,
|
|
"forwarding_rules": [],
|
|
"continuous_scraping": {
|
|
"enabled": False,
|
|
"interval_minutes": 1,
|
|
"channels": [],
|
|
"run_all_tracked": True,
|
|
},
|
|
}
|
|
|
|
# ── StateStore (thread-safe JSON) ─────────────────────────────────────
|
|
|
|
|
|
class StateStore:
|
|
"""Thread-safe JSON state store with TTL cache."""
|
|
|
|
def __init__(self, path: Path, defaults: Optional[Dict[str, Any]] = None):
|
|
self.path = path
|
|
self.defaults = defaults or {}
|
|
self.lock = threading.RLock()
|
|
self._cache: Optional[Dict[str, Any]] = None
|
|
self._cache_time: float = 0
|
|
self._cache_ttl: float = 1.0
|
|
|
|
def load(self) -> Dict[str, Any]:
|
|
with self.lock:
|
|
now = time.time()
|
|
if self._cache is not None and (now - self._cache_time) < self._cache_ttl:
|
|
return deepcopy(self._cache)
|
|
if not self.path.exists():
|
|
result = deepcopy(self.defaults)
|
|
self._cache = result
|
|
self._cache_time = now
|
|
return deepcopy(result)
|
|
try:
|
|
with self.path.open("r", encoding="utf-8") as handle:
|
|
state: Dict[str, Any] = json.load(handle)
|
|
except (json.JSONDecodeError, OSError):
|
|
result = deepcopy(self.defaults)
|
|
self._cache = result
|
|
self._cache_time = now
|
|
return deepcopy(result)
|
|
result = self._merge_defaults(state)
|
|
self._cache = result
|
|
self._cache_time = now
|
|
return deepcopy(result)
|
|
|
|
def save(self, state: Dict[str, Any]) -> None:
|
|
with self.lock:
|
|
self.path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
|
|
merged = self._merge_defaults(state)
|
|
|
|
# Unique temp file so two writers to the same path cannot clobber
|
|
# each other's in-progress file, plus fsync before atomic rename so
|
|
# a power loss cannot leave an empty/corrupt state file behind.
|
|
fd, tmp_name = tempfile.mkstemp(
|
|
dir=str(self.path.parent),
|
|
prefix=self.path.name + ".tmp-",
|
|
suffix=".tmp",
|
|
)
|
|
tmp_path = Path(tmp_name)
|
|
try:
|
|
with os.fdopen(fd, "w", encoding="utf-8") as handle:
|
|
json.dump(merged, handle, ensure_ascii=False, indent=2)
|
|
handle.write("\n")
|
|
handle.flush()
|
|
os.fsync(handle.fileno())
|
|
# Restrict permissions on the state file to the owning user.
|
|
try:
|
|
os.chmod(tmp_path, 0o600)
|
|
except OSError:
|
|
pass
|
|
os.replace(tmp_path, self.path)
|
|
finally:
|
|
# Ensure no leftover stale temp file if something went wrong.
|
|
if tmp_path.exists():
|
|
try:
|
|
tmp_path.unlink()
|
|
except OSError:
|
|
pass
|
|
|
|
# Best-effort directory fsync for full durability (POSIX only).
|
|
try:
|
|
dir_fd = os.open(str(self.path.parent), os.O_RDONLY)
|
|
try:
|
|
os.fsync(dir_fd)
|
|
finally:
|
|
os.close(dir_fd)
|
|
except OSError:
|
|
pass
|
|
|
|
# After a successful save, best-effort sweep leftover stale tmp
|
|
# files (e.g. from a crashed writer) older than an hour. Throttled
|
|
# to avoid scanning the dir on every save.
|
|
self._sweep_stale_tmp(now=time.time())
|
|
|
|
self._cache = None
|
|
|
|
_STALE_TMP_MAX_AGE = 3600.0 # 1 hour
|
|
_STALE_TMP_SWEEP_INTERVAL = 60.0
|
|
_stale_sweep_last: float = 0.0
|
|
|
|
def _sweep_stale_tmp(self, now: float) -> None:
|
|
if (now - self._stale_sweep_last) < self._STALE_TMP_SWEEP_INTERVAL:
|
|
return
|
|
self._stale_sweep_last = now
|
|
try:
|
|
cutoff = now - self._STALE_TMP_MAX_AGE
|
|
for stale in self.path.parent.glob(self.path.name + ".tmp-*.tmp"):
|
|
try:
|
|
if stale.stat().st_mtime < cutoff:
|
|
stale.unlink()
|
|
except OSError:
|
|
pass
|
|
except OSError:
|
|
pass
|
|
|
|
def update(self, mutator: Callable[[Dict[str, Any]], None]) -> Dict[str, Any]:
|
|
with self.lock:
|
|
state = self.load()
|
|
mutator(state)
|
|
self.save(state)
|
|
return state
|
|
|
|
def continuous_config(self) -> Dict[str, Any]:
|
|
state = self.load()
|
|
return deepcopy(state.get("continuous_scraping") or ACCOUNT_DEFAULTS["continuous_scraping"])
|
|
|
|
def save_continuous_config(self, config: Dict[str, Any]) -> Dict[str, Any]:
|
|
def mutate(state: Dict[str, Any]) -> None:
|
|
state["continuous_scraping"] = {
|
|
"enabled": bool(config.get("enabled", True)),
|
|
"interval_minutes": max(1, int(config.get("interval_minutes", 1) or 1)),
|
|
"channels": [
|
|
str(item).strip().lstrip("@")
|
|
for item in config.get("channels", [])
|
|
if str(item).strip()
|
|
],
|
|
"run_all_tracked": bool(config.get("run_all_tracked", True)),
|
|
}
|
|
|
|
return self.update(mutate)["continuous_scraping"]
|
|
|
|
def _merge_defaults(self, state: Dict[str, Any]) -> Dict[str, Any]:
|
|
merged = deepcopy(self.defaults)
|
|
for key, value in state.items():
|
|
if isinstance(value, dict) and isinstance(merged.get(key), dict):
|
|
nested = deepcopy(merged[key])
|
|
nested.update(value)
|
|
merged[key] = nested
|
|
else:
|
|
merged[key] = value
|
|
return merged
|
|
|
|
|
|
# ── global state helpers ──────────────────────────────────────────────
|
|
|
|
_GLOBAL_STORE: Optional[StateStore] = None
|
|
|
|
|
|
def get_global_store(data_dir: Path) -> StateStore:
|
|
"""Return the process-wide global state store singleton.
|
|
|
|
The store is cached in a module-global and the SAME instance is returned
|
|
for the same process, regardless of how many times this is called with the
|
|
same (or any) data_dir. webui_server's ``STATE_STORE`` should delegate to
|
|
this function so there is exactly one authoritative global store per
|
|
process rather than a second, potentially divergent instance.
|
|
|
|
Backward-compat note: callers that cache their own ``_GLOBAL_STORE=None``
|
|
sentinel (e.g. tests resetting state between cases) still work because we
|
|
re-create the singleton lazily on first call.
|
|
"""
|
|
global _GLOBAL_STORE
|
|
if _GLOBAL_STORE is None:
|
|
_GLOBAL_STORE = StateStore(data_dir / "state.json", defaults=GLOBAL_DEFAULTS)
|
|
return _GLOBAL_STORE
|
|
|
|
|
|
def load_global(data_dir: Path) -> Dict[str, Any]:
|
|
return get_global_store(data_dir).load()
|
|
|
|
|
|
def save_global(data_dir: Path, state: Dict[str, Any]) -> None:
|
|
get_global_store(data_dir).save(state)
|
|
|
|
|
|
def list_accounts(data_dir: Path) -> List[str]:
|
|
return list(load_global(data_dir).get("accounts", []))
|
|
|
|
|
|
def account_exists(data_dir: Path, account_id: str) -> bool:
|
|
return account_id in list_accounts(data_dir)
|
|
|
|
|
|
# ── per-account state helpers ─────────────────────────────────────────
|
|
|
|
_ACCOUNT_STORES: Dict[str, StateStore] = {}
|
|
_account_stores_lock = threading.Lock()
|
|
|
|
|
|
def get_account_store(data_dir: Path, account_id: str) -> StateStore:
|
|
global _ACCOUNT_STORES
|
|
with _account_stores_lock:
|
|
key = f"{data_dir}:{account_id}"
|
|
if key not in _ACCOUNT_STORES:
|
|
store = StateStore(
|
|
data_dir / "accounts" / account_id / "state.json",
|
|
defaults=ACCOUNT_DEFAULTS,
|
|
)
|
|
_ACCOUNT_STORES[key] = store
|
|
return _ACCOUNT_STORES[key]
|
|
|
|
|
|
def load_account(data_dir: Path, account_id: str) -> Dict[str, Any]:
|
|
return get_account_store(data_dir, account_id).load()
|
|
|
|
|
|
def save_account(data_dir: Path, account_id: str, state: Dict[str, Any]) -> None:
|
|
get_account_store(data_dir, account_id).save(state)
|
|
|
|
|
|
def account_data_dir(data_dir: Path, account_id: str) -> Path:
|
|
"""Per-account directory for channel DBs, media, exports."""
|
|
return data_dir / "accounts" / account_id
|
|
|
|
|
|
def account_session_path(session_dir: Path, account_id: str) -> str:
|
|
"""Per-account Telethon session file. Always under session/<id>.session."""
|
|
return str(session_dir / f"{account_id}.session")
|
|
|
|
|
|
# ── MIGRATION (with data copy) ─────────────────────────────────────────
|
|
|
|
|
|
def _is_valid_channel_id(channel_id: str) -> bool:
|
|
"""Path-safe channel id check.
|
|
|
|
Mirrors webui_server.normalize_channel_id's validation: reject entries
|
|
containing ``/`` or ``\\``, control chars, ``.``/``..``, empty; keep
|
|
numbers and plain names. Defined locally (not imported from webui_server,
|
|
which would be circular) so it can be shared by migration.
|
|
"""
|
|
channel_id = str(channel_id or "").strip()
|
|
if (
|
|
not channel_id
|
|
or "/" in channel_id
|
|
or "\\" in channel_id
|
|
or channel_id in {".", ".."}
|
|
or any(ord(ch) < 32 for ch in channel_id)
|
|
):
|
|
return False
|
|
return True
|
|
|
|
|
|
def _clean_continuous_channels(channels: Any) -> List[str]:
|
|
"""Normalize/drop invalid continuous-scraping channel entries during
|
|
migration. Mirrors webui_server.clean_continuous_channels: strips a
|
|
leading ``@``, keeps numbers/names, and drops unsafe entries so they can
|
|
never become a path-traversal vector.
|
|
"""
|
|
cleaned: List[str] = []
|
|
if not isinstance(channels, list):
|
|
return cleaned
|
|
for item in channels:
|
|
cleaned.append(str(item).strip().lstrip("@"))
|
|
return [c for c in cleaned if _is_valid_channel_id(c)]
|
|
|
|
|
|
def _copy_channel_data(src_root: Path, dst_root: Path, channel_id: str) -> None:
|
|
"""Copy a single channel's DB + media from src_root to dst_root."""
|
|
src_ch = src_root / channel_id
|
|
dst_ch = dst_root / channel_id
|
|
if not src_ch.exists():
|
|
return
|
|
dst_ch.mkdir(parents=True, exist_ok=True, mode=0o700)
|
|
|
|
# SQLite DB
|
|
db_name = f"{channel_id}.db"
|
|
src_db = src_ch / db_name
|
|
dst_db = dst_ch / db_name
|
|
if src_db.exists() and not dst_db.exists():
|
|
shutil.copy2(src_db, dst_db)
|
|
logger.info(" copied %s", db_name)
|
|
|
|
# Media directory
|
|
src_media = src_ch / "media"
|
|
dst_media = dst_ch / "media"
|
|
if src_media.exists() and src_media.is_dir():
|
|
dst_media.mkdir(parents=True, exist_ok=True, mode=0o700)
|
|
for item in src_media.iterdir():
|
|
if item.is_file():
|
|
dst_file = dst_media / item.name
|
|
if not dst_file.exists():
|
|
shutil.copy2(item, dst_file)
|
|
|
|
# Exports (csv, json)
|
|
for ext in (".csv", ".json"):
|
|
for f in src_ch.glob(f"*{ext}"):
|
|
dst_f = dst_ch / f.name
|
|
if not dst_f.exists():
|
|
shutil.copy2(f, dst_f)
|
|
|
|
|
|
def migrate_legacy_state(data_dir: Path, session_dir: Path) -> bool:
|
|
"""
|
|
Run once at startup. Detects legacy data/state.json that has api_id
|
|
and migrates to multi-account format, COPYING channel DBs/media from
|
|
data/<channel_id>/ → data/accounts/default/<channel_id>/ .
|
|
|
|
Returns True if migration was performed.
|
|
"""
|
|
state_path = data_dir / "state.json"
|
|
if not state_path.exists():
|
|
return False
|
|
|
|
try:
|
|
with state_path.open("r", encoding="utf-8") as f:
|
|
raw: Dict[str, Any] = json.load(f)
|
|
except (json.JSONDecodeError, OSError):
|
|
return False
|
|
|
|
# Already migrated or fresh install with no legacy api_id
|
|
if raw.get("version") == 2 or not raw.get("api_id"):
|
|
return False
|
|
|
|
logger.info(
|
|
"=== Legacy state detected — migrating to multi-account format ==="
|
|
)
|
|
logger.info("This may take a while if you have many channels with media.")
|
|
|
|
# ── 1. Create per-account state for "default" ──────────────────────
|
|
acc_dir = data_dir / "accounts" / "default"
|
|
acc_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
|
|
|
|
# Normalize/drop invalid continuous-scraping channel entries during
|
|
# migration so unsafe values can never become a path-traversal vector.
|
|
continuous_cfg = raw.get(
|
|
"continuous_scraping",
|
|
{
|
|
"enabled": True,
|
|
"interval_minutes": 1,
|
|
"channels": [],
|
|
"run_all_tracked": True,
|
|
},
|
|
)
|
|
if not isinstance(continuous_cfg, dict):
|
|
continuous_cfg = {}
|
|
continuous_cfg["channels"] = _clean_continuous_channels(
|
|
continuous_cfg.get("channels")
|
|
)
|
|
|
|
acc_state = {
|
|
"label": "Default",
|
|
"api_id": raw.get("api_id"),
|
|
"api_hash": raw.get("api_hash"),
|
|
"channels": raw.get("channels", {}),
|
|
"channel_names": raw.get("channel_names", {}),
|
|
"scrape_media": raw.get("scrape_media", True),
|
|
"forwarding_rules": raw.get("forwarding_rules", []),
|
|
"continuous_scraping": continuous_cfg,
|
|
}
|
|
|
|
acc_state_path = acc_dir / "state.json"
|
|
with open(acc_state_path, "w", encoding="utf-8") as f:
|
|
json.dump(acc_state, f, ensure_ascii=False, indent=2)
|
|
f.write("\n")
|
|
logger.info(" ✓ data/accounts/default/state.json created")
|
|
|
|
# ── 2. Copy channel data ───────────────────────────────────────────
|
|
legacy_channels = list(raw.get("channels", {}).keys())
|
|
if legacy_channels:
|
|
logger.info(
|
|
" Copying %d channel(s) from data/<id>/ → data/accounts/default/<id>/ ...",
|
|
len(legacy_channels),
|
|
)
|
|
for i, channel_id in enumerate(legacy_channels, 1):
|
|
logger.info(" [%d/%d] channel %s", i, len(legacy_channels), channel_id)
|
|
_copy_channel_data(data_dir, acc_dir, channel_id)
|
|
logger.info(" ✓ Channel data copied")
|
|
else:
|
|
logger.info(" (no tracked channels to copy)")
|
|
|
|
# ── 3. Rewrite global state.json ───────────────────────────────────
|
|
global_state = {
|
|
"accounts": ["default"],
|
|
"version": 2,
|
|
}
|
|
with open(state_path, "w", encoding="utf-8") as f:
|
|
json.dump(global_state, f, ensure_ascii=False, indent=2)
|
|
f.write("\n")
|
|
logger.info(" ✓ data/state.json rewritten (accounts list only)")
|
|
|
|
# ── 4. Rename legacy session file ─────────────────────────────────
|
|
legacy_session = session_dir / "session.session"
|
|
if legacy_session.exists():
|
|
new_session = session_dir / "default.session"
|
|
if not new_session.exists():
|
|
legacy_session.rename(new_session)
|
|
logger.info(" ✓ session/session.session → session/default.session")
|
|
|
|
logger.info("=== Migration complete ===")
|
|
return True
|