287 lines
13 KiB
Python
287 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
import os
|
|
import sys
|
|
import json
|
|
import re
|
|
try:
|
|
from jinja2 import Environment, FileSystemLoader
|
|
except ImportError:
|
|
print("FATAL: jinja2 is not installed. Please run 'pip install jinja2'.", file=sys.stderr)
|
|
exit(1)
|
|
import logging
|
|
import ipaddress
|
|
from typing import Optional
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
|
|
|
def is_ip_address(address: str) -> bool:
|
|
"""Checks if a given string is a valid IP address (IPv4 or IPv6)."""
|
|
if not address:
|
|
return False
|
|
try:
|
|
ipaddress.ip_address(address)
|
|
return True
|
|
except ValueError:
|
|
return False
|
|
|
|
|
|
def load_dotenv(dotenv_path):
|
|
"""
|
|
Loads environment variables from a .env file.
|
|
Does not override existing environment variables from the system.
|
|
"""
|
|
if not os.path.exists(dotenv_path):
|
|
logging.warning(f".env file not found at {dotenv_path}. Using system environment variables or defaults.")
|
|
return False
|
|
try:
|
|
with open(dotenv_path) as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line and not line.startswith('#') and '=' in line:
|
|
key, value = line.split('=', 1)
|
|
key = key.strip()
|
|
value = value.strip()
|
|
# Remove surrounding quotes which are common in .env files
|
|
# Handle both single and double quotes
|
|
if (value.startswith('"') and value.endswith('"')) or \
|
|
(value.startswith("'") and value.endswith("'")):
|
|
value = value[1:-1]
|
|
|
|
# os.environ only takes strings
|
|
value = str(value)
|
|
|
|
if key not in os.environ:
|
|
os.environ[key] = value
|
|
logging.info(f"Successfully loaded variables from {dotenv_path}")
|
|
return True
|
|
except Exception as e:
|
|
logging.error(f"Failed to read or parse {dotenv_path}: {e}")
|
|
return False
|
|
|
|
def _get_port_from_proxy_url(url: str) -> Optional[str]:
|
|
"""Extracts the port from a proxy URL string."""
|
|
if not url or not isinstance(url, str):
|
|
return None
|
|
match = re.search(r':(\d+)$', url.strip())
|
|
return match.group(1) if match else None
|
|
|
|
def expand_env_vars(value: str) -> str:
|
|
"""
|
|
Expands environment variables in a string, including default values.
|
|
Supports ${VAR} and ${VAR:-default}.
|
|
"""
|
|
if not isinstance(value, str):
|
|
return value
|
|
|
|
# Regex to find ${VAR:-default} or ${VAR}
|
|
pattern = re.compile(r'\$\{(?P<var>\w+)(?::-(?P<default>.*?))?\}')
|
|
|
|
def replacer(match):
|
|
var_name = match.group('var')
|
|
default_value = match.group('default')
|
|
# Get value from os.environ, or use default, or empty string
|
|
return os.getenv(var_name, default_value if default_value is not None else '')
|
|
|
|
return pattern.sub(replacer, value)
|
|
|
|
def generate_configs():
|
|
"""
|
|
Generates envoy.yaml, docker-compose.camoufox.yaml, and camoufox_endpoints.json
|
|
from Jinja2 templates and environment variables.
|
|
"""
|
|
try:
|
|
# --- Setup Paths ---
|
|
# The script runs from /app. Configs and templates are in /app/configs.
|
|
project_root = os.path.dirname(os.path.abspath(__file__)) # This will be /app
|
|
configs_dir = os.path.join(project_root, 'configs')
|
|
|
|
# Load .env from the project root ONLY - no fallback
|
|
dotenv_path = os.path.join(project_root, '.env')
|
|
logging.info(f"Looking for .env file at: {dotenv_path}")
|
|
|
|
if os.path.exists(dotenv_path):
|
|
if load_dotenv(dotenv_path):
|
|
logging.info(f"Using .env file from: {dotenv_path}")
|
|
else:
|
|
logging.error(f"Failed to load .env file from: {dotenv_path}")
|
|
exit(1)
|
|
else:
|
|
logging.warning(f".env file not found at {dotenv_path}. Using system environment variables or defaults.")
|
|
|
|
# --- Common Configuration ---
|
|
ytdlp_workers_str = os.getenv('YTDLP_WORKERS', '3').strip()
|
|
try:
|
|
# Handle empty string case by defaulting to 3, otherwise convert to int.
|
|
worker_count = int(ytdlp_workers_str) if ytdlp_workers_str else 3
|
|
except (ValueError, TypeError):
|
|
logging.warning(f"Invalid value for YTDLP_WORKERS: '{ytdlp_workers_str}'. Defaulting to 3.")
|
|
worker_count = 3
|
|
|
|
if worker_count == 0:
|
|
worker_count = os.cpu_count() or 1
|
|
logging.info(f"YTDLP_WORKERS is 0, auto-detected {worker_count} CPU cores for worker and camoufox config.")
|
|
|
|
# The templates are in the 'configs' directory.
|
|
env = Environment(loader=FileSystemLoader(configs_dir), trim_blocks=True, lstrip_blocks=True)
|
|
# Make the helper function available to Jinja2 templates
|
|
env.globals['_get_port_from_proxy_url'] = _get_port_from_proxy_url
|
|
|
|
# Get service role from environment to determine what to generate
|
|
# Ensure we strip any remaining quotes that might have slipped through
|
|
service_role = os.getenv('service_role', 'management')
|
|
# Additional stripping of quotes for robustness
|
|
if (service_role.startswith('"') and service_role.endswith('"')) or \
|
|
(service_role.startswith("'") and service_role.endswith("'")):
|
|
service_role = service_role[1:-1]
|
|
|
|
logging.info(f"Service role for generation: '{service_role}'")
|
|
|
|
# --- Camoufox Configuration (only for worker/all-in-one roles) ---
|
|
camoufox_proxies = []
|
|
expanded_camoufox_proxies_str = ""
|
|
if service_role != 'management':
|
|
logging.info("--- Generating Camoufox (Remote Browser) Configuration ---")
|
|
camoufox_proxies_str = os.getenv('CAMOUFOX_PROXIES')
|
|
if not camoufox_proxies_str:
|
|
logging.warning("CAMOUFOX_PROXIES environment variable not set. No camoufox instances will be generated.")
|
|
else:
|
|
# Expand environment variables within the string before splitting
|
|
expanded_camoufox_proxies_str = expand_env_vars(camoufox_proxies_str)
|
|
logging.info(f"Expanded CAMOUFOX_PROXIES from '{camoufox_proxies_str}' to '{expanded_camoufox_proxies_str}'")
|
|
camoufox_proxies = [{'url': p.strip()} for p in expanded_camoufox_proxies_str.split(',') if p.strip()]
|
|
logging.info(f"Found {len(camoufox_proxies)} proxy/proxies for Camoufox.")
|
|
logging.info(f"Each Camoufox instance will support {worker_count} concurrent browser sessions.")
|
|
logging.info(f"Total browser sessions supported on this worker: {len(camoufox_proxies) * worker_count}")
|
|
|
|
|
|
vnc_password = os.getenv('VNC_PASSWORD', 'supersecret')
|
|
base_vnc_port = int(os.getenv('CAMOUFOX_BASE_VNC_PORT', 5901))
|
|
camoufox_port = int(os.getenv('CAMOUFOX_PORT', 12345))
|
|
camoufox_backend_prefix = os.getenv('CAMOUFOX_BACKEND_PREFIX', 'camoufox-')
|
|
|
|
# --- Generate docker-compose.camoufox.yaml ---
|
|
compose_template = env.get_template('docker-compose.camoufox.yaml.j2')
|
|
compose_output_file = os.path.join(configs_dir, 'docker-compose.camoufox.yaml')
|
|
|
|
camoufox_config_data = {
|
|
'camoufox_proxies': camoufox_proxies,
|
|
'vnc_password': vnc_password,
|
|
'camoufox_port': camoufox_port,
|
|
'worker_count': worker_count,
|
|
}
|
|
|
|
rendered_compose_config = compose_template.render(camoufox_config_data)
|
|
with open(compose_output_file, 'w') as f:
|
|
f.write(rendered_compose_config)
|
|
|
|
logging.info(f"Successfully generated {compose_output_file} with {len(camoufox_proxies)} camoufox service(s).")
|
|
logging.info("This docker-compose file defines the remote browser services, one for each proxy.")
|
|
logging.info("----------------------------------------------------------")
|
|
|
|
# --- Generate camoufox_endpoints.json ---
|
|
endpoints_map = {}
|
|
for i, proxy in enumerate(camoufox_proxies):
|
|
proxy_port = _get_port_from_proxy_url(proxy['url'])
|
|
if proxy_port:
|
|
# Use the correct container name pattern that matches the docker-compose template
|
|
# The container name in the template is: ytdlp-ops-camoufox-{{ proxy_port }}-{{ loop.index }}-1
|
|
container_name = f"ytdlp-ops-camoufox-{proxy_port}-{i+1}-1"
|
|
container_base_port = camoufox_port + i * worker_count
|
|
endpoints = []
|
|
for j in range(worker_count):
|
|
port = container_base_port + j
|
|
endpoints.append(f"ws://{container_name}:{port}/mypath")
|
|
|
|
endpoints_map[proxy_port] = {
|
|
"ws_endpoints": endpoints
|
|
}
|
|
else:
|
|
logging.warning(f"Could not extract port from proxy URL: {proxy['url']}. Skipping for endpoint map.")
|
|
|
|
endpoints_data = {"endpoints": endpoints_map}
|
|
# The camoufox directory is at the root of the project context, not under 'airflow'.
|
|
# camoufox_dir = os.path.join(project_root, 'camoufox')
|
|
# os.makedirs(camoufox_dir, exist_ok=True)
|
|
endpoints_output_file = os.path.join(configs_dir, 'camoufox_endpoints.json')
|
|
with open(endpoints_output_file, 'w') as f:
|
|
json.dump(endpoints_data, f, indent=2)
|
|
logging.info(f"Successfully generated {endpoints_output_file} with {len(endpoints_map)} port-keyed endpoint(s).")
|
|
logging.info("This file maps each proxy to a list of WebSocket endpoints for Camoufox.")
|
|
logging.info("The token_generator uses this map to connect to the correct remote browser.")
|
|
else:
|
|
logging.info("Skipping Camoufox configuration generation for 'management' role.")
|
|
|
|
# --- Generate docker-compose-ytdlp-ops.yaml ---
|
|
ytdlp_ops_template = env.get_template('docker-compose-ytdlp-ops.yaml.j2')
|
|
ytdlp_ops_output_file = os.path.join(configs_dir, 'docker-compose-ytdlp-ops.yaml')
|
|
|
|
# Combine all proxies (camoufox and general) into a single string for the server.
|
|
all_proxies = []
|
|
|
|
# Track if we have any explicit proxy configuration
|
|
has_explicit_proxies = False
|
|
|
|
# Add camoufox proxies if they exist
|
|
if expanded_camoufox_proxies_str:
|
|
camoufox_proxy_list = [p.strip() for p in expanded_camoufox_proxies_str.split(',') if p.strip()]
|
|
all_proxies.extend(camoufox_proxy_list)
|
|
if camoufox_proxy_list:
|
|
has_explicit_proxies = True
|
|
logging.info(f"Added {len(camoufox_proxy_list)} camoufox proxies: {camoufox_proxy_list}")
|
|
|
|
|
|
combined_proxies_str = ",".join(all_proxies)
|
|
logging.info(f"Combined proxy string for ytdlp-ops-service: '{combined_proxies_str}'")
|
|
|
|
ytdlp_ops_config_data = {
|
|
'combined_proxies_str': combined_proxies_str,
|
|
'service_role': service_role,
|
|
'camoufox_proxies': camoufox_proxies,
|
|
}
|
|
rendered_ytdlp_ops_config = ytdlp_ops_template.render(ytdlp_ops_config_data)
|
|
with open(ytdlp_ops_output_file, 'w') as f:
|
|
f.write(rendered_ytdlp_ops_config)
|
|
logging.info(f"Successfully generated {ytdlp_ops_output_file}")
|
|
|
|
# --- Envoy Configuration ---
|
|
envoy_port = int(os.getenv('ENVOY_PORT', 9080))
|
|
base_port = int(os.getenv('YTDLP_BASE_PORT', 9090))
|
|
envoy_admin_port = int(os.getenv('ENVOY_ADMIN_PORT', 9901))
|
|
# For local dev, ENVOY_BACKEND_ADDRESS is set to 127.0.0.1. For Docker, it's unset, so we default to the service name.
|
|
backend_address = os.getenv('ENVOY_BACKEND_ADDRESS', 'ytdlp-ops-service')
|
|
# Use STATIC for IP addresses, and STRICT_DNS for anything else (hostnames).
|
|
envoy_cluster_type = 'STATIC' if is_ip_address(backend_address) else 'STRICT_DNS'
|
|
|
|
# --- Generate envoy.yaml ---
|
|
envoy_template = env.get_template('envoy.yaml.j2')
|
|
# Output envoy.yaml to the configs directory, where other generated files are.
|
|
envoy_output_file = os.path.join(configs_dir, 'envoy.yaml')
|
|
|
|
logging.info("--- Generating Envoy Configuration ---")
|
|
logging.info(f"Envoy will listen on public port: {envoy_port}")
|
|
logging.info(f"It will load balance requests across {worker_count} internal gRPC endpoints of the 'ytdlp-ops-service'.")
|
|
logging.info(f"The backend service is located at: '{backend_address}' (type: {envoy_cluster_type})")
|
|
|
|
envoy_config_data = {
|
|
'envoy_port': envoy_port,
|
|
'worker_count': worker_count,
|
|
'base_port': base_port,
|
|
'envoy_admin_port': envoy_admin_port,
|
|
'backend_address': backend_address,
|
|
'envoy_cluster_type': envoy_cluster_type,
|
|
}
|
|
|
|
rendered_envoy_config = envoy_template.render(envoy_config_data)
|
|
with open(envoy_output_file, 'w') as f:
|
|
f.write(rendered_envoy_config)
|
|
logging.info(f"Successfully generated {envoy_output_file}")
|
|
logging.info("--- Configuration Generation Complete ---")
|
|
|
|
except Exception as e:
|
|
logging.error(f"Failed to generate configurations: {e}", exc_info=True)
|
|
exit(1)
|
|
|
|
if __name__ == '__main__':
|
|
generate_configs()
|