DDNS/ddns/__main__.py

# -*- coding:utf-8 -*-
"""
DDNS
@author: NewFuture, rufengsuixing
"""

import sys
from io import TextIOWrapper
from logging import getLogger
from subprocess import check_output

from . import ip
from .__init__ import __description__, __version__, build_date
from .cache import Cache
from .config import Config, load_configs  # noqa: F401
from .provider import SimpleProvider, get_provider_class  # noqa: F401

logger = getLogger()


class UpdateCancelled(Exception):
    """Raised when a caller cancels an in-progress DDNS update."""


def _raise_if_cancelled(cancelled):
    # type: (object | None) -> None
    if cancelled is not None and cancelled():
        raise UpdateCancelled("DDNS update cancelled.")


def _get_ip_from_rule(ip_type, rule):
    """
    Resolve an IP address from a single rule.
    """
    rule_text = str(rule)
    if rule_text.isdigit():
        return getattr(ip, "local_v" + ip_type)(rule)
    if rule_text.startswith("cmd:"):
        return str(check_output(rule_text[4:]).strip().decode("utf-8"))
    if rule_text.startswith("shell:"):
        return str(check_output(rule_text[6:], shell=True).strip().decode("utf-8"))
    if rule_text.startswith("url:"):
        return getattr(ip, "public_v" + ip_type)(rule_text[4:])
    if rule_text.startswith("regex:"):
        return getattr(ip, "regex_v" + ip_type)(rule_text[6:])
    return getattr(ip, rule_text + "_v" + ip_type)()


def get_ip(ip_type, rules, cancelled=None):
    """
    get IP address
    """
    if rules is False:  # disabled
        return False
    for rule in rules:
        _raise_if_cancelled(cancelled)
        try:
            logger.debug("get_ip:(%s, %s)", ip_type, rule)
            result = _get_ip_from_rule(ip_type, rule)
        except Exception as e:
            logger.error("Failed to get %s address: %s", ip_type, e)
            continue

        _raise_if_cancelled(cancelled)
        if result:
            return result

        logger.debug("Rule %s did not return a usable IPv%s address", rule, ip_type)
    return None


def update_ip(dns, cache, index_rule, domains, record_type, config, cancelled=None):
    # type: (SimpleProvider, Cache | None, list[str]|bool, list[str], str, Config, object | None) -> bool | None
    """
    更新IP并变更DNS记录
    """
    _raise_if_cancelled(cancelled)
    if not domains or index_rule is False:
        return None

    ip_type = "4" if record_type == "A" else "6"
    address = get_ip(ip_type, index_rule, cancelled=cancelled)
    if not address:
        logger.error("Fail to get %s address!", ip_type)
        return False

    update_success = True

    for domain in domains:
        _raise_if_cancelled(cancelled)
        domain = domain.lower()
        cache_key = "{}:{}".format(domain, record_type)
        if cache and cache.get(cache_key) == address:
            logger.info("%s[%s] address not changed, using cache: %s", domain, record_type, address)
            continue

        try:
            result = dns.set_record(
                domain, address, record_type=record_type, ttl=config.ttl, line=config.line, **config.extra
            )
            if result:
                logger.warning("set %s[IPv%s]: %s successfully.", domain, ip_type, address)
                if isinstance(cache, dict):
                    cache[cache_key] = address
            else:
                update_success = False
                logger.error("Failed to update %s record for %s", record_type, domain)
        except Exception as e:
            update_success = False
            logger.exception("Failed to update %s record for %s: %s", record_type, domain, e)
    _raise_if_cancelled(cancelled)
    return update_success


def run(config, cancelled=None):
    # type: (Config, object | None) -> bool
    """
    Run the DDNS update process
    """
    # 设置IP模块的SSL验证配置
    ip.ssl_verify = config.ssl

    # dns provider class
    provider_class = get_provider_class(config.dns)
    dns = provider_class(
        config.id, config.token, endpoint=config.endpoint, logger=logger, proxy=config.proxy, ssl=config.ssl
    )
    cache = Cache.new(config.cache, config.md5(), logger, config.cache_max_age)
    result4 = update_ip(dns, cache, config.index4, config.ipv4, "A", config, cancelled=cancelled)
    result6 = update_ip(dns, cache, config.index6, config.ipv6, "AAAA", config, cancelled=cancelled)
    return result4 is not False and result6 is not False


def main():
    stdout = sys.stdout  # pythonw 模式无 stdout
    mcp_mode = len(sys.argv) > 1 and sys.argv[1] == "mcp"
    if stdout and stdout.encoding and stdout.encoding.lower() != "utf-8" and hasattr(stdout, "buffer"):
        # 兼容windows 和部分ASCII编码的老旧系统
        sys.stdout = TextIOWrapper(sys.stdout.buffer, encoding="utf-8")
        sys.stderr = TextIOWrapper(sys.stderr.buffer, encoding="utf-8")

    # Windows 下输出一个空行
    if stdout and sys.platform.startswith("win") and not mcp_mode:
        stdout.write("\r\n")

    logger.name = "ddns"

    # 使用多配置加载器，它会自动处理单个和多个配置
    configs = load_configs(__description__, __version__, build_date)

    if len(configs) == 1:
        # 单个配置，使用原有逻辑（向后兼容）
        config = configs[0]
        success = run(config)
        if not success:
            sys.exit(1)
    else:
        # 多个配置，使用新的批处理逻辑
        overall_success = True
        for i, config in enumerate(configs):
            # 如果log_level有值则设置setLevel
            if hasattr(config, "log_level") and config.log_level:
                logger.setLevel(config.log_level)
            logger.info("Running configuration %d/%d", i + 1, len(configs))
            # 记录当前provider
            logger.info("Using DNS provider: %s", config.dns)
            success = run(config)
            if not success:
                overall_success = False
                logger.error("Configuration %d failed", i + 1)
            else:
                logger.info("Configuration %d completed successfully", i + 1)

        if not overall_success:
            logger.error("Some configurations failed")
            sys.exit(1)
        else:
            logger.info("All configurations completed successfully")


if __name__ == "__main__":
    main()