296 lines
14 KiB
Python
296 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
CLI tool to set up profiles from a YAML policy file.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from typing import List
|
|
|
|
try:
|
|
import yaml
|
|
except ImportError:
|
|
print("PyYAML is not installed. Please install it with: pip install PyYAML", file=sys.stderr)
|
|
yaml = None
|
|
|
|
try:
|
|
from dotenv import load_dotenv
|
|
except ImportError:
|
|
load_dotenv = None
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def run_command(cmd: List[str], capture: bool = False):
|
|
"""Runs a command and raises an exception on failure."""
|
|
logger.debug(f"Running command: {' '.join(cmd)}")
|
|
# check=True will raise CalledProcessError on non-zero exit codes
|
|
result = subprocess.run(
|
|
cmd,
|
|
capture_output=capture,
|
|
text=True,
|
|
check=True
|
|
)
|
|
return result
|
|
|
|
def add_setup_profiles_parser(subparsers):
|
|
"""Adds the parser for the 'setup-profiles' command."""
|
|
parser = subparsers.add_parser(
|
|
'setup-profiles',
|
|
description="Set up profiles for a simulation or test run based on a policy file.",
|
|
formatter_class=argparse.RawTextHelpFormatter,
|
|
help="Set up profiles from a policy file."
|
|
)
|
|
parser.add_argument('--policy', '--policy-file', dest='policy_file', required=True, help="Path to the YAML profile setup policy file.")
|
|
parser.add_argument('--env', help="Override the environment name from the policy file. For multi-setup files, this will override the 'env' for ALL setups being run.")
|
|
parser.add_argument('--env-file', help="Override the env_file setting in the policy.")
|
|
parser.add_argument('--auth-only', action='store_true', help='In a multi-setup policy file, run only the auth_profile_setup.')
|
|
parser.add_argument('--download-only', action='store_true', help='In a multi-setup policy file, run only the download_profile_setup.')
|
|
parser.add_argument('--redis-host', default=None, help='Redis host. Overrides policy and .env file.')
|
|
parser.add_argument('--redis-port', type=int, default=None, help='Redis port. Overrides policy and .env file.')
|
|
parser.add_argument('--redis-password', default=None, help='Redis password. Overrides policy and .env file.')
|
|
parser.add_argument('--preserve-profiles', action='store_true', help="Do not clean up existing profiles; create only what is missing.")
|
|
parser.add_argument('--cleanup-prefix', action='append', help="Prefix of profiles to delete before setup. Can be specified multiple times. Overrides policy-based cleanup.")
|
|
parser.add_argument('--cleanup-all', action='store_true', help="(Destructive) Delete ALL data for the environment (profiles, proxies, counters) before setup. Overrides all other cleanup options.")
|
|
parser.add_argument('--reset-global-counters', action='store_true', help="Reset global counters like 'failed_lock_attempts'.")
|
|
parser.add_argument('--verbose', action='store_true', help="Enable verbose logging.")
|
|
return parser
|
|
|
|
def _run_setup_for_env(profile_setup: dict, common_args: list, args: argparse.Namespace) -> int:
|
|
"""Runs the profile setup logic for a given configuration block."""
|
|
if args.cleanup_all:
|
|
logger.info("--- (DESTRUCTIVE) Cleaning up all data for the environment via --cleanup-all ---")
|
|
try:
|
|
cleanup_cmd = ['bin/ytops-client', 'profile', 'delete-all', '--confirm'] + common_args
|
|
run_command(cleanup_cmd)
|
|
except subprocess.CalledProcessError as e:
|
|
logger.error(f"Failed to clean up all data for the environment: {e}")
|
|
return 1 # Stop if cleanup fails
|
|
|
|
# Disable other cleanup logic
|
|
profile_setup['cleanup_before_run'] = False
|
|
args.cleanup_prefix = None
|
|
|
|
# If --cleanup-prefix is provided, it takes precedence over policy settings
|
|
if args.cleanup_prefix:
|
|
logger.info("--- Cleaning up profiles based on --cleanup-prefix ---")
|
|
for prefix in args.cleanup_prefix:
|
|
try:
|
|
list_cmd = ['bin/ytops-client', 'profile', 'list', '--format', 'json'] + common_args
|
|
result = run_command(list_cmd, capture=True)
|
|
profiles_to_delete = [p for p in json.loads(result.stdout) if p['name'].startswith(prefix)]
|
|
|
|
if not profiles_to_delete:
|
|
logger.info(f"No profiles with prefix '{prefix}' found to delete.")
|
|
continue
|
|
|
|
logger.info(f"Found {len(profiles_to_delete)} profiles with prefix '{prefix}' to delete.")
|
|
for p in profiles_to_delete:
|
|
delete_cmd = ['bin/ytops-client', 'profile', 'delete', p['name'], '--confirm'] + common_args
|
|
run_command(delete_cmd)
|
|
except (subprocess.CalledProcessError, json.JSONDecodeError) as e:
|
|
logger.warning(f"Could not list or parse existing profiles with prefix '{prefix}' for cleanup. Error: {e}")
|
|
|
|
# Disable policy-based cleanup as we've handled it via CLI
|
|
profile_setup['cleanup_before_run'] = False
|
|
|
|
if args.preserve_profiles:
|
|
if profile_setup.get('cleanup_before_run'):
|
|
logger.info("--preserve-profiles is set, overriding 'cleanup_before_run: true' from policy.")
|
|
profile_setup['cleanup_before_run'] = False
|
|
|
|
if args.reset_global_counters:
|
|
logger.info("--- Resetting global counters ---")
|
|
try:
|
|
reset_cmd = ['bin/ytops-client', 'profile', 'reset-global-counters'] + common_args
|
|
run_command(reset_cmd)
|
|
except subprocess.CalledProcessError as e:
|
|
logger.error(f"Failed to reset global counters: {e}")
|
|
|
|
if profile_setup.get('cleanup_before_run'):
|
|
logger.info("--- Cleaning up old profiles ---")
|
|
for pool in profile_setup.get('pools', []):
|
|
prefix = pool.get('prefix')
|
|
if prefix:
|
|
try:
|
|
list_cmd = ['bin/ytops-client', 'profile', 'list', '--format', 'json'] + common_args
|
|
result = run_command(list_cmd, capture=True)
|
|
profiles = [p for p in json.loads(result.stdout) if p['name'].startswith(prefix)]
|
|
|
|
if not profiles:
|
|
logger.info(f"No profiles with prefix '{prefix}' found to delete.")
|
|
continue
|
|
|
|
logger.info(f"Found {len(profiles)} profiles with prefix '{prefix}' to delete.")
|
|
for p in profiles:
|
|
delete_cmd = ['bin/ytops-client', 'profile', 'delete', p['name'], '--confirm'] + common_args
|
|
run_command(delete_cmd)
|
|
except (subprocess.CalledProcessError, json.JSONDecodeError) as e:
|
|
logger.warning(f"Could not list existing profiles with prefix '{prefix}' for cleanup. Assuming none exist. Error: {e}")
|
|
|
|
existing_profiles = set()
|
|
if not profile_setup.get('cleanup_before_run'):
|
|
logger.info("--- Checking for existing profiles ---")
|
|
try:
|
|
list_cmd = ['bin/ytops-client', 'profile', 'list', '--format', 'json'] + common_args
|
|
result = run_command(list_cmd, capture=True)
|
|
profiles_data = json.loads(result.stdout)
|
|
existing_profiles = {p['name'] for p in profiles_data}
|
|
logger.info(f"Found {len(existing_profiles)} existing profiles.")
|
|
except (subprocess.CalledProcessError, json.JSONDecodeError) as e:
|
|
logger.error(f"Failed to list existing profiles. Will attempt to create all profiles. Error: {e}")
|
|
|
|
logger.info("--- Creating new profiles (if needed) ---")
|
|
for pool in profile_setup.get('pools', []):
|
|
prefix = pool.get('prefix')
|
|
proxy = pool.get('proxy')
|
|
count = pool.get('count', 0)
|
|
start_in_rest_minutes = pool.get('start_in_rest_for_minutes')
|
|
|
|
profiles_in_pool_created = 0
|
|
# Get a list of all profiles that should exist for this pool
|
|
profile_names_in_pool = [f"{prefix}_{i}" for i in range(count)]
|
|
|
|
for profile_name in profile_names_in_pool:
|
|
if profile_name in existing_profiles:
|
|
logger.debug(f"Profile '{profile_name}' already exists, preserving.")
|
|
continue
|
|
|
|
try:
|
|
create_cmd = ['bin/ytops-client', 'profile', 'create', profile_name, proxy] + common_args
|
|
run_command(create_cmd)
|
|
profiles_in_pool_created += 1
|
|
except subprocess.CalledProcessError as e:
|
|
logger.error(f"Failed to create profile '{profile_name}': {e}")
|
|
|
|
if profiles_in_pool_created > 0:
|
|
logger.info(f"Created {profiles_in_pool_created} new profile(s) for pool '{prefix}'.")
|
|
elif count > 0:
|
|
logger.info(f"No new profiles needed for pool '{prefix}'. All {count} profile(s) already exist.")
|
|
|
|
# If requested, put the proxy for this pool into a RESTING state.
|
|
# This is done even for existing profiles when --preserve-profiles is used.
|
|
if start_in_rest_minutes and proxy:
|
|
logger.info(f"Setting proxy '{proxy}' for pool '{prefix}' to start in RESTING state for {start_in_rest_minutes} minutes.")
|
|
try:
|
|
# The 'set-proxy-state' command takes '--duration-minutes'
|
|
set_state_cmd = ['bin/ytops-client', 'profile', 'set-proxy-state', proxy, 'RESTING',
|
|
'--duration-minutes', str(start_in_rest_minutes)] + common_args
|
|
run_command(set_state_cmd)
|
|
except subprocess.CalledProcessError as e:
|
|
logger.error(f"Failed to set initial REST state for proxy '{proxy}': {e}")
|
|
|
|
return 0
|
|
|
|
|
|
def main_setup_profiles(args):
|
|
"""Main logic for the 'setup-profiles' command."""
|
|
if not yaml: return 1
|
|
|
|
if args.verbose:
|
|
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
|
logger.setLevel(logging.DEBUG)
|
|
else:
|
|
if not logging.getLogger().handlers:
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
|
logger.setLevel(logging.INFO)
|
|
|
|
try:
|
|
with open(args.policy_file, 'r', encoding='utf-8') as f:
|
|
policy = yaml.safe_load(f) or {}
|
|
except (IOError, yaml.YAMLError) as e:
|
|
logger.error(f"Failed to load or parse policy file {args.policy_file}: {e}")
|
|
return 1
|
|
|
|
# --- Expand common_pools if defined ---
|
|
common_pools_def = policy.get('common_pools')
|
|
if common_pools_def:
|
|
logger.info("Found 'common_pools' definition. Expanding into setup blocks.")
|
|
expanded_pools = []
|
|
for pool_def in common_pools_def:
|
|
# Handle both 'prefix' (string) and 'prefixes' (list) for flexibility
|
|
prefixes = pool_def.get('prefixes', [])
|
|
if 'prefix' in pool_def and pool_def['prefix'] not in prefixes:
|
|
prefixes.append(pool_def['prefix'])
|
|
|
|
if not prefixes:
|
|
logger.warning(f"A pool in 'common_pools' is missing a 'prefix' or 'prefixes' key. Skipping: {pool_def}")
|
|
continue
|
|
|
|
for p in prefixes:
|
|
new_pool = pool_def.copy()
|
|
new_pool['prefix'] = p
|
|
if 'prefixes' in new_pool: del new_pool['prefixes']
|
|
expanded_pools.append(new_pool)
|
|
|
|
for setup_key in ['auth_profile_setup', 'download_profile_setup']:
|
|
if setup_key in policy and policy[setup_key].get('use_common_pools'):
|
|
logger.debug(f"Applying {len(expanded_pools)} common pool definitions to '{setup_key}'.")
|
|
policy[setup_key]['pools'] = expanded_pools
|
|
|
|
sim_params = policy.get('simulation_parameters', {})
|
|
setups_to_run = []
|
|
|
|
if not args.download_only and 'auth_profile_setup' in policy:
|
|
setups_to_run.append(('Auth', policy['auth_profile_setup']))
|
|
|
|
if not args.auth_only and 'download_profile_setup' in policy:
|
|
setups_to_run.append(('Download', policy['download_profile_setup']))
|
|
|
|
# Backward compatibility for old single-block format
|
|
if not setups_to_run and 'profile_setup' in policy:
|
|
legacy_config = policy['profile_setup']
|
|
# Synthesize the env from the global section for the legacy block
|
|
legacy_config['env'] = args.env or sim_params.get('env')
|
|
setups_to_run.append(('Legacy', legacy_config))
|
|
|
|
if not setups_to_run:
|
|
logger.error("No 'auth_profile_setup', 'download_profile_setup', or legacy 'profile_setup' block found in policy file.")
|
|
return 1
|
|
|
|
env_file = args.env_file or sim_params.get('env_file')
|
|
if load_dotenv:
|
|
if load_dotenv(env_file):
|
|
print(f"Loaded environment variables from {env_file or '.env file'}", file=sys.stderr)
|
|
elif args.env_file and not os.path.exists(args.env_file):
|
|
print(f"Error: The specified env_file was not found: {args.env_file}", file=sys.stderr)
|
|
return 1
|
|
|
|
base_common_args = []
|
|
if env_file: base_common_args.extend(['--env-file', env_file])
|
|
|
|
redis_host = args.redis_host or os.getenv('REDIS_HOST') or os.getenv('MASTER_HOST_IP') or sim_params.get('redis_host')
|
|
if redis_host: base_common_args.extend(['--redis-host', redis_host])
|
|
|
|
redis_port = args.redis_port
|
|
if redis_port is None:
|
|
redis_port_env = os.getenv('REDIS_PORT')
|
|
redis_port = int(redis_port_env) if redis_port_env and redis_port_env.isdigit() else sim_params.get('redis_port')
|
|
if redis_port: base_common_args.extend(['--redis-port', str(redis_port)])
|
|
|
|
redis_password = args.redis_password or os.getenv('REDIS_PASSWORD') or sim_params.get('redis_password')
|
|
if redis_password: base_common_args.extend(['--redis-password', redis_password])
|
|
|
|
if args.verbose: base_common_args.append('--verbose')
|
|
|
|
for setup_name, setup_config in setups_to_run:
|
|
logger.info(f"--- Running setup for {setup_name} simulation ---")
|
|
|
|
effective_env = args.env or setup_config.get('env')
|
|
if not effective_env:
|
|
logger.error(f"Could not determine environment for '{setup_name}' setup. Please specify 'env' in the policy block or via --env.")
|
|
return 1
|
|
|
|
env_common_args = base_common_args + ['--env', effective_env]
|
|
|
|
if _run_setup_for_env(setup_config, env_common_args, args) != 0:
|
|
return 1
|
|
|
|
logger.info("\n--- All profile setups complete. ---")
|
|
logger.info("You can now run the policy enforcer to manage the profiles:")
|
|
logger.info("e.g., bin/ytops-client policy-enforcer --policy-file policies/8_unified_simulation_enforcer.yaml --live")
|
|
return 0
|