forust
6e3966aeb5
- GlobalStateStore for account registry in global state.json - AccountStateStore for per-account settings (channels, credentials, continuous config) - migrate_legacy_state() copies legacy data/ to data/accounts/default/ with progress logging; renames session.session to default.session - Helpers: account_data_dir(), account_session_path(), list_accounts(), load_account(), save_account(), get_account_store(), get_global_store() - StateStore accepts defaults parameter for flexible default values
314 lines
11 KiB
Python
314 lines
11 KiB
Python
import json
|
|
import logging
|
|
import shutil
|
|
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": True,
|
|
"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 dict(self._cache)
|
|
if not self.path.exists():
|
|
result = deepcopy(self.defaults)
|
|
self._cache = result
|
|
self._cache_time = now
|
|
return 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 result
|
|
result = self._merge_defaults(state)
|
|
self._cache = result
|
|
self._cache_time = now
|
|
return result
|
|
|
|
def save(self, state: Dict[str, Any]) -> None:
|
|
with self.lock:
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp_path = self.path.with_suffix(self.path.suffix + ".tmp")
|
|
merged = self._merge_defaults(state)
|
|
with tmp_path.open("w", encoding="utf-8") as handle:
|
|
json.dump(merged, handle, ensure_ascii=False, indent=2)
|
|
handle.write("\n")
|
|
tmp_path.replace(self.path)
|
|
self._cache = None
|
|
|
|
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()
|
|
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:
|
|
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 _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)
|
|
|
|
# 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)
|
|
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)
|
|
|
|
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": raw.get(
|
|
"continuous_scraping",
|
|
{
|
|
"enabled": True,
|
|
"interval_minutes": 1,
|
|
"channels": [],
|
|
"run_all_tracked": True,
|
|
},
|
|
),
|
|
}
|
|
|
|
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
|