558 lines
23 KiB
Python
558 lines
23 KiB
Python
# -*- encoding=utf8 -*-
|
||
"""
|
||
Task dispatcher configuration.
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import re
|
||
from typing import Any, Dict, List, Optional, Tuple
|
||
|
||
try:
|
||
import yaml
|
||
except ImportError: # pragma: no cover - exercised only when dependency is absent.
|
||
yaml = None
|
||
|
||
|
||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
CONFIG_DIR = os.path.join(BASE_DIR, "config")
|
||
CONFIG_PATH_ENV = os.environ.get("AUTOOL_DISPATCHER_CONFIG")
|
||
CONFIG_PATH = CONFIG_PATH_ENV or os.path.join(CONFIG_DIR, "config.json")
|
||
CONFIG_YAML_PATH = os.path.join(BASE_DIR, "config.yaml")
|
||
CONFIG_EXAMPLE_PATH = os.path.join(CONFIG_DIR, "config.example.json")
|
||
CONFIG_EXAMPLE_YAML_PATH = os.path.join(BASE_DIR, "config.example.yaml")
|
||
CURRENT_ENV_PATH = os.path.join(CONFIG_DIR, "current_env.txt")
|
||
LOCAL_CONFIG_PATH = os.path.join(CONFIG_DIR, "local.json")
|
||
ENV_CONFIG_PATHS = {
|
||
"prod": os.path.join(CONFIG_DIR, "prod.json"),
|
||
"test": os.path.join(CONFIG_DIR, "test.json"),
|
||
}
|
||
ENV_ALIASES = {
|
||
"": "prod",
|
||
"prod": "prod",
|
||
"production": "prod",
|
||
"main": "prod",
|
||
"master": "prod",
|
||
"test": "test",
|
||
"testing": "test",
|
||
}
|
||
|
||
|
||
def _load_config_payload(path: str, *, required: bool) -> Any:
|
||
if not os.path.exists(path):
|
||
if required:
|
||
raise FileNotFoundError(f"config file not found: {path}")
|
||
return {}
|
||
|
||
with open(path, "r", encoding="utf-8") as handle:
|
||
if path.endswith((".yaml", ".yml")):
|
||
if yaml is None:
|
||
raise ImportError("PyYAML is required to read YAML config files. Install: pip install PyYAML")
|
||
payload = yaml.safe_load(handle) or {}
|
||
else:
|
||
payload = json.load(handle)
|
||
|
||
return payload
|
||
|
||
|
||
def _load_config_dict(path: str, *, required: bool) -> Dict[str, Any]:
|
||
payload = _load_config_payload(path, required=required)
|
||
|
||
if not isinstance(payload, dict):
|
||
raise ValueError(f"config payload must be an object: {path}")
|
||
return payload
|
||
|
||
|
||
def _load_json_dict(path: str, *, required: bool) -> Dict[str, Any]:
|
||
return _load_config_dict(path, required=required)
|
||
|
||
|
||
def _merge_dict(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]:
|
||
merged = dict(base)
|
||
for key, value in override.items():
|
||
current = merged.get(key)
|
||
if isinstance(current, dict) and isinstance(value, dict):
|
||
merged[key] = _merge_dict(current, value)
|
||
continue
|
||
merged[key] = value
|
||
return merged
|
||
|
||
|
||
def _read_current_env() -> str:
|
||
if not os.path.exists(CURRENT_ENV_PATH):
|
||
return "prod"
|
||
with open(CURRENT_ENV_PATH, "r", encoding="utf-8") as handle:
|
||
return handle.read().strip().lower()
|
||
|
||
|
||
def normalize_env_name(env_name: Optional[str] = None) -> str:
|
||
raw = str("prod" if env_name is None else env_name).strip().lower()
|
||
normalized = ENV_ALIASES.get(raw)
|
||
if normalized:
|
||
return normalized
|
||
raise ValueError(f"unsupported config env: {env_name or raw}")
|
||
|
||
|
||
def get_runtime_env_name() -> str:
|
||
return CONFIG_ENV
|
||
|
||
|
||
def _load_legacy_layered_config(env_name: Optional[str] = None) -> Dict[str, Any]:
|
||
normalized_env = normalize_env_name(env_name)
|
||
config = _load_config_dict(ENV_CONFIG_PATHS[normalized_env], required=True)
|
||
config = _merge_dict(config, _load_config_dict(LOCAL_CONFIG_PATH, required=False))
|
||
return config
|
||
|
||
|
||
# 嵌套小写 → 扁平大写键名映射
|
||
_NESTED_TO_FLAT_MAP = {
|
||
("environment", "name"): "CONFIG_ENV",
|
||
("instance_name",): "INSTANCE_NAME",
|
||
("redis", "host"): "REDIS_HOST",
|
||
("redis", "port"): "REDIS_PORT",
|
||
("redis", "db"): "REDIS_DB",
|
||
("redis", "channel_namespace"): "CHANNEL_NAMESPACE",
|
||
("redis", "max_connections"): "REDIS_MAX_CONNECTIONS_DISPATCHER",
|
||
("tasks", "max_retry_count"): "MAX_RETRY_COUNT",
|
||
("tasks", "csv_file"): "TASK_CSV_FILE",
|
||
("worker", "stale_timeout"): "WORKER_STALE_TIMEOUT",
|
||
("worker", "reinit_timeout"): "WORKER_REINIT_TIMEOUT",
|
||
("worker", "init_monitor_window"): "WORKER_INIT_MONITOR_WINDOW",
|
||
("worker", "action_max_parallel"): "WORKER_ACTION_MAX_PARALLEL",
|
||
("worker", "managed_only_can_dispatch"): "ONLY_MANAGED_WORKERS_CAN_DISPATCH",
|
||
("worker", "allowed_ip_patterns"): "WORKER_ALLOWED_IP_PATTERNS",
|
||
("ssh", "connect_timeout"): "SSH_CONNECT_TIMEOUT_SECONDS",
|
||
("ssh", "action_timeout"): "SSH_ACTION_TIMEOUT_SECONDS",
|
||
("ssh", "default_user"): "SSH_DEFAULT_USER",
|
||
("ssh", "default_password"): "SSH_DEFAULT_PASSWORD",
|
||
("ssh", "default_port"): "SSH_DEFAULT_PORT",
|
||
("ssh", "psexec_session_id"): "PSEXEC_DEFAULT_SESSION_ID",
|
||
("git", "repo_url"): "AUTOOL_GIT_REPO_URL",
|
||
("mumu", "manager_path"): "MUMU_MANAGER_PATH",
|
||
("mumu", "vm_index"): "MUMU_VM_INDEX",
|
||
("mumu", "max_vm_index"): "MUMU_MAX_VM_INDEX",
|
||
("mumu", "restart_settle_seconds"): "MUMU_RESTART_SETTLE_SECONDS",
|
||
("mumu", "recover_script_path"): "MUMU_RECOVER_SCRIPT_PATH",
|
||
("mumu", "recover_ahk_exe"): "MUMU_RECOVER_AHK_EXE",
|
||
("mumu", "clean_backup_path"): "CLEAN_BACKUP_PATH",
|
||
("mumu", "local_import_dir"): "LOCAL_IMPORT_DIR",
|
||
("mumu", "network_bridge_card"): "NET_BRIDGE_CARD",
|
||
("mumu", "adb_ip_offset"): "EMULATOR_ADB_IP_OFFSET",
|
||
("mumu", "bridge_ip_offset"): "MUMU_BRIDGE_IP_OFFSET",
|
||
("mumu", "network_script_path"): "MUMU_NETWORK_SCRIPT_PATH",
|
||
("mumu", "network_gateways"): "MUMU_NETWORK_GATEWAYS",
|
||
("share", "smb_user"): "SHARE_SMB_USER",
|
||
("share", "smb_password"): "SHARE_SMB_PASSWORD",
|
||
("logging", "file"): "LOG_FILE",
|
||
("logging", "max_bytes"): "LOG_MAX_BYTES",
|
||
("logging", "backup_count"): "LOG_BACKUP_COUNT",
|
||
("logging", "date_format"): "LOG_DATE_FORMAT",
|
||
("notifications", "wechat_tokens"): "WECHAT_TOKENS",
|
||
("notifications", "wecom_tokens"): "WECOM_TOKENS",
|
||
("alert", "window_seconds"): "ALERT_WINDOW_SECONDS",
|
||
("alert", "threshold"): "ALERT_THRESHOLD",
|
||
("alert", "cooldown_seconds"): "ALERT_COOLDOWN_SECONDS",
|
||
("monitoring", "status_push_interval"): "WORKER_STATUS_PUSH_INTERVAL",
|
||
("monitoring", "timezone"): "MONITORING_TIMEZONE",
|
||
("monitoring", "timeline_bucket_minutes"): "MONITORING_TIMELINE_BUCKET_MINUTES",
|
||
("pipeline", "run_option_keys"): "RUN_PIPELINE_OPTION_KEYS",
|
||
("pipeline", "default_options"): "RUN_PIPELINE_DEFAULT_OPTIONS",
|
||
("pipeline", "auto_reboot_recovery_delay"): "AUTO_REBOOT_RECOVERY_DELAY_SECONDS",
|
||
("pipeline", "legacy_end_worker_ips"): "RUNBATCH_LEGACY_END_WORKER_IPS",
|
||
("pipeline", "force_kill_images"): "RUNBATCH_FORCE_KILL_IMAGES",
|
||
("dashboard", "host"): "DASHBOARD_HOST",
|
||
("dashboard", "port"): "DASHBOARD_PORT",
|
||
("analytics", "traffic_root"): "ANALYTICS_TRAFFIC_ROOT",
|
||
("analytics", "traffic_root_block"): "ANALYTICS_TRAFFIC_ROOT_BLOCK",
|
||
("analytics", "traversal_root"): "ANALYTICS_TRAVERSAL_ROOT",
|
||
("analytics", "tpdpi_app_list_file"): "ANALYTICS_TPDPI_APP_LIST_FILE",
|
||
("analytics", "tpdpi_url_lib_file"): "ANALYTICS_TPDPI_URL_LIB_FILE",
|
||
("analytics", "artifact_wait_seconds"): "ANALYTICS_ARTIFACT_WAIT_SECONDS",
|
||
("analytics", "artifact_file_wait_seconds"): "ANALYTICS_ARTIFACT_FILE_WAIT_SECONDS",
|
||
("analytics", "job_poll_seconds"): "ANALYTICS_JOB_POLL_SECONDS",
|
||
("analytics", "model_traffic_threshold"): "MODEL_TRAFFIC_THRESHOLD",
|
||
("minio", "enabled"): "MINIO_ENABLED",
|
||
("minio", "endpoint"): "MINIO_ENDPOINT",
|
||
("minio", "access_key"): "MINIO_ACCESS_KEY",
|
||
("minio", "secret_key"): "MINIO_SECRET_KEY",
|
||
("minio", "bucket"): "MINIO_BUCKET",
|
||
("minio", "secure"): "MINIO_SECURE",
|
||
("apk", "download_mode"): "APK_DOWNLOAD_MODE",
|
||
("apk", "local_storage_dir"): "APK_LOCAL_STORAGE_DIR",
|
||
("apk", "smb_dir"): "APK_SMB_DIR",
|
||
("apk", "us_export_dir"): "APK_US_EXPORT_DIR",
|
||
("apk", "us_device_a_serial"): "APK_US_DEVICE_A_SERIAL",
|
||
("apk", "us_device_b_serial"): "APK_US_DEVICE_B_SERIAL",
|
||
("apk", "us_device_a_vm_index"): "APK_US_DEVICE_A_VM_INDEX",
|
||
("apk", "us_device_b_vm_index"): "APK_US_DEVICE_B_VM_INDEX",
|
||
("apk", "prefetch_poll_interval"): "APK_PREFETCH_POLL_INTERVAL",
|
||
("apk", "download_queue_interval"): "APK_DOWNLOAD_QUEUE_INTERVAL",
|
||
("apk", "minio_max_bytes"): "APK_MINIO_MAX_BYTES",
|
||
("apk", "emulator_cleanup_enabled"): "APK_EMULATOR_CLEANUP_ENABLED",
|
||
("task_routing_rules",): "TASK_ROUTING_RULES",
|
||
("worker_inventory",): "WORKER_INVENTORY",
|
||
}
|
||
|
||
|
||
def _is_nested_format(config: Dict[str, Any]) -> bool:
|
||
"""检测是否为嵌套小写格式：顶层有 environment 键且无 CONFIG_ENV"""
|
||
return "environment" in config and "CONFIG_ENV" not in config
|
||
|
||
|
||
def _dig(config: Dict[str, Any], *keys: str) -> Any:
|
||
"""从嵌套字典中按路径取值"""
|
||
value = config
|
||
for key in keys:
|
||
if not isinstance(value, dict):
|
||
return None
|
||
value = value.get(key)
|
||
return value
|
||
|
||
|
||
def _normalize_nested_config(config: Dict[str, Any]) -> Dict[str, Any]:
|
||
"""将嵌套小写格式转为扁平大写格式，已是扁平格式则直接返回"""
|
||
if not _is_nested_format(config):
|
||
return config
|
||
|
||
flat = {}
|
||
for nested_keys, flat_key in _NESTED_TO_FLAT_MAP.items():
|
||
value = _dig(config, *nested_keys)
|
||
if value is not None:
|
||
flat[flat_key] = value
|
||
return flat
|
||
|
||
|
||
def load_dispatcher_config(
|
||
env_name: Optional[str] = None,
|
||
config_path: Optional[str] = None,
|
||
) -> Dict[str, Any]:
|
||
if config_path:
|
||
return _normalize_nested_config(_load_config_dict(os.path.abspath(config_path), required=True))
|
||
|
||
if env_name is not None:
|
||
return _load_legacy_layered_config(env_name)
|
||
|
||
if CONFIG_PATH_ENV:
|
||
return _normalize_nested_config(_load_config_dict(os.path.abspath(CONFIG_PATH_ENV), required=True))
|
||
|
||
if os.path.exists(CONFIG_YAML_PATH):
|
||
return _normalize_nested_config(_load_config_dict(CONFIG_YAML_PATH, required=True))
|
||
|
||
if os.path.exists(CONFIG_PATH):
|
||
return _normalize_nested_config(_load_config_dict(CONFIG_PATH, required=True))
|
||
|
||
if os.path.exists(CONFIG_EXAMPLE_YAML_PATH):
|
||
return _normalize_nested_config(_load_config_dict(CONFIG_EXAMPLE_YAML_PATH, required=True))
|
||
|
||
if os.path.exists(CONFIG_EXAMPLE_PATH):
|
||
return _normalize_nested_config(_load_config_dict(CONFIG_EXAMPLE_PATH, required=True))
|
||
|
||
legacy_env = _read_current_env()
|
||
return _load_legacy_layered_config(legacy_env)
|
||
|
||
|
||
CONFIG = load_dispatcher_config()
|
||
CONFIG_ENV = normalize_env_name(CONFIG.get("CONFIG_ENV") or CONFIG.get("INSTANCE_NAME") or "prod")
|
||
|
||
|
||
# ==================== Instance ====================
|
||
INSTANCE_NAME = str(CONFIG["INSTANCE_NAME"]).strip()
|
||
INSTANCE_DIR = os.path.join(BASE_DIR, "runtime", INSTANCE_NAME)
|
||
REPORT_DIR = os.path.join(INSTANCE_DIR, "reports")
|
||
|
||
|
||
# ==================== Redis ====================
|
||
REDIS_HOST = str(CONFIG["REDIS_HOST"]).strip()
|
||
REDIS_PORT = int(CONFIG["REDIS_PORT"])
|
||
REDIS_DB = int(CONFIG["REDIS_DB"])
|
||
CHANNEL_NAMESPACE = str(CONFIG["CHANNEL_NAMESPACE"]).strip()
|
||
REDIS_MAX_CONNECTIONS_DISPATCHER = int(CONFIG["REDIS_MAX_CONNECTIONS_DISPATCHER"])
|
||
|
||
# ==================== Tasks ====================
|
||
MAX_RETRY_COUNT = int(CONFIG["MAX_RETRY_COUNT"])
|
||
TASK_CSV_PATH = os.path.join(BASE_DIR, str(CONFIG["TASK_CSV_FILE"]).strip())
|
||
|
||
# ==================== Worker ====================
|
||
WORKER_STALE_TIMEOUT = int(CONFIG["WORKER_STALE_TIMEOUT"])
|
||
WORKER_REINIT_TIMEOUT = int(CONFIG["WORKER_REINIT_TIMEOUT"])
|
||
WORKER_INIT_MONITOR_WINDOW = int(CONFIG["WORKER_INIT_MONITOR_WINDOW"])
|
||
WORKER_ACTION_MAX_PARALLEL = int(CONFIG["WORKER_ACTION_MAX_PARALLEL"])
|
||
SSH_CONNECT_TIMEOUT_SECONDS = int(CONFIG["SSH_CONNECT_TIMEOUT_SECONDS"])
|
||
SSH_ACTION_TIMEOUT_SECONDS = int(CONFIG["SSH_ACTION_TIMEOUT_SECONDS"])
|
||
AUTOOL_GIT_REPO_URL = str(CONFIG["AUTOOL_GIT_REPO_URL"]).strip()
|
||
MUMU_MANAGER_PATH = str(CONFIG["MUMU_MANAGER_PATH"]).strip()
|
||
MUMU_VM_INDEX = int(CONFIG["MUMU_VM_INDEX"])
|
||
MUMU_MAX_VM_INDEX = int(CONFIG["MUMU_MAX_VM_INDEX"])
|
||
CLEAN_BACKUP_PATH = str(CONFIG["CLEAN_BACKUP_PATH"]).strip()
|
||
LOCAL_IMPORT_DIR = str(CONFIG["LOCAL_IMPORT_DIR"]).strip()
|
||
NET_BRIDGE_CARD = str(CONFIG["NET_BRIDGE_CARD"]).strip()
|
||
EMULATOR_ADB_IP_OFFSET = int(CONFIG.get("EMULATOR_ADB_IP_OFFSET", 100))
|
||
MUMU_BRIDGE_IP_OFFSET = int(CONFIG.get("MUMU_BRIDGE_IP_OFFSET", 100))
|
||
MUMU_NETWORK_GATEWAYS = dict(
|
||
CONFIG.get("MUMU_NETWORK_GATEWAYS")
|
||
or {"192.168.1": "192.168.1.1", "192.168.2": "192.168.2.1"}
|
||
)
|
||
WORKER_ALLOWED_IP_PATTERNS = tuple(str(item).strip() for item in CONFIG.get("WORKER_ALLOWED_IP_PATTERNS", []) if str(item).strip())
|
||
SSH_DEFAULT_USER = str(CONFIG["SSH_DEFAULT_USER"]).strip()
|
||
SSH_DEFAULT_PASSWORD = str(CONFIG["SSH_DEFAULT_PASSWORD"]).strip()
|
||
SSH_DEFAULT_PORT = int(CONFIG["SSH_DEFAULT_PORT"])
|
||
PSEXEC_DEFAULT_SESSION_ID = int(CONFIG["PSEXEC_DEFAULT_SESSION_ID"])
|
||
MUMU_RESTART_SETTLE_SECONDS = int(CONFIG.get("MUMU_RESTART_SETTLE_SECONDS", 15))
|
||
RUNBATCH_LEGACY_END_WORKER_IPS = frozenset(
|
||
str(item).strip()
|
||
for item in CONFIG.get("RUNBATCH_LEGACY_END_WORKER_IPS", ["192.168.1.51", "192.168.1.61", "192.168.2.101"])
|
||
if str(item).strip()
|
||
)
|
||
RUNBATCH_FORCE_KILL_IMAGES = tuple(
|
||
str(item).strip()
|
||
for item in CONFIG.get("RUNBATCH_FORCE_KILL_IMAGES", ["OpenConsole.exe", "WindowsTerminal.exe", "conhost.exe", "powershell.exe"])
|
||
if str(item).strip()
|
||
)
|
||
MUMU_RECOVER_SCRIPT_PATH = str(CONFIG.get("MUMU_RECOVER_SCRIPT_PATH", r"\\your-file-server\share\recover_mumu_image.cmd")).strip()
|
||
MUMU_NETWORK_SCRIPT_PATH = str(CONFIG.get("MUMU_NETWORK_SCRIPT_PATH", r"\\your-file-server\share\recover_mumu_network.cmd")).strip()
|
||
MUMU_RECOVER_AHK_EXE = str(CONFIG.get("MUMU_RECOVER_AHK_EXE", r"\\your-file-server\share\mumu_recover.exe")).strip()
|
||
|
||
# ==================== Backward compatibility ====================
|
||
SHARE_SMB_USER = str(CONFIG.get("SHARE_SMB_USER") or "").strip()
|
||
SHARE_SMB_PASSWORD = str(CONFIG.get("SHARE_SMB_PASSWORD") or "").strip()
|
||
|
||
|
||
def get_repo_url(worker_ip: str) -> str:
|
||
del worker_ip
|
||
return AUTOOL_GIT_REPO_URL
|
||
|
||
|
||
def get_clean_backup_path(worker_ip: str) -> str:
|
||
del worker_ip
|
||
return CLEAN_BACKUP_PATH
|
||
|
||
|
||
def get_share_smb_target(worker_ip: str) -> str:
|
||
del worker_ip
|
||
path = CLEAN_BACKUP_PATH
|
||
if not path.startswith("\\\\"):
|
||
return ""
|
||
parts = [part for part in path.split("\\") if part]
|
||
if not parts:
|
||
return ""
|
||
return f"\\\\{parts[0]}"
|
||
|
||
# ==================== Output files ====================
|
||
FAILED_TASKS_CSV = os.path.join(REPORT_DIR, "failed_tasks.csv")
|
||
RETRY_TASKS_CSV = os.path.join(REPORT_DIR, "retry_tasks.csv")
|
||
SUCCESS_TASKS_CSV = os.path.join(REPORT_DIR, "success_tasks.csv")
|
||
FINAL_REPORT_CSV = os.path.join(REPORT_DIR, "final_report.csv")
|
||
WORKER_REPORT_CSV = os.path.join(REPORT_DIR, "worker_report.csv")
|
||
|
||
# ==================== Logging ====================
|
||
LOG_DIR = os.path.join(INSTANCE_DIR, "logs")
|
||
LOG_FILE = str(CONFIG["LOG_FILE"]).strip()
|
||
LOG_MAX_BYTES = int(CONFIG["LOG_MAX_BYTES"])
|
||
LOG_BACKUP_COUNT = int(CONFIG["LOG_BACKUP_COUNT"])
|
||
LOG_FORMAT = f"%(asctime)s [{INSTANCE_NAME}] [%(levelname)s] %(message)s"
|
||
LOG_DATE_FORMAT = str(CONFIG["LOG_DATE_FORMAT"]).strip()
|
||
|
||
# ==================== Notifications ====================
|
||
WECHAT_TOKENS = dict(CONFIG.get("WECHAT_TOKENS") or {})
|
||
WECOM_TOKENS = dict(CONFIG.get("WECOM_TOKENS") or {})
|
||
|
||
# ==================== Alert strategy ====================
|
||
ALERT_WINDOW_SECONDS = int(CONFIG["ALERT_WINDOW_SECONDS"])
|
||
ALERT_THRESHOLD = int(CONFIG["ALERT_THRESHOLD"])
|
||
ALERT_COOLDOWN_SECONDS = int(CONFIG["ALERT_COOLDOWN_SECONDS"])
|
||
|
||
# ==================== Status push ====================
|
||
WORKER_STATUS_PUSH_INTERVAL = int(CONFIG["WORKER_STATUS_PUSH_INTERVAL"])
|
||
|
||
RUN_PIPELINE_OPTION_KEYS = tuple(CONFIG.get("RUN_PIPELINE_OPTION_KEYS") or ())
|
||
RUN_PIPELINE_DEFAULT_OPTIONS = dict(CONFIG.get("RUN_PIPELINE_DEFAULT_OPTIONS") or {})
|
||
AUTO_REBOOT_RECOVERY_DELAY_SECONDS = int(CONFIG["AUTO_REBOOT_RECOVERY_DELAY_SECONDS"])
|
||
|
||
|
||
# ==================== Worker registry ====================
|
||
WORKER_INVENTORY_CONFIG_KEY = "WORKER_INVENTORY"
|
||
WORKER_INVENTORY_PATH = CONFIG_PATH
|
||
|
||
|
||
def _resolve_ssh_endpoint(ssh_target: str, ssh_port: int) -> Tuple[str, int]:
|
||
"""直连模式：解析 SSH 目标地址，验证合法性后直接返回"""
|
||
host = str(ssh_target).split("@")[-1].strip()
|
||
ip = extract_worker_ip(host)
|
||
if ip != host:
|
||
return host, ssh_port
|
||
if WORKER_ALLOWED_IP_PATTERNS and not any(re.fullmatch(pattern, ip) for pattern in WORKER_ALLOWED_IP_PATTERNS):
|
||
raise ValueError(f"worker ip is not allowed by WORKER_ALLOWED_IP_PATTERNS: {ssh_target}")
|
||
return host, ssh_port
|
||
|
||
|
||
def extract_worker_ip(value: Any) -> str:
|
||
text = str(value or "").strip()
|
||
match = re.search(r"(?<!\d)(?:\d{1,3}\.){3}\d{1,3}(?!\d)", text)
|
||
if not match:
|
||
return ""
|
||
octets = match.group(0).split(".")
|
||
if all(0 <= int(octet) <= 255 for octet in octets):
|
||
return match.group(0)
|
||
return ""
|
||
|
||
|
||
def normalize_worker_id(worker_id: Any = "", ip_address: Any = "") -> str:
|
||
return extract_worker_ip(ip_address) or extract_worker_ip(worker_id) or str(worker_id or "").strip()
|
||
|
||
|
||
def _normalize_worker_entry(raw: Dict[str, Any]) -> Dict[str, Any]:
|
||
ssh_target = str(raw["ssh_target"]).strip()
|
||
worker_id = normalize_worker_id(raw.get("worker_id", ""), ssh_target)
|
||
if not worker_id:
|
||
raise ValueError("worker inventory entry missing worker_id")
|
||
if not ssh_target:
|
||
raise ValueError(f"worker inventory entry missing ssh_target: {worker_id}")
|
||
ssh_user = str(raw.get("ssh_user") or SSH_DEFAULT_USER).strip()
|
||
ssh_password = str(raw.get("ssh_password") or "").strip()
|
||
ssh_port = int(raw.get("ssh_port") or SSH_DEFAULT_PORT)
|
||
ssh_host, ssh_port = _resolve_ssh_endpoint(ssh_target, ssh_port)
|
||
repo_dir = str(raw["repo_dir"]).strip()
|
||
if not repo_dir:
|
||
raise ValueError(f"worker inventory entry missing repo_dir: {worker_id}")
|
||
python_exe = str(raw["python_exe"]).strip()
|
||
if not python_exe:
|
||
raise ValueError(f"worker inventory entry missing python_exe: {worker_id}")
|
||
psexec_session_id = int(raw.get("psexec_session_id") or PSEXEC_DEFAULT_SESSION_ID)
|
||
tags = []
|
||
seen = set()
|
||
for tag in raw.get("tags", []) or []:
|
||
candidate = str(tag).strip()
|
||
if candidate and candidate not in seen:
|
||
seen.add(candidate)
|
||
tags.append(candidate)
|
||
|
||
item = {
|
||
"worker_id": worker_id,
|
||
"ssh_target": ssh_target,
|
||
"ssh_host": ssh_host,
|
||
"ssh_port": ssh_port,
|
||
"ssh_user": ssh_user,
|
||
"repo_dir": repo_dir,
|
||
"python_exe": python_exe,
|
||
"psexec_session_id": psexec_session_id,
|
||
"tags": tags,
|
||
}
|
||
if ssh_password:
|
||
item["ssh_password"] = ssh_password
|
||
return item
|
||
|
||
|
||
def _worker_matches_env(worker: Dict[str, Any], env_name: str) -> bool:
|
||
if env_name == "test":
|
||
return "test" in worker.get("tags", [])
|
||
return True
|
||
|
||
|
||
def load_worker_inventory(path: Optional[str] = None, env_name: Optional[str] = None) -> List[Dict[str, Any]]:
|
||
normalized_env = normalize_env_name(CONFIG_ENV if env_name is None else env_name)
|
||
target_path = path
|
||
if target_path:
|
||
if not os.path.exists(target_path):
|
||
return []
|
||
payload = _load_config_payload(target_path, required=True)
|
||
if isinstance(payload, dict):
|
||
payload = payload.get(WORKER_INVENTORY_CONFIG_KEY, [])
|
||
else:
|
||
payload = CONFIG.get(WORKER_INVENTORY_CONFIG_KEY, [])
|
||
|
||
if not isinstance(payload, list):
|
||
raise ValueError(f"worker inventory must be a list: {target_path}")
|
||
|
||
normalized: List[Dict[str, Any]] = []
|
||
seen = set()
|
||
for item in payload:
|
||
if not isinstance(item, dict):
|
||
raise ValueError(f"worker inventory entry must be an object: {item!r}")
|
||
normalized_item = _normalize_worker_entry(item)
|
||
worker_id = normalized_item["worker_id"]
|
||
if worker_id in seen:
|
||
raise ValueError(f"duplicate worker_id in worker inventory: {worker_id}")
|
||
seen.add(worker_id)
|
||
if not _worker_matches_env(normalized_item, normalized_env):
|
||
continue
|
||
normalized.append(normalized_item)
|
||
return normalized
|
||
|
||
|
||
WORKER_INVENTORY = load_worker_inventory()
|
||
MANAGED_WORKER_IDS = [item["worker_id"] for item in WORKER_INVENTORY]
|
||
ONLY_MANAGED_WORKERS_CAN_DISPATCH = bool(CONFIG.get("ONLY_MANAGED_WORKERS_CAN_DISPATCH", False))
|
||
|
||
# Exact-match routing rules for targeted dispatch.
|
||
# Values must be worker_id strings; worker IDs are normalized to the worker IP.
|
||
TASK_ROUTING_RULES = dict(CONFIG.get("TASK_ROUTING_RULES") or {"package_name": {}, "task_key": {}})
|
||
|
||
# ==================== Dashboard ====================
|
||
DASHBOARD_HOST = str(CONFIG["DASHBOARD_HOST"]).strip()
|
||
DASHBOARD_PORT = int(CONFIG["DASHBOARD_PORT"])
|
||
WORKER_ONLINE_TIMEOUT_SECONDS = WORKER_STALE_TIMEOUT
|
||
|
||
# ==================== Monitoring ====================
|
||
MONITORING_DB_PATH = os.path.join(INSTANCE_DIR, "monitoring.sqlite3")
|
||
MONITORING_TIMEZONE = str(CONFIG["MONITORING_TIMEZONE"]).strip()
|
||
MONITORING_TIMELINE_BUCKET_MINUTES = int(CONFIG["MONITORING_TIMELINE_BUCKET_MINUTES"])
|
||
|
||
# ==================== Analytics ====================
|
||
ANALYTICS_TRAFFIC_ROOT = str(CONFIG["ANALYTICS_TRAFFIC_ROOT"]).strip()
|
||
ANALYTICS_TRAFFIC_ROOT_BLOCK = str(CONFIG["ANALYTICS_TRAFFIC_ROOT_BLOCK"]).strip()
|
||
ANALYTICS_TRAVERSAL_ROOT = str(CONFIG["ANALYTICS_TRAVERSAL_ROOT"]).strip()
|
||
ANALYTICS_TPDPI_APP_LIST = os.path.join(
|
||
BASE_DIR,
|
||
"calculate",
|
||
str(CONFIG["ANALYTICS_TPDPI_APP_LIST_FILE"]).strip(),
|
||
)
|
||
ANALYTICS_TPDPI_URL_LIB = os.path.join(
|
||
BASE_DIR,
|
||
"calculate",
|
||
str(CONFIG["ANALYTICS_TPDPI_URL_LIB_FILE"]).strip(),
|
||
)
|
||
ANALYTICS_ARTIFACT_WAIT_SECONDS = int(CONFIG["ANALYTICS_ARTIFACT_WAIT_SECONDS"])
|
||
ANALYTICS_ARTIFACT_FILE_WAIT_SECONDS = int(CONFIG["ANALYTICS_ARTIFACT_FILE_WAIT_SECONDS"])
|
||
ANALYTICS_JOB_POLL_SECONDS = int(CONFIG["ANALYTICS_JOB_POLL_SECONDS"])
|
||
MODEL_TRAFFIC_THRESHOLD = int(CONFIG["MODEL_TRAFFIC_THRESHOLD"])
|
||
|
||
|
||
def channel_name(name: str) -> str:
|
||
return f"{CHANNEL_NAMESPACE}:{name}" if CHANNEL_NAMESPACE else name
|
||
|
||
|
||
# ==================== MinIO / APK Cloud Storage ====================
|
||
# 功能开关：是否启用 Minio 上传下载功能
|
||
# - True: 使用 Minio 进行 APK 上传和分发（原团队内网模式）
|
||
# - False: 禁用 Minio，使用直接下载模式（跨团队部署推荐）
|
||
MINIO_ENABLED = bool(CONFIG.get("MINIO_ENABLED", True))
|
||
|
||
# APK 下载模式（当 MINIO_ENABLED=False 时生效）
|
||
# - "direct": Worker 直接从下载链接获取 APK（跨团队部署推荐）
|
||
# - "smb": 从 SMB 共享目录读取 APK（需要内网访问权限）
|
||
APK_DOWNLOAD_MODE = str(CONFIG.get("APK_DOWNLOAD_MODE", "direct")).strip().lower()
|
||
|
||
# Minio 连接配置（仅当 MINIO_ENABLED=True 时必需）
|
||
MINIO_ENDPOINT = str(CONFIG.get("MINIO_ENDPOINT", "")).strip()
|
||
MINIO_ACCESS_KEY = str(CONFIG.get("MINIO_ACCESS_KEY", "")).strip()
|
||
MINIO_SECRET_KEY = str(CONFIG.get("MINIO_SECRET_KEY", "")).strip()
|
||
MINIO_BUCKET = str(CONFIG.get("MINIO_BUCKET", "autool-apk")).strip()
|
||
MINIO_SECURE = bool(CONFIG.get("MINIO_SECURE", False))
|
||
|
||
# APK 存储路径配置
|
||
APK_LOCAL_STORAGE_DIR = str(CONFIG.get("APK_LOCAL_STORAGE_DIR", "")).strip()
|
||
APK_SMB_DIR = str(CONFIG.get("APK_SMB_DIR", "")).strip()
|
||
APK_US_EXPORT_DIR = str(CONFIG.get("APK_US_EXPORT_DIR", "D:\\mumu_apks")).strip()
|
||
|
||
# US 拓扑设备配置（仅当使用 US 拓扑模式时需要）
|
||
APK_US_DEVICE_A_SERIAL = str(CONFIG.get("APK_US_DEVICE_A_SERIAL", "127.0.0.1:7555")).strip()
|
||
APK_US_DEVICE_B_SERIAL = str(CONFIG.get("APK_US_DEVICE_B_SERIAL", "127.0.0.1:7556")).strip()
|
||
APK_US_DEVICE_A_VM_INDEX = int(CONFIG.get("APK_US_DEVICE_A_VM_INDEX", 2))
|
||
APK_US_DEVICE_B_VM_INDEX = int(CONFIG.get("APK_US_DEVICE_B_VM_INDEX", 1))
|
||
|
||
# APK 管理配置
|
||
APK_PREFETCH_POLL_INTERVAL = int(CONFIG.get("APK_PREFETCH_POLL_INTERVAL", 120))
|
||
APK_DOWNLOAD_QUEUE_INTERVAL = int(CONFIG.get("APK_DOWNLOAD_QUEUE_INTERVAL", 60))
|
||
APK_MINIO_MAX_BYTES = int(CONFIG.get("APK_MINIO_MAX_BYTES", 107374182400)) # 100GB
|
||
APK_EMULATOR_CLEANUP_ENABLED = bool(CONFIG.get("APK_EMULATOR_CLEANUP_ENABLED", True))
|