douyin-sparkflow/DouYinSparkFlow/utils/config.py

import json
import logging
import os
import secrets
import sys
from copy import deepcopy
from enum import Enum
from pathlib import Path

from utils.logger import setup_logger


logger = setup_logger(level=logging.DEBUG)

DEBUG = False
CONFIGFILE = "config.json"
USERDATAFILE = "usersData.json"
APPSETTINGSFILE = "webui_settings.json"

DEFAULT_CONFIG = {
    "multiTask": True,
    "taskCount": 5,
    "proxyAddress": "",
    "messageTemplate": "【AI续火花】",
    "saveDebugArtifacts": False,
    "useProtocolSender": True,
    "protocolDryRun": False,
    "browserSenderAccounts": [],
    "sendStrategy": {
        "shuffleTargets": True,
        "accountStartDelaySecondsMin": 0,
        "accountStartDelaySecondsMax": 20,
        "messageIntervalSecondsMin": 18,
        "messageIntervalSecondsMax": 45,
        "messageVariants": [],
    },
    "dailySendWindow": {
        "enabled": False,
        "startHour": 10,
        "endHour": 18,
        "scheduleIntervalMinutes": 10,
    },
    "hitokotoTypes": ["文学", "影视", "诗词", "哲学"],
    "happyNewYear": {
        "enabled": False,
        "messageTemplate": "【[data]|[data_lunar]】\n[API]",
    },
}

DEFAULT_APP_SETTINGS = {
    "admin_username": "admin",
    "admin_password_hash": "",
    "session_secret": "",
    "session_max_age_seconds": 8 * 60 * 60,
    "compose_root": "",
    "ui_host": "0.0.0.0",
    "ui_port": 8787,
    "login_poll_interval_seconds": 1,
    "ops_log_file": "/app/logs/douyin-sparkflow.log",
    "proxy_refresh_script": "/opt/douyin-sparkflow/refresh_proxy.sh",
    "local_login_helper_url": "http://127.0.0.1:18765",
    "login_desktop_api_url": "http://127.0.0.1:18090",
    "server_host": "",
    "server_username": "",
    "server_password": "",
}

config = None
userData = None
appSettings = None


class Environment(Enum):
    GITHUBACTION = "GITHUB_ACTION"
    LOCAL = "LOCAL"
    PACKED = "PACKED"

    def __str__(self):
        return self.value


def get_environment():
    if getattr(sys, "frozen", False) and hasattr(sys, "_MEIPASS"):
        return Environment.PACKED
    if os.getenv("GITHUB_ACTIONS") == "true":
        return Environment.GITHUBACTION
    return Environment.LOCAL


def repo_root():
    return Path(__file__).resolve().parents[1]


def _runtime_root():
    env = get_environment()
    if env == Environment.PACKED:
        return Path(sys.executable).resolve().parent
    return repo_root()


def config_path():
    return _runtime_root() / CONFIGFILE


def users_data_path():
    return _runtime_root() / USERDATAFILE


def app_settings_path():
    return _runtime_root() / APPSETTINGSFILE


def default_compose_root():
    root = repo_root()
    parent = root.parent
    if (parent / "docker-compose.yml").exists():
        return str(parent)
    return str(root)


def _merge_defaults(data, defaults):
    merged = deepcopy(defaults)
    for key, value in data.items():
        if key in merged and isinstance(merged[key], dict) and isinstance(value, dict):
            merged[key].update(value)
        else:
            merged[key] = value
    return merged


def _load_json_file(path, defaults=None):
    if not path.exists():
        if defaults is None:
            raise FileNotFoundError(path)
        path.write_text(json.dumps(defaults, ensure_ascii=False, indent=2), encoding="utf-8")
        return deepcopy(defaults)

    text = path.read_text(encoding="utf-8")
    if not text.strip():
        return deepcopy(defaults) if defaults is not None else None

    data = json.loads(text)
    if defaults is None:
        return data
    # _merge_defaults only works with dicts; for list-shaped data
    # (e.g. usersData.json) just return the parsed data as-is.
    if not isinstance(data, dict) or not isinstance(defaults, dict):
        return data
    return _merge_defaults(data, defaults)


def _save_json_file(path, data):
    path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")


def get_config(force_reload=False):
    global config
    if config is None or force_reload:
        config = _load_json_file(config_path(), DEFAULT_CONFIG)
    return deepcopy(config)


def save_config(new_config):
    global config
    config = _merge_defaults(new_config, DEFAULT_CONFIG)
    _save_json_file(config_path(), config)
    return deepcopy(config)


def get_userData(force_reload=False):
    global userData
    if userData is not None and not force_reload:
        return deepcopy(userData)

    env = get_environment()
    if env == Environment.GITHUBACTION:
        raw = os.getenv("USER_DATA", "")
        if not raw:
            logger.error("Environment variable USER_DATA is not set")
            raise RuntimeError("USER_DATA is required in GITHUB_ACTIONS mode")
        userData = json.loads(raw)
    else:
        userData = _load_json_file(users_data_path(), [])

    return deepcopy(userData)


def save_userData(accounts):
    global userData
    normalized = list(accounts)
    userData = normalized
    _save_json_file(users_data_path(), normalized)
    return deepcopy(userData)


def normalize_unique_id(unique_id):
    if not unique_id:
        return ""
    digits = "".join(ch for ch in str(unique_id) if ch.isdigit())
    return digits or str(unique_id).strip()


def upsert_user_account(unique_id, username, cookies, targets, extra=None):
    unique_id = normalize_unique_id(unique_id)
    accounts = get_userData(force_reload=True)
    payload = {
        "unique_id": unique_id,
        "username": username,
        "cookies": cookies,
        "targets": list(targets),
    }
    if extra:
        payload.update(extra)

    for account in accounts:
        if normalize_unique_id(account.get("unique_id")) == unique_id:
            if "enabled" not in payload:
                payload["enabled"] = account.get("enabled", True)
            account.update(payload)
            save_userData(accounts)
            return payload

    if "enabled" not in payload:
        payload["enabled"] = True
    accounts.append(payload)
    save_userData(accounts)
    return payload


def delete_user_account(unique_id):
    normalized_id = normalize_unique_id(unique_id)
    accounts = get_userData(force_reload=True)
    remaining = [item for item in accounts if normalize_unique_id(item.get("unique_id")) != normalized_id]
    removed = len(accounts) != len(remaining)
    if removed:
        save_userData(remaining)
    return removed


def get_app_settings(force_reload=False):
    global appSettings
    if appSettings is None or force_reload:
        appSettings = _load_json_file(app_settings_path(), DEFAULT_APP_SETTINGS)
        if not appSettings.get("session_secret"):
            appSettings["session_secret"] = secrets.token_urlsafe(32)
        if not appSettings.get("compose_root"):
            appSettings["compose_root"] = default_compose_root()
        _save_json_file(app_settings_path(), appSettings)
    return deepcopy(appSettings)


def save_app_settings(new_settings):
    global appSettings
    appSettings = _merge_defaults(new_settings, DEFAULT_APP_SETTINGS)
    if not appSettings.get("session_secret"):
        appSettings["session_secret"] = secrets.token_urlsafe(32)
    if not appSettings.get("compose_root"):
        appSettings["compose_root"] = default_compose_root()
    _save_json_file(app_settings_path(), appSettings)
    return deepcopy(appSettings)