137 lines
7.2 KiB
Python
137 lines
7.2 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
CLI tool to orchestrate multi-stage profile simulations.
|
|
"""
|
|
import argparse
|
|
import logging
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
try:
|
|
import yaml
|
|
except ImportError:
|
|
print("PyYAML is not installed. Please install it with: pip install PyYAML", file=sys.stderr)
|
|
yaml = None
|
|
|
|
# Import the main functions from the tools we are wrapping
|
|
from .profile_setup_tool import main_setup_profiles
|
|
from .stress_policy_tool import main_stress_policy
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Define default policy paths relative to the project root
|
|
PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
|
POLICY_DIR = PROJECT_ROOT / 'policies'
|
|
POLICY_FILE_SETUP = str(POLICY_DIR / '6_simulation_policy.yaml')
|
|
POLICY_FILE_AUTH = str(POLICY_DIR / '7_continuous_auth.yaml')
|
|
POLICY_FILE_DOWNLOAD = str(POLICY_DIR / '8_continuous_download.yaml')
|
|
|
|
def add_simulation_parser(subparsers):
|
|
"""Adds the parser for the 'simulation' command."""
|
|
parser = subparsers.add_parser(
|
|
'simulation',
|
|
description="Run multi-stage profile simulations (setup, auth, download). This provides a unified interface for the simulation workflow.",
|
|
formatter_class=argparse.RawTextHelpFormatter,
|
|
help="Run multi-stage profile simulations."
|
|
)
|
|
|
|
# Common arguments for all simulation subcommands
|
|
common_parser = argparse.ArgumentParser(add_help=False)
|
|
common_parser.add_argument('--env-file', default=None, help="Path to a .env file to load. Overrides setting from policy file.")
|
|
common_parser.add_argument('--redis-host', default=None, help='Redis host. Overrides policy and .env file.')
|
|
common_parser.add_argument('--redis-port', type=int, default=None, help='Redis port. Overrides policy and .env file.')
|
|
common_parser.add_argument('--redis-password', default=None, help='Redis password. Overrides policy and .env file.')
|
|
common_parser.add_argument('--env', default='sim', help="Environment name for Redis key prefix. Default: 'sim'.")
|
|
common_parser.add_argument('--expire-time-shift-minutes', type=int, default=None, help="Consider URLs expiring in N minutes as expired. Overrides policy.")
|
|
common_parser.add_argument('--verbose', action='store_true', help="Enable verbose logging.")
|
|
|
|
sim_subparsers = parser.add_subparsers(dest='simulation_command', help='Simulation stage to run', required=True)
|
|
|
|
# --- Setup ---
|
|
setup_parser = sim_subparsers.add_parser('setup', help='Set up profiles for a simulation.', parents=[common_parser])
|
|
setup_parser.add_argument('--policy-file', dest='policy', default=POLICY_FILE_SETUP, help=f'Path to the setup policy YAML file. Default: {POLICY_FILE_SETUP}')
|
|
setup_parser.add_argument('--preserve-profiles', action='store_true', help="Do not clean up existing profiles.")
|
|
setup_parser.add_argument('--reset-global-counters', action='store_true', help="Reset global counters like 'failed_lock_attempts'.")
|
|
|
|
# --- Auth ---
|
|
auth_parser = sim_subparsers.add_parser('auth', help='Run the authentication (get-info) part of the simulation.', parents=[common_parser])
|
|
auth_parser.add_argument('--policy-file', dest='policy', default=POLICY_FILE_AUTH, help=f'Path to the auth simulation policy file. Default: {POLICY_FILE_AUTH}')
|
|
auth_parser.add_argument('--set', action='append', default=[], help="Override a policy setting using 'key.subkey=value' format.")
|
|
|
|
# --- Download ---
|
|
download_parser = sim_subparsers.add_parser('download', help='Run the download part of the simulation.', parents=[common_parser])
|
|
download_parser.add_argument('--policy-file', dest='policy', default=POLICY_FILE_DOWNLOAD, help=f'Path to the download simulation policy file. Default: {POLICY_FILE_DOWNLOAD}')
|
|
download_parser.add_argument('--set', action='append', default=[], help="Override a policy setting using 'key.subkey=value' format.")
|
|
|
|
def main_simulation(args):
|
|
"""Main dispatcher for 'simulation' command."""
|
|
# --- Load policy to get simulation parameters ---
|
|
policy = {}
|
|
# The 'policy' attribute is guaranteed to exist by the arg parser for all subcommands
|
|
if not yaml:
|
|
logger.error("Cannot load policy file because PyYAML is not installed.")
|
|
return 1
|
|
try:
|
|
with open(args.policy, 'r') as f:
|
|
# We only need the first document if it's a multi-policy file
|
|
policy = yaml.safe_load(f) or {}
|
|
except (IOError, yaml.YAMLError) as e:
|
|
logger.error(f"Failed to load or parse policy file {args.policy}: {e}")
|
|
return 1
|
|
|
|
sim_params = policy.get('simulation_parameters', {})
|
|
effective_env_file = args.env_file or sim_params.get('env_file')
|
|
|
|
if args.simulation_command == 'setup':
|
|
# Create an args object that main_setup_profiles expects
|
|
setup_args = SimpleNamespace(
|
|
policy_file=args.policy,
|
|
env_file=effective_env_file,
|
|
preserve_profiles=args.preserve_profiles,
|
|
reset_global_counters=args.reset_global_counters,
|
|
verbose=args.verbose,
|
|
redis_host=args.redis_host,
|
|
redis_port=args.redis_port,
|
|
redis_password=args.redis_password
|
|
)
|
|
return main_setup_profiles(setup_args)
|
|
|
|
elif args.simulation_command == 'auth':
|
|
# This command runs the stress tool in auth (fetch_only) mode.
|
|
# It is expected that the policy-enforcer is run as a separate process.
|
|
stress_args = SimpleNamespace(
|
|
policy=args.policy, policy_name=None, list_policies=False, show_overrides=False,
|
|
set=args.set, profile_prefix=None, start_from_url_index=None, auto_merge_fragments=None,
|
|
remove_fragments_after_merge=None, fragments_dir=None, remote_dir=None, cleanup=None,
|
|
verbose=args.verbose, dry_run=False, disable_log_writing=False,
|
|
# Redis connection args
|
|
env_file=effective_env_file, redis_host=args.redis_host, redis_port=args.redis_port,
|
|
redis_password=args.redis_password, env=args.env, key_prefix=None,
|
|
expire_time_shift_minutes=args.expire_time_shift_minutes
|
|
)
|
|
|
|
logger.info("\n--- Starting Auth Simulation (stress-policy) ---")
|
|
return main_stress_policy(stress_args)
|
|
|
|
elif args.simulation_command == 'download':
|
|
# This is simpler, just runs the stress tool in download mode.
|
|
stress_args = SimpleNamespace(
|
|
policy=args.policy, policy_name=None, list_policies=False, show_overrides=False,
|
|
set=args.set, profile_prefix=None, start_from_url_index=None, auto_merge_fragments=None,
|
|
remove_fragments_after_merge=None, fragments_dir=None, remote_dir=None, cleanup=None,
|
|
verbose=args.verbose, dry_run=False, disable_log_writing=False,
|
|
# Redis connection args
|
|
env_file=effective_env_file, redis_host=args.redis_host, redis_port=args.redis_port,
|
|
redis_password=args.redis_password, env=args.env, key_prefix=None,
|
|
expire_time_shift_minutes=args.expire_time_shift_minutes
|
|
)
|
|
logger.info("\n--- Starting Download Simulation (stress-policy) ---")
|
|
return main_stress_policy(stress_args)
|
|
|
|
return 1 # Should not be reached
|