1333 lines
75 KiB
Python
1333 lines
75 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
CLI tool to enforce policies on profiles.
|
|
"""
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import os
|
|
import signal
|
|
import sys
|
|
import time
|
|
|
|
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
|
|
|
|
from .profile_manager_tool import ProfileManager
|
|
|
|
# Configure logging
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Graceful shutdown handler
|
|
shutdown_event = False
|
|
def handle_shutdown(sig, frame):
|
|
global shutdown_event
|
|
logger.info("Shutdown signal received. Stopping policy enforcer...")
|
|
shutdown_event = True
|
|
|
|
class PolicyEnforcer:
|
|
def __init__(self, manager, dry_run=False):
|
|
self.manager = manager
|
|
self.dry_run = dry_run
|
|
self.actions_taken_this_cycle = 0
|
|
|
|
PROXY_REST_REASON = "Proxy resting"
|
|
|
|
def apply_policies(self, args):
|
|
self.actions_taken_this_cycle = 0
|
|
logger.debug(f"Applying policies... (Dry run: {self.dry_run})")
|
|
|
|
# --- Phase 1: Policies that don't depend on a consistent profile state snapshot ---
|
|
# Manage proxy states and clean up stale locks before we fetch profile states.
|
|
self.enforce_proxy_group_rotation(getattr(args, 'proxy_groups', []))
|
|
self.enforce_proxy_work_rest_cycle(args)
|
|
self.enforce_max_proxy_active_time(args)
|
|
self.enforce_proxy_policies(args)
|
|
if args.unlock_stale_locks_after_seconds and args.unlock_stale_locks_after_seconds > 0:
|
|
self.enforce_stale_lock_cleanup(args.unlock_stale_locks_after_seconds)
|
|
|
|
# --- Phase 2: Policies that require a consistent, shared view of profile states ---
|
|
# Fetch all profile states ONCE to create a consistent snapshot for this cycle.
|
|
all_profiles_list = self.manager.list_profiles()
|
|
all_profiles_map = {p['name']: p for p in all_profiles_list}
|
|
|
|
# Apply profile group policies (rotation, max_active). This will modify the local `all_profiles_map`.
|
|
self.enforce_profile_group_policies(getattr(args, 'profile_groups', []), all_profiles_map)
|
|
|
|
# Un-rest profiles. This also reads from and modifies the local `all_profiles_map`.
|
|
self.enforce_unrest_policy(getattr(args, 'profile_groups', []), all_profiles_map)
|
|
|
|
# --- Phase 3: Apply policies to individual active profiles ---
|
|
# Use the now-updated snapshot to determine which profiles are active.
|
|
active_profiles = [p for p in all_profiles_map.values() if p['state'] == self.manager.STATE_ACTIVE]
|
|
|
|
# Filter out profiles that are managed by a profile group, as their state is handled separately.
|
|
profile_groups = getattr(args, 'profile_groups', [])
|
|
if profile_groups:
|
|
grouped_profiles = set()
|
|
for group in profile_groups:
|
|
if 'profiles' in group:
|
|
for p_name in group['profiles']:
|
|
grouped_profiles.add(p_name)
|
|
elif 'prefix' in group:
|
|
prefix = group['prefix']
|
|
for p in all_profiles_list:
|
|
if p['name'].startswith(prefix):
|
|
grouped_profiles.add(p['name'])
|
|
|
|
original_count = len(active_profiles)
|
|
active_profiles = [p for p in active_profiles if p['name'] not in grouped_profiles]
|
|
if len(active_profiles) != original_count:
|
|
logger.debug(f"Filtered out {original_count - len(active_profiles)} profile(s) managed by profile groups.")
|
|
|
|
for profile in active_profiles:
|
|
# Check for failure burst first, as it's more severe.
|
|
# If it's banned, no need to check other rules for it.
|
|
if self.enforce_failure_burst_policy(profile, args.ban_on_failures, args.ban_on_failures_window_minutes):
|
|
continue
|
|
if self.enforce_rate_limit_policy(profile, getattr(args, 'rate_limit_requests', 0), getattr(args, 'rate_limit_window_minutes', 0), getattr(args, 'rate_limit_rest_duration_minutes', 0)):
|
|
continue
|
|
self.enforce_failure_rate_policy(profile, args.max_failure_rate, args.min_requests_for_rate)
|
|
self.enforce_rest_policy(profile, args.rest_after_requests, args.rest_duration_minutes)
|
|
|
|
return self.actions_taken_this_cycle > 0
|
|
|
|
def enforce_failure_burst_policy(self, profile, max_failures, window_minutes):
|
|
if not max_failures or not window_minutes or max_failures <= 0 or window_minutes <= 0:
|
|
return False
|
|
|
|
window_seconds = window_minutes * 60
|
|
# Count only fatal error types (auth, download) for the ban policy.
|
|
# Tolerated errors are excluded from this check.
|
|
error_count = (
|
|
self.manager.get_activity_rate(profile['name'], 'failure', window_seconds) +
|
|
self.manager.get_activity_rate(profile['name'], 'download_error', window_seconds)
|
|
)
|
|
|
|
if error_count >= max_failures:
|
|
reason = f"Error burst detected: {error_count} errors in the last {window_minutes} minute(s) (threshold: {max_failures})"
|
|
logger.warning(f"Banning profile '{profile['name']}' due to error burst: {reason}")
|
|
if not self.dry_run:
|
|
self.manager.update_profile_state(profile['name'], self.manager.STATE_BANNED, reason)
|
|
self.actions_taken_this_cycle += 1
|
|
return True # Indicates profile was banned
|
|
return False
|
|
|
|
def enforce_rate_limit_policy(self, profile, max_requests, window_minutes, rest_duration_minutes):
|
|
if not max_requests or not window_minutes or max_requests <= 0 or window_minutes <= 0:
|
|
return False
|
|
|
|
window_seconds = window_minutes * 60
|
|
# Count all successful activities (auth, download) for rate limiting.
|
|
# We don't count failures, as they often don't hit the target server in the same way.
|
|
activity_count = (
|
|
self.manager.get_activity_rate(profile['name'], 'success', window_seconds) +
|
|
self.manager.get_activity_rate(profile['name'], 'download', window_seconds)
|
|
)
|
|
|
|
if activity_count >= max_requests:
|
|
reason = f"Rate limit hit: {activity_count} requests in last {window_minutes} minute(s) (limit: {max_requests})"
|
|
logger.info(f"Resting profile '{profile['name']}' for {rest_duration_minutes}m: {reason}")
|
|
if not self.dry_run:
|
|
self.manager.update_profile_state(profile['name'], self.manager.STATE_RESTING, reason)
|
|
rest_until = time.time() + rest_duration_minutes * 60
|
|
self.manager.update_profile_field(profile['name'], 'rest_until', str(rest_until))
|
|
self.actions_taken_this_cycle += 1
|
|
return True # Indicates profile was rested
|
|
return False
|
|
|
|
def enforce_unrest_policy(self, profile_groups, all_profiles_map):
|
|
all_profiles_list = list(all_profiles_map.values())
|
|
resting_profiles = [p for p in all_profiles_list if p['state'] == self.manager.STATE_RESTING]
|
|
cooldown_profiles = [p for p in all_profiles_list if p['state'] == self.manager.STATE_COOLDOWN]
|
|
profiles_to_check = resting_profiles + cooldown_profiles
|
|
now = time.time()
|
|
|
|
if not profiles_to_check:
|
|
return
|
|
|
|
# Sort profiles to check by their rest_until timestamp, then by name.
|
|
# This creates a deterministic FIFO queue for activation.
|
|
profiles_to_check.sort(key=lambda p: (p.get('rest_until', 0), p.get('name', '')))
|
|
|
|
# --- Group-aware unrest logic ---
|
|
profile_to_group_map = {}
|
|
group_to_profiles_map = {}
|
|
if profile_groups:
|
|
for group in profile_groups:
|
|
group_name = group.get('name')
|
|
if not group_name: continue
|
|
|
|
profiles_in_group = []
|
|
if 'profiles' in group:
|
|
profiles_in_group = sorted(group['profiles'])
|
|
elif 'prefix' in group:
|
|
prefix = group['prefix']
|
|
profiles_in_group = sorted([p['name'] for p in all_profiles_list if p['name'].startswith(prefix)])
|
|
|
|
group_to_profiles_map[group_name] = profiles_in_group
|
|
for p_name in profiles_in_group:
|
|
profile_to_group_map[p_name] = group_name
|
|
|
|
# This will store the live count of active profiles for each group,
|
|
# preventing race conditions within a single enforcer run.
|
|
live_active_counts = {}
|
|
if profile_groups:
|
|
for group_name, profiles_in_group in group_to_profiles_map.items():
|
|
count = 0
|
|
for p_name in profiles_in_group:
|
|
profile_state = all_profiles_map.get(p_name, {}).get('state')
|
|
if profile_state in [self.manager.STATE_ACTIVE, self.manager.STATE_LOCKED, self.manager.STATE_COOLDOWN]:
|
|
count += 1
|
|
live_active_counts[group_name] = count
|
|
# --- End group logic setup ---
|
|
|
|
# --- New logic: Identify groups with waiting profiles ---
|
|
groups_with_waiting_profiles = {}
|
|
if profile_groups:
|
|
for group in profile_groups:
|
|
group_name = group.get('name')
|
|
if not group_name: continue
|
|
|
|
defer_activation = group.get('defer_activation_if_any_waiting', False)
|
|
if not defer_activation: continue
|
|
|
|
profiles_in_group = group_to_profiles_map.get(group_name, [])
|
|
waiting_profile = next(
|
|
(p for p_name, p in all_profiles_map.items()
|
|
if p_name in profiles_in_group and p.get('rest_reason') == 'waiting_downloads'),
|
|
None
|
|
)
|
|
if waiting_profile:
|
|
groups_with_waiting_profiles[group_name] = waiting_profile['name']
|
|
# --- End new logic ---
|
|
|
|
unique_proxies = sorted(list(set(p['proxy'] for p in profiles_to_check if p.get('proxy'))))
|
|
proxy_states = self.manager.get_proxy_states(unique_proxies)
|
|
|
|
for profile in profiles_to_check:
|
|
profile_name = profile['name']
|
|
group_name = profile_to_group_map.get(profile_name)
|
|
|
|
# --- New logic: Defer activation if group has a waiting profile ---
|
|
if group_name in groups_with_waiting_profiles:
|
|
waiting_profile_name = groups_with_waiting_profiles[group_name]
|
|
if profile_name != waiting_profile_name:
|
|
logger.debug(f"Profile '{profile_name}' activation deferred because profile '{waiting_profile_name}' in group '{group_name}' is waiting for downloads.")
|
|
continue
|
|
# --- End new logic ---
|
|
|
|
# --- New logic for waiting_downloads ---
|
|
if profile.get('rest_reason') == 'waiting_downloads':
|
|
profile_name = profile['name']
|
|
group_name = profile_to_group_map.get(profile_name)
|
|
group_policy = next((g for g in profile_groups if g.get('name') == group_name), {})
|
|
|
|
max_wait_minutes = group_policy.get('max_wait_for_downloads_minutes', 240)
|
|
wait_started_at = profile.get('wait_started_at', 0)
|
|
|
|
downloads_pending = self.manager.get_pending_downloads(profile_name)
|
|
is_timed_out = (time.time() - wait_started_at) > (max_wait_minutes * 60) if wait_started_at > 0 else False
|
|
|
|
if downloads_pending <= 0 or is_timed_out:
|
|
if is_timed_out:
|
|
logger.warning(f"Profile '{profile_name}' download wait timed out after {max_wait_minutes}m. Forcing rotation.")
|
|
else:
|
|
logger.info(f"All pending downloads for profile '{profile_name}' are complete. Proceeding with rotation.")
|
|
|
|
self.actions_taken_this_cycle += 1
|
|
|
|
# Transition to a normal post-rotation rest period.
|
|
new_reason = "Rotation complete (downloads finished)"
|
|
rest_duration_minutes = group_policy.get('rest_duration_minutes_on_rotation', 0)
|
|
rest_until_ts = time.time() + (rest_duration_minutes * 60)
|
|
|
|
if not self.dry_run:
|
|
self.manager.update_profile_field(profile['name'], 'rest_reason', new_reason)
|
|
self.manager.update_profile_field(profile['name'], 'rest_until', str(rest_until_ts))
|
|
self.manager.clear_pending_downloads(profile_name)
|
|
# Reset counters to ensure a fresh start after the wait period.
|
|
self.manager.reset_profile_counters(profile_name)
|
|
|
|
# Update local map so it can be activated in the same cycle if rest_duration is 0
|
|
all_profiles_map[profile_name]['rest_reason'] = new_reason
|
|
all_profiles_map[profile_name]['rest_until'] = rest_until_ts
|
|
all_profiles_map[profile_name]['success_count'] = 0
|
|
all_profiles_map[profile_name]['failure_count'] = 0
|
|
all_profiles_map[profile_name]['tolerated_error_count'] = 0
|
|
all_profiles_map[profile_name]['download_count'] = 0
|
|
all_profiles_map[profile_name]['download_error_count'] = 0
|
|
|
|
# Let the rest of the unrest logic handle the activation now that rest_until is set.
|
|
profile['rest_until'] = rest_until_ts # Update profile in loop
|
|
else:
|
|
logger.debug(f"Profile '{profile_name}' is still waiting for {downloads_pending} download(s) to complete.")
|
|
continue # Skip to next profile, do not attempt to activate.
|
|
# --- End new logic ---
|
|
|
|
rest_until = profile.get('rest_until', 0)
|
|
if now >= rest_until:
|
|
profile_name = profile['name']
|
|
group_name = profile_to_group_map.get(profile_name)
|
|
|
|
# --- Group-aware unrest check ---
|
|
if group_name:
|
|
group_policy = next((g for g in profile_groups if g.get('name') == group_name), None)
|
|
if not group_policy:
|
|
continue # Should not happen if maps are built correctly
|
|
|
|
max_active = group_policy.get('max_active_profiles', 1)
|
|
|
|
# Check if the group is already at its capacity for active profiles.
|
|
# We use the live counter which is updated during this enforcer cycle.
|
|
|
|
# Special handling for COOLDOWN profiles: they should be allowed to become ACTIVE
|
|
# even if the group is at capacity, because they are already counted as "active".
|
|
# We check if the group would be over capacity *without* this profile.
|
|
is_cooldown_profile = profile['state'] == self.manager.STATE_COOLDOWN
|
|
effective_active_count = live_active_counts.get(group_name, 0)
|
|
|
|
# If we are considering a COOLDOWN profile, it's already in the count.
|
|
# The check should be if activating it would exceed the limit, assuming
|
|
# no *other* profile is active.
|
|
capacity_check_count = effective_active_count
|
|
if is_cooldown_profile:
|
|
capacity_check_count -= 1
|
|
|
|
if capacity_check_count >= max_active:
|
|
logger.debug(f"Profile '{profile_name}' rest ended, but group '{group_name}' is at capacity ({effective_active_count}/{max_active}). Deferring activation.")
|
|
|
|
# If a profile's COOLDOWN ends but it can't be activated (because another
|
|
# profile is active), move it to RESTING so it's clear it's waiting for capacity.
|
|
if is_cooldown_profile:
|
|
reason = "Waiting for group capacity"
|
|
logger.info(f"Profile '{profile_name}' cooldown ended but group is full. Moving to RESTING to wait for a slot.")
|
|
self.actions_taken_this_cycle += 1
|
|
if not self.dry_run:
|
|
# We need to update state but keep rest_until in the past.
|
|
self.manager.update_profile_state(profile['name'], self.manager.STATE_RESTING, reason)
|
|
self.manager.update_profile_field(profile['name'], 'rest_until', '0')
|
|
|
|
# Update local map
|
|
all_profiles_map[profile_name]['state'] = self.manager.STATE_RESTING
|
|
all_profiles_map[profile_name]['rest_until'] = 0
|
|
all_profiles_map[profile_name]['rest_reason'] = reason
|
|
|
|
continue # Do not activate, group is full.
|
|
else:
|
|
# Defensive check for orphaned profiles that should be in a group.
|
|
# This can happen if list_profiles() returns an incomplete list for one cycle,
|
|
# causing the group maps to be incomplete. This check prevents a "stampede"
|
|
# of activations that would violate group limits.
|
|
is_orphan = False
|
|
for group in profile_groups:
|
|
prefix = group.get('prefix')
|
|
if prefix and profile_name.startswith(prefix):
|
|
is_orphan = True
|
|
logger.warning(
|
|
f"Profile '{profile_name}' appears to belong to group '{group.get('name')}' "
|
|
f"but was not found in the initial scan. Deferring activation to prevent violating group limits."
|
|
)
|
|
break
|
|
if is_orphan:
|
|
continue # Skip activation for this profile
|
|
# --- End group check ---
|
|
|
|
# Before activating, ensure the profile's proxy is not resting.
|
|
proxy_url = profile.get('proxy')
|
|
if proxy_url:
|
|
proxy_state_data = proxy_states.get(proxy_url, {})
|
|
if proxy_state_data.get('state') == self.manager.STATE_RESTING:
|
|
logger.debug(f"Profile '{profile['name']}' rest period ended, but its proxy '{proxy_url}' is still resting. Deferring activation.")
|
|
|
|
# Update reason for clarity in the UI when a profile is blocked by its proxy.
|
|
new_reason = "Waiting for proxy"
|
|
if profile.get('rest_reason') != new_reason:
|
|
logger.info(f"Updating profile '{profile['name']}' reason to '{new_reason}'.")
|
|
self.actions_taken_this_cycle += 1
|
|
if not self.dry_run:
|
|
self.manager.update_profile_field(profile['name'], 'rest_reason', new_reason)
|
|
# Update local map for consistency within this enforcer cycle.
|
|
all_profiles_map[profile_name]['rest_reason'] = new_reason
|
|
|
|
continue # Do not activate this profile yet.
|
|
|
|
# Update group counter BEFORE making any changes, so subsequent checks in this cycle use the updated count
|
|
if group_name and profile['state'] == self.manager.STATE_RESTING:
|
|
# For RESTING profiles, they're becoming active, so increment the count
|
|
live_active_counts[group_name] = live_active_counts.get(group_name, 0) + 1
|
|
# COOLDOWN profiles are already counted, no change needed
|
|
|
|
logger.info(f"Activating profile '{profile['name']}' (rest period completed).")
|
|
self.actions_taken_this_cycle += 1
|
|
|
|
# Determine if this was a true rest or just a cooldown that was waiting for a slot.
|
|
is_waiting_after_cooldown = profile.get('rest_reason') == "Waiting for group capacity"
|
|
|
|
if not self.dry_run:
|
|
# When un-resting from a long rest, reset counters to give it a fresh start.
|
|
# Do not reset for COOLDOWN or a profile that was waiting after cooldown.
|
|
self.manager.update_profile_state(profile['name'], self.manager.STATE_ACTIVE, "Rest period completed")
|
|
if profile['state'] == self.manager.STATE_RESTING and not is_waiting_after_cooldown:
|
|
self.manager.reset_profile_counters(profile['name'])
|
|
|
|
# Update the shared map to reflect the change immediately for this cycle.
|
|
all_profiles_map[profile_name]['state'] = self.manager.STATE_ACTIVE
|
|
if profile['state'] == self.manager.STATE_RESTING and not is_waiting_after_cooldown:
|
|
all_profiles_map[profile_name]['success_count'] = 0
|
|
all_profiles_map[profile_name]['failure_count'] = 0
|
|
all_profiles_map[profile_name]['tolerated_error_count'] = 0
|
|
all_profiles_map[profile_name]['download_count'] = 0
|
|
all_profiles_map[profile_name]['download_error_count'] = 0
|
|
|
|
def enforce_failure_rate_policy(self, profile, max_failure_rate, min_requests):
|
|
if max_failure_rate <= 0:
|
|
return
|
|
|
|
success = profile.get('global_success_count', 0)
|
|
failure = profile.get('global_failure_count', 0)
|
|
total = success + failure
|
|
|
|
if total < min_requests:
|
|
return
|
|
|
|
current_failure_rate = failure / total if total > 0 else 0
|
|
|
|
if current_failure_rate >= max_failure_rate:
|
|
reason = f"Global failure rate {current_failure_rate:.2f} >= threshold {max_failure_rate} ({int(failure)}/{int(total)} failures)"
|
|
logger.warning(f"Banning profile '{profile['name']}' due to high failure rate: {reason}")
|
|
if not self.dry_run:
|
|
self.manager.update_profile_state(profile['name'], self.manager.STATE_BANNED, reason)
|
|
self.actions_taken_this_cycle += 1
|
|
|
|
def enforce_rest_policy(self, profile, rest_after_requests, rest_duration_minutes):
|
|
if not rest_after_requests or rest_after_requests <= 0 or not rest_duration_minutes or rest_duration_minutes <= 0:
|
|
return
|
|
|
|
total_requests = (
|
|
int(profile.get('success_count', 0)) +
|
|
int(profile.get('failure_count', 0)) +
|
|
int(profile.get('tolerated_error_count', 0)) +
|
|
int(profile.get('download_count', 0)) +
|
|
int(profile.get('download_error_count', 0))
|
|
)
|
|
|
|
if total_requests >= rest_after_requests:
|
|
reason = f"Request count {total_requests} >= threshold {rest_after_requests}"
|
|
logger.info(f"Resting profile '{profile['name']}' for {rest_duration_minutes}m: {reason}")
|
|
self.actions_taken_this_cycle += 1
|
|
if not self.dry_run:
|
|
self.manager.update_profile_state(profile['name'], self.manager.STATE_RESTING, reason)
|
|
rest_until = time.time() + rest_duration_minutes * 60
|
|
self.manager.update_profile_field(profile['name'], 'rest_until', str(rest_until))
|
|
|
|
def enforce_stale_lock_cleanup(self, max_lock_seconds):
|
|
"""Finds and unlocks profiles with stale locks."""
|
|
if self.dry_run:
|
|
logger.info(f"[Dry Run] Would check for and clean up locks older than {max_lock_seconds} seconds.")
|
|
return
|
|
|
|
cleaned_count = self.manager.cleanup_stale_locks(max_lock_seconds)
|
|
if cleaned_count > 0:
|
|
self.actions_taken_this_cycle += cleaned_count
|
|
|
|
def enforce_profile_group_policies(self, profile_groups, all_profiles_map):
|
|
"""
|
|
Manages profiles within defined groups. This includes:
|
|
1. Rotating out profiles that have met their request limit.
|
|
2. Healing the group by ensuring no more than `max_active_profiles` are active.
|
|
3. Initializing the group by activating a profile if none are active.
|
|
|
|
This method operates on and modifies the `all_profiles_map` passed to it.
|
|
"""
|
|
if not profile_groups:
|
|
return
|
|
|
|
all_profiles_list = list(all_profiles_map.values())
|
|
|
|
for group in profile_groups:
|
|
group_name = group.get('name')
|
|
if not group_name:
|
|
logger.warning("Found a profile group without a 'name'. Skipping.")
|
|
continue
|
|
|
|
profiles_in_group = set()
|
|
if 'profiles' in group:
|
|
profiles_in_group = set(group['profiles'])
|
|
elif 'prefix' in group:
|
|
prefix = group['prefix']
|
|
profiles_in_group = {p['name'] for p in all_profiles_list if p['name'].startswith(prefix)}
|
|
|
|
if not profiles_in_group:
|
|
logger.warning(f"Profile group '{group_name}' has no matching profiles. Skipping.")
|
|
continue
|
|
|
|
# --- Persist group policy to Redis for observability ---
|
|
rotate_after_requests = group.get('rotate_after_requests')
|
|
max_active_profiles = group.get('max_active_profiles')
|
|
if not self.dry_run:
|
|
# This is a non-critical update, so we don't need to check for existence.
|
|
# We just update it on every cycle to ensure it's fresh.
|
|
self.manager.set_profile_group_state(group_name, {
|
|
'rotate_after_requests': rotate_after_requests,
|
|
'max_active_profiles': max_active_profiles,
|
|
'prefix': group.get('prefix') # Store prefix for observability
|
|
})
|
|
|
|
# --- 1. Handle Rotation for Active Profiles ---
|
|
rotate_after_requests = group.get('rotate_after_requests')
|
|
if rotate_after_requests and rotate_after_requests > 0:
|
|
# Consider ACTIVE, LOCKED, and COOLDOWN profiles for rotation eligibility.
|
|
eligible_for_rotation_check = [
|
|
p for p in all_profiles_list
|
|
if p['name'] in profiles_in_group and p['state'] in [self.manager.STATE_ACTIVE, self.manager.STATE_LOCKED, self.manager.STATE_COOLDOWN]
|
|
]
|
|
|
|
for profile in eligible_for_rotation_check:
|
|
total_requests = (
|
|
int(profile.get('success_count', 0)) +
|
|
int(profile.get('failure_count', 0)) +
|
|
int(profile.get('tolerated_error_count', 0)) +
|
|
int(profile.get('download_count', 0)) +
|
|
int(profile.get('download_error_count', 0))
|
|
)
|
|
if total_requests >= rotate_after_requests:
|
|
# If a profile is LOCKED, we can't rotate it yet.
|
|
# Instead, we update its reason to show that a rotation is pending.
|
|
if profile['state'] == self.manager.STATE_LOCKED:
|
|
pending_reason = f"Pending Rotation (requests: {total_requests}/{rotate_after_requests})"
|
|
# Only update if the reason is not already set, to avoid spamming Redis.
|
|
if profile.get('reason') != pending_reason:
|
|
logger.info(f"Profile '{profile['name']}' in group '{group_name}' is due for rotation but is LOCKED. Marking as pending.")
|
|
self.actions_taken_this_cycle += 1
|
|
if not self.dry_run:
|
|
self.manager.update_profile_field(profile['name'], 'reason', pending_reason)
|
|
else:
|
|
logger.debug(f"Profile '{profile['name']}' in group '{group_name}' is due for rotation but is currently LOCKED. Already marked as pending.")
|
|
continue
|
|
|
|
# If the profile is ACTIVE or in COOLDOWN, we can rotate it immediately.
|
|
reason = f"Rotated after {total_requests} requests (limit: {rotate_after_requests})"
|
|
logger.info(f"Rotating profile '{profile['name']}' in group '{group_name}': {reason}")
|
|
self.actions_taken_this_cycle += 1
|
|
|
|
wait_for_downloads = group.get('wait_download_finish_per_profile', False)
|
|
|
|
new_reason = reason
|
|
rest_until_ts = 0
|
|
|
|
if wait_for_downloads:
|
|
new_reason = "waiting_downloads"
|
|
logger.info(f"Profile '{profile['name']}' will wait for pending downloads to complete.")
|
|
else:
|
|
rest_duration_minutes = group.get('rest_duration_minutes_on_rotation')
|
|
if rest_duration_minutes and rest_duration_minutes > 0:
|
|
rest_until_ts = time.time() + rest_duration_minutes * 60
|
|
|
|
if not self.dry_run:
|
|
self.manager.update_profile_state(profile['name'], self.manager.STATE_RESTING, new_reason)
|
|
|
|
if wait_for_downloads:
|
|
self.manager.update_profile_field(profile['name'], 'wait_started_at', str(time.time()))
|
|
# Set rest_until to 0 to indicate it's not a time-based rest
|
|
self.manager.update_profile_field(profile['name'], 'rest_until', '0')
|
|
elif rest_until_ts > 0:
|
|
self.manager.update_profile_field(profile['name'], 'rest_until', str(rest_until_ts))
|
|
|
|
# Reset all session counters for the next active cycle
|
|
self.manager.reset_profile_counters(profile['name'])
|
|
|
|
# Update our local map so subsequent policies in this cycle see the change immediately.
|
|
all_profiles_map[profile['name']]['state'] = self.manager.STATE_RESTING
|
|
all_profiles_map[profile['name']]['rest_reason'] = new_reason
|
|
if wait_for_downloads:
|
|
all_profiles_map[profile['name']]['wait_started_at'] = time.time()
|
|
all_profiles_map[profile['name']]['rest_until'] = 0
|
|
else:
|
|
all_profiles_map[profile['name']]['rest_until'] = rest_until_ts
|
|
all_profiles_map[profile['name']]['success_count'] = 0
|
|
all_profiles_map[profile['name']]['failure_count'] = 0
|
|
all_profiles_map[profile['name']]['tolerated_error_count'] = 0
|
|
all_profiles_map[profile['name']]['download_count'] = 0
|
|
all_profiles_map[profile['name']]['download_error_count'] = 0
|
|
|
|
# --- 2. Self-Healing: Enforce max_active_profiles ---
|
|
max_active = group.get('max_active_profiles', 1)
|
|
|
|
# Get the current list of active/locked profiles from our potentially modified local map
|
|
# A profile is considered "active" for group limits if it is ACTIVE, LOCKED, or in COOLDOWN.
|
|
current_active_or_locked_profiles = [
|
|
p for name, p in all_profiles_map.items()
|
|
if name in profiles_in_group and p['state'] in [self.manager.STATE_ACTIVE, self.manager.STATE_LOCKED, self.manager.STATE_COOLDOWN]
|
|
]
|
|
|
|
num_active_or_locked = len(current_active_or_locked_profiles)
|
|
if num_active_or_locked > max_active:
|
|
logger.warning(f"Healing group '{group_name}': Found {num_active_or_locked} active/locked profiles, but max is {max_active}. Resting excess ACTIVE profiles.")
|
|
|
|
# We can only rest profiles that are in the ACTIVE state, not LOCKED.
|
|
profiles_that_can_be_rested = [p for p in current_active_or_locked_profiles if p['state'] == self.manager.STATE_ACTIVE]
|
|
|
|
# Sort to determine which profiles to rest. We prefer to rest profiles
|
|
# that have been used more. As a tie-breaker (especially for profiles
|
|
# with 0 requests), we rest the one that has been active the longest
|
|
# (oldest last_used timestamp).
|
|
profiles_that_can_be_rested.sort(key=lambda p: p.get('last_used', 0)) # Oldest first
|
|
profiles_that_can_be_rested.sort(key=lambda p: (
|
|
p.get('success_count', 0) + p.get('failure_count', 0) +
|
|
p.get('tolerated_error_count', 0) +
|
|
p.get('download_count', 0) + p.get('download_error_count', 0)
|
|
), reverse=True) # Most requests first
|
|
|
|
num_to_rest = num_active_or_locked - max_active
|
|
profiles_to_rest = profiles_that_can_be_rested[:num_to_rest]
|
|
for profile in profiles_to_rest:
|
|
req_count = (
|
|
profile.get('success_count', 0) + profile.get('failure_count', 0) +
|
|
profile.get('tolerated_error_count', 0) +
|
|
profile.get('download_count', 0) +
|
|
profile.get('download_error_count', 0)
|
|
)
|
|
logger.warning(f"Healing group '{group_name}': Resting profile '{profile['name']}' (request count: {req_count}).")
|
|
self.actions_taken_this_cycle += 1
|
|
if not self.dry_run:
|
|
self.manager.update_profile_state(profile['name'], self.manager.STATE_RESTING, "Group max_active healing")
|
|
# Give it a rest time of 0, so it's immediately eligible for activation
|
|
# by the unrest logic if the group has capacity.
|
|
rest_until_ts = 0
|
|
self.manager.update_profile_field(profile['name'], 'rest_until', str(rest_until_ts))
|
|
|
|
# --- 3. Initialization: Activate profiles if below capacity ---
|
|
# This is a fallback for initialization or if all profiles were rested/banned.
|
|
# The primary activation mechanism is in `enforce_unrest_policy`.
|
|
elif num_active_or_locked < max_active:
|
|
# Check if there are any non-active, non-banned, non-locked profiles to activate.
|
|
eligible_profiles = [p for name, p in all_profiles_map.items() if name in profiles_in_group and p['state'] not in [self.manager.STATE_ACTIVE, self.manager.STATE_BANNED, self.manager.STATE_LOCKED]]
|
|
if eligible_profiles:
|
|
# This is a simple initialization case. We don't activate here because
|
|
# `enforce_unrest_policy` will handle it more intelligently based on rest times.
|
|
# This block ensures that on the very first run, a group doesn't sit empty.
|
|
if num_active_or_locked == 0:
|
|
logger.debug(f"Group '{group_name}' has no active profiles. `enforce_unrest_policy` will attempt to activate one.")
|
|
|
|
def enforce_proxy_group_rotation(self, proxy_groups):
|
|
"""Manages mutually exclusive work cycles for proxies within defined groups."""
|
|
if not proxy_groups:
|
|
return
|
|
|
|
group_names = [g['name'] for g in proxy_groups if g.get('name')]
|
|
if not group_names:
|
|
return
|
|
|
|
group_states = self.manager.get_proxy_group_states(group_names)
|
|
now = time.time()
|
|
|
|
for group in proxy_groups:
|
|
group_name = group.get('name')
|
|
if not group_name:
|
|
logger.warning("Found a proxy group without a 'name'. Skipping.")
|
|
continue
|
|
|
|
proxies_in_group = group.get('proxies', [])
|
|
if not proxies_in_group:
|
|
logger.warning(f"Proxy group '{group_name}' has no proxies defined. Skipping.")
|
|
continue
|
|
|
|
work_minutes = group.get('work_minutes_per_proxy')
|
|
if not work_minutes or work_minutes <= 0:
|
|
logger.warning(f"Proxy group '{group_name}' is missing 'work_minutes_per_proxy'. Skipping.")
|
|
continue
|
|
|
|
if not self.dry_run:
|
|
for proxy_url in proxies_in_group:
|
|
self.manager.set_proxy_group_membership(proxy_url, group_name, work_minutes)
|
|
|
|
work_duration_seconds = work_minutes * 60
|
|
state = group_states.get(group_name, {})
|
|
|
|
if not state:
|
|
# First run for this group, initialize it
|
|
logger.info(f"Initializing new proxy group '{group_name}'. Activating first proxy '{proxies_in_group[0]}'.")
|
|
self.actions_taken_this_cycle += 1
|
|
active_proxy_index = 0
|
|
next_rotation_ts = now + work_duration_seconds
|
|
|
|
if not self.dry_run:
|
|
# Activate the first, rest the others
|
|
self.manager.set_proxy_state(proxies_in_group[0], self.manager.STATE_ACTIVE)
|
|
for i, proxy_url in enumerate(proxies_in_group):
|
|
if i != active_proxy_index:
|
|
# Rest indefinitely; group logic will activate it when its turn comes.
|
|
self.manager.set_proxy_state(proxy_url, self.manager.STATE_RESTING, rest_duration_minutes=99999)
|
|
|
|
self.manager.set_proxy_group_state(group_name, active_proxy_index, next_rotation_ts)
|
|
|
|
elif now >= state.get('next_rotation_timestamp', 0):
|
|
# Time to rotate
|
|
current_active_index = state.get('active_proxy_index', 0)
|
|
next_active_index = (current_active_index + 1) % len(proxies_in_group)
|
|
|
|
old_active_proxy = proxies_in_group[current_active_index]
|
|
new_active_proxy = proxies_in_group[next_active_index]
|
|
|
|
logger.info(f"Rotating proxy group '{group_name}': Deactivating '{old_active_proxy}', Activating '{new_active_proxy}'.")
|
|
self.actions_taken_this_cycle += 1
|
|
|
|
next_rotation_ts = now + work_duration_seconds
|
|
|
|
if not self.dry_run:
|
|
# Rest the old proxy
|
|
self.manager.set_proxy_state(old_active_proxy, self.manager.STATE_RESTING, rest_duration_minutes=99999)
|
|
# Activate the new one
|
|
self.manager.set_proxy_state(new_active_proxy, self.manager.STATE_ACTIVE)
|
|
# Update group state
|
|
self.manager.set_proxy_group_state(group_name, next_active_index, next_rotation_ts)
|
|
|
|
def enforce_proxy_work_rest_cycle(self, args):
|
|
"""Enforces a work/rest cycle on proxies based on time."""
|
|
work_minutes = args.proxy_work_minutes
|
|
rest_minutes = args.proxy_rest_duration_minutes
|
|
|
|
if not work_minutes or work_minutes <= 0 or not rest_minutes or rest_minutes <= 0:
|
|
return
|
|
|
|
# Get a flat list of all proxies managed by groups, so we can ignore them.
|
|
proxy_groups = getattr(args, 'proxy_groups', [])
|
|
grouped_proxies = set()
|
|
if proxy_groups:
|
|
for group in proxy_groups:
|
|
for proxy_url in group.get('proxies', []):
|
|
grouped_proxies.add(proxy_url)
|
|
|
|
all_profiles = self.manager.list_profiles()
|
|
if not all_profiles:
|
|
return
|
|
|
|
unique_proxies = sorted(list(set(p['proxy'] for p in all_profiles if p.get('proxy'))))
|
|
|
|
# Filter out proxies that are managed by the group rotation logic
|
|
proxies_to_manage = [p for p in unique_proxies if p not in grouped_proxies]
|
|
if not proxies_to_manage:
|
|
logger.debug("All unique proxies are managed by proxy groups. Skipping individual work/rest cycle enforcement.")
|
|
return
|
|
|
|
proxy_states = self.manager.get_proxy_states(proxies_to_manage)
|
|
now = time.time()
|
|
|
|
for proxy_url, state_data in proxy_states.items():
|
|
state = state_data.get('state', self.manager.STATE_ACTIVE)
|
|
|
|
# Un-rest logic
|
|
if state == self.manager.STATE_RESTING:
|
|
rest_until = state_data.get('rest_until', 0)
|
|
if now >= rest_until:
|
|
logger.info(f"Activating proxy '{proxy_url}' (rest period complete).")
|
|
self.actions_taken_this_cycle += 1
|
|
if not self.dry_run:
|
|
self.manager.set_proxy_state(proxy_url, self.manager.STATE_ACTIVE)
|
|
|
|
# Also activate any profiles that were resting due to this proxy
|
|
profiles_for_proxy = [p for p in all_profiles if p.get('proxy') == proxy_url]
|
|
for profile in profiles_for_proxy:
|
|
if profile['state'] == self.manager.STATE_RESTING and profile.get('rest_reason') == self.PROXY_REST_REASON:
|
|
logger.info(f"Activating profile '{profile['name']}' as its proxy '{proxy_url}' is now active.")
|
|
self.actions_taken_this_cycle += 1
|
|
if not self.dry_run:
|
|
self.manager.update_profile_state(profile['name'], self.manager.STATE_ACTIVE, "Proxy activated")
|
|
else:
|
|
# Proxy is still resting. Ensure any of its profiles that are ACTIVE are moved to RESTING.
|
|
# This catches profiles that were unlocked while their proxy was resting.
|
|
rest_until_ts = state_data.get('rest_until', 0)
|
|
profiles_for_proxy = [p for p in all_profiles if p.get('proxy') == proxy_url]
|
|
for profile in profiles_for_proxy:
|
|
if profile['state'] == self.manager.STATE_ACTIVE:
|
|
logger.info(f"Resting profile '{profile['name']}' as its proxy '{proxy_url}' is resting.")
|
|
self.actions_taken_this_cycle += 1
|
|
if not self.dry_run:
|
|
self.manager.update_profile_state(profile['name'], self.manager.STATE_RESTING, self.PROXY_REST_REASON)
|
|
self.manager.update_profile_field(profile['name'], 'rest_until', str(rest_until_ts))
|
|
|
|
# Rest logic
|
|
elif state == self.manager.STATE_ACTIVE:
|
|
work_start = state_data.get('work_start_timestamp', 0)
|
|
if work_start == 0: # Proxy was just created, start its work cycle
|
|
if not self.dry_run:
|
|
self.manager.set_proxy_state(proxy_url, self.manager.STATE_ACTIVE)
|
|
continue
|
|
|
|
work_duration_seconds = work_minutes * 60
|
|
active_duration = now - work_start
|
|
logger.debug(f"Proxy '{proxy_url}' has been active for {active_duration:.0f}s (limit: {work_duration_seconds}s).")
|
|
if active_duration >= work_duration_seconds:
|
|
logger.info(f"Resting proxy '{proxy_url}' for {rest_minutes}m (work period of {work_minutes}m complete).")
|
|
self.actions_taken_this_cycle += 1
|
|
|
|
rest_until_ts = time.time() + rest_minutes * 60
|
|
if not self.dry_run:
|
|
self.manager.set_proxy_state(proxy_url, self.manager.STATE_RESTING, rest_minutes)
|
|
|
|
# Also rest any active profiles using this proxy
|
|
profiles_for_proxy = [p for p in all_profiles if p.get('proxy') == proxy_url]
|
|
for profile in profiles_for_proxy:
|
|
if profile['state'] == self.manager.STATE_ACTIVE:
|
|
logger.info(f"Resting profile '{profile['name']}' as its proxy '{proxy_url}' is resting.")
|
|
self.actions_taken_this_cycle += 1
|
|
if not self.dry_run:
|
|
self.manager.update_profile_state(profile['name'], self.manager.STATE_RESTING, self.PROXY_REST_REASON)
|
|
self.manager.update_profile_field(profile['name'], 'rest_until', str(rest_until_ts))
|
|
|
|
def enforce_max_proxy_active_time(self, args):
|
|
"""
|
|
Enforces a global maximum active time for any proxy, regardless of group membership.
|
|
This acts as a safety net to prevent a proxy from being stuck in an ACTIVE state.
|
|
"""
|
|
max_active_minutes = args.max_global_proxy_active_minutes
|
|
rest_minutes = args.rest_duration_on_max_active
|
|
|
|
if not max_active_minutes or max_active_minutes <= 0:
|
|
return
|
|
|
|
all_profiles = self.manager.list_profiles()
|
|
if not all_profiles:
|
|
return
|
|
|
|
unique_proxies = sorted(list(set(p['proxy'] for p in all_profiles if p.get('proxy'))))
|
|
if not unique_proxies:
|
|
return
|
|
|
|
proxy_states = self.manager.get_proxy_states(unique_proxies)
|
|
now = time.time()
|
|
|
|
for proxy_url, state_data in proxy_states.items():
|
|
if state_data.get('state') == self.manager.STATE_ACTIVE:
|
|
work_start = state_data.get('work_start_timestamp', 0)
|
|
if work_start == 0:
|
|
continue # Just activated, timestamp not set yet.
|
|
|
|
active_duration_seconds = now - work_start
|
|
max_active_seconds = max_active_minutes * 60
|
|
|
|
if active_duration_seconds >= max_active_seconds:
|
|
reason = f"Exceeded max active time of {max_active_minutes}m"
|
|
logger.warning(f"Resting proxy '{proxy_url}' for {rest_minutes}m: {reason}")
|
|
self.actions_taken_this_cycle += 1
|
|
|
|
rest_until_ts = now + rest_minutes * 60
|
|
if not self.dry_run:
|
|
self.manager.set_proxy_state(proxy_url, self.manager.STATE_RESTING, rest_minutes)
|
|
|
|
# Also rest any active profiles using this proxy
|
|
profiles_for_proxy = [p for p in all_profiles if p.get('proxy') == proxy_url]
|
|
for profile in profiles_for_proxy:
|
|
if profile['state'] == self.manager.STATE_ACTIVE:
|
|
logger.info(f"Resting profile '{profile['name']}' as its proxy '{proxy_url}' is resting due to max active time.")
|
|
self.actions_taken_this_cycle += 1
|
|
if not self.dry_run:
|
|
self.manager.update_profile_state(profile['name'], self.manager.STATE_RESTING, self.PROXY_REST_REASON)
|
|
self.manager.update_profile_field(profile['name'], 'rest_until', str(rest_until_ts))
|
|
|
|
def enforce_proxy_policies(self, args):
|
|
proxy_ban_enabled = args.proxy_ban_on_failures and args.proxy_ban_on_failures > 0
|
|
proxy_rate_limit_enabled = getattr(args, 'proxy_rate_limit_requests', 0) > 0
|
|
if not proxy_ban_enabled and not proxy_rate_limit_enabled:
|
|
return
|
|
|
|
all_profiles = self.manager.list_profiles()
|
|
if not all_profiles:
|
|
return
|
|
|
|
unique_proxies = sorted(list(set(p['proxy'] for p in all_profiles if p.get('proxy'))))
|
|
|
|
if not unique_proxies:
|
|
return
|
|
|
|
logger.debug(f"Checking proxy policies for {len(unique_proxies)} unique proxies...")
|
|
|
|
for proxy_url in unique_proxies:
|
|
profiles_for_proxy = [p for p in all_profiles if p.get('proxy') == proxy_url]
|
|
if self.enforce_proxy_failure_burst_policy(
|
|
proxy_url,
|
|
profiles_for_proxy,
|
|
args.proxy_ban_on_failures,
|
|
args.proxy_ban_window_minutes
|
|
):
|
|
continue # Banned, no need for other checks
|
|
|
|
self.enforce_proxy_rate_limit_policy(
|
|
proxy_url,
|
|
profiles_for_proxy,
|
|
getattr(args, 'proxy_rate_limit_requests', 0),
|
|
getattr(args, 'proxy_rate_limit_window_minutes', 0),
|
|
getattr(args, 'proxy_rate_limit_rest_duration_minutes', 0)
|
|
)
|
|
|
|
def enforce_proxy_failure_burst_policy(self, proxy_url, profiles_for_proxy, max_failures, window_minutes):
|
|
if not max_failures or not window_minutes or max_failures <= 0 or window_minutes <= 0:
|
|
return False
|
|
|
|
window_seconds = window_minutes * 60
|
|
failure_count = self.manager.get_proxy_activity_rate(proxy_url, 'failure', window_seconds)
|
|
|
|
if failure_count >= max_failures:
|
|
reason = f"Proxy failure burst: {failure_count} failures in last {window_minutes}m (threshold: {max_failures})"
|
|
logger.warning(f"Banning {len(profiles_for_proxy)} profile(s) on proxy '{proxy_url}' due to failure burst: {reason}")
|
|
self.actions_taken_this_cycle += 1
|
|
|
|
if not self.dry_run:
|
|
for profile in profiles_for_proxy:
|
|
# Don't re-ban already banned profiles
|
|
if profile['state'] != self.manager.STATE_BANNED:
|
|
self.manager.update_profile_state(profile['name'], self.manager.STATE_BANNED, reason)
|
|
return True # Indicates action was taken
|
|
return False
|
|
|
|
def enforce_proxy_rate_limit_policy(self, proxy_url, profiles_for_proxy, max_requests, window_minutes, rest_duration_minutes):
|
|
if not max_requests or not window_minutes or max_requests <= 0 or window_minutes <= 0:
|
|
return False
|
|
|
|
window_seconds = window_minutes * 60
|
|
# Count all successful activities for the proxy
|
|
activity_count = (
|
|
self.manager.get_proxy_activity_rate(proxy_url, 'success', window_seconds) +
|
|
self.manager.get_proxy_activity_rate(proxy_url, 'download', window_seconds)
|
|
)
|
|
|
|
if activity_count >= max_requests:
|
|
reason = f"Proxy rate limit hit: {activity_count} requests in last {window_minutes}m (limit: {max_requests})"
|
|
logger.info(f"Resting proxy '{proxy_url}' for {rest_duration_minutes}m: {reason}")
|
|
self.actions_taken_this_cycle += 1
|
|
|
|
if not self.dry_run:
|
|
self.manager.set_proxy_state(proxy_url, self.manager.STATE_RESTING, rest_duration_minutes)
|
|
return True # Indicates action was taken
|
|
return False
|
|
|
|
def add_policy_enforcer_parser(subparsers):
|
|
"""Adds the parser for the 'policy-enforcer' command."""
|
|
parser = subparsers.add_parser(
|
|
'policy-enforcer',
|
|
description='Apply policies to profiles (ban, rest, etc.).',
|
|
formatter_class=argparse.RawTextHelpFormatter,
|
|
help='Apply policies to profiles (ban, rest, etc.).'
|
|
)
|
|
|
|
parser.add_argument('--policy', '--policy-file', dest='policy_file', help='Path to a YAML policy file to load default settings from.')
|
|
parser.add_argument('--env-file', help='Path to a .env file to load environment variables from.')
|
|
parser.add_argument('--redis-host', default=None, help='Redis host. Defaults to REDIS_HOST or MASTER_HOST_IP env var, or localhost.')
|
|
parser.add_argument('--redis-port', type=int, default=None, help='Redis port. Defaults to REDIS_PORT env var, or 6379.')
|
|
parser.add_argument('--redis-password', default=None, help='Redis password. Defaults to REDIS_PASSWORD env var.')
|
|
parser.add_argument('--env', default=None, help="Default environment name for Redis key prefix. Used if --auth-env or --download-env are not specified. Overrides policy file setting.")
|
|
parser.add_argument('--auth-env', help="Override the environment for the Auth simulation.")
|
|
parser.add_argument('--download-env', help="Override the environment for the Download simulation.")
|
|
parser.add_argument('--legacy', action='store_true', help="Use legacy key prefix ('profile_mgmt_') without environment.")
|
|
parser.add_argument('--key-prefix', default=None, help='Explicit key prefix for Redis. Overrides --env, --legacy and any defaults.')
|
|
parser.add_argument('--verbose', action='store_true', help='Enable verbose logging')
|
|
parser.add_argument('--dry-run', action='store_true', help='Show what would be done without making changes.')
|
|
|
|
# Policy arguments
|
|
policy_group = parser.add_argument_group('Policy Rules')
|
|
policy_group.add_argument('--max-failure-rate', type=float, default=None,
|
|
help='Ban a profile if its failure rate exceeds this value (0.0 to 1.0). Default: 0.5')
|
|
policy_group.add_argument('--min-requests-for-rate', type=int, default=None,
|
|
help='Minimum total requests before failure rate is calculated. Default: 20')
|
|
policy_group.add_argument('--ban-on-failures', type=int, default=None,
|
|
help='Ban a profile if it has this many failures within the time window (0 to disable). Default: 0')
|
|
policy_group.add_argument('--ban-on-failures-window-minutes', type=int, default=None,
|
|
help='The time window in minutes for the failure burst check. Default: 5')
|
|
policy_group.add_argument('--rest-after-requests', type=int, default=None,
|
|
help='Move a profile to RESTING after this many total requests (0 to disable). Default: 0')
|
|
policy_group.add_argument('--rest-duration-minutes', type=int, default=None,
|
|
help='How long a profile should rest. Default: 15')
|
|
policy_group.add_argument('--rate-limit-requests', type=int, default=None,
|
|
help='Rest a profile if it exceeds this many requests in the time window (0 to disable).')
|
|
policy_group.add_argument('--rate-limit-window-minutes', type=int, default=None,
|
|
help='The time window in minutes for the rate limit check.')
|
|
policy_group.add_argument('--rate-limit-rest-duration-minutes', type=int, default=None,
|
|
help='How long a profile should rest after hitting the rate limit.')
|
|
policy_group.add_argument('--unlock-stale-locks-after-seconds', type=int, default=None,
|
|
help='Unlock profiles that have been in a LOCKED state for more than this many seconds (0 to disable). Default: 120')
|
|
|
|
proxy_policy_group = parser.add_argument_group('Proxy Policy Rules')
|
|
proxy_policy_group.add_argument('--proxy-work-minutes', type=int, default=None,
|
|
help='Work duration for a proxy before it rests (0 to disable). Default: 0')
|
|
proxy_policy_group.add_argument('--proxy-rest-duration-minutes', type=int, default=None,
|
|
help='Rest duration for a proxy after its work period. Default: 0')
|
|
proxy_policy_group.add_argument('--proxy-ban-on-failures', type=int, default=None,
|
|
help='Ban a proxy (and all its profiles) if it has this many failures within the time window (0 to disable). Default: 0')
|
|
proxy_policy_group.add_argument('--proxy-ban-window-minutes', type=int, default=None,
|
|
help='The time window in minutes for the proxy failure burst check. Default: 10')
|
|
proxy_policy_group.add_argument('--proxy-rate-limit-requests', type=int, default=None,
|
|
help='Rest a proxy if it exceeds this many requests in the time window (0 to disable).')
|
|
proxy_policy_group.add_argument('--proxy-rate-limit-window-minutes', type=int, default=None,
|
|
help='The time window in minutes for the proxy rate limit check.')
|
|
proxy_policy_group.add_argument('--proxy-rate-limit-rest-duration-minutes', type=int, default=None,
|
|
help='How long a proxy should rest after hitting the rate limit.')
|
|
proxy_policy_group.add_argument('--max-global-proxy-active-minutes', type=int, default=None,
|
|
help='Global maximum time a proxy can be active before being rested (0 to disable). Acts as a safety net. Default: 0')
|
|
proxy_policy_group.add_argument('--rest-duration-on-max-active', type=int, default=None,
|
|
help='How long a proxy should rest after hitting the global max active time. Default: 10')
|
|
|
|
# Execution control
|
|
exec_group = parser.add_argument_group('Execution Control')
|
|
exec_group.add_argument('--live', action='store_true', help='Run continuously, applying policies periodically.')
|
|
exec_group.add_argument('--interval-seconds', type=int, default=None,
|
|
help='When in --live mode, how often to apply policies. Default: 60')
|
|
exec_group.add_argument('--auth-only', action='store_true', help='Run enforcer for the auth simulation only.')
|
|
exec_group.add_argument('--download-only', action='store_true', help='Run enforcer for the download simulation only.')
|
|
|
|
return parser
|
|
|
|
def sync_cross_simulation(auth_manager, download_manager, sync_config, dry_run=False):
|
|
"""Synchronize profile states between auth and download simulations."""
|
|
if not sync_config:
|
|
return
|
|
|
|
profile_links = sync_config.get('profile_links', [])
|
|
sync_states = sync_config.get('sync_states', [])
|
|
sync_rotation = sync_config.get('sync_rotation', False)
|
|
enforce_auth_lead = sync_config.get('enforce_auth_lead', False)
|
|
|
|
if not profile_links:
|
|
return
|
|
|
|
# --- Get all profiles once for efficiency ---
|
|
all_auth_profiles = {p['name']: p for p in auth_manager.list_profiles()}
|
|
all_download_profiles = {p['name']: p for p in download_manager.list_profiles()}
|
|
|
|
# --- State and Rotation Sync (handles prefixes correctly) ---
|
|
for link in profile_links:
|
|
auth_prefix = link.get('auth')
|
|
download_prefix = link.get('download')
|
|
if not auth_prefix or not download_prefix:
|
|
continue
|
|
|
|
auth_profiles_in_group = [p for name, p in all_auth_profiles.items() if name.startswith(auth_prefix)]
|
|
|
|
for auth_profile in auth_profiles_in_group:
|
|
# Assume 1-to-1 name mapping (e.g., auth 'user1_0' maps to download 'user1_0')
|
|
download_profile_name = auth_profile['name']
|
|
download_profile = all_download_profiles.get(download_profile_name)
|
|
|
|
if not download_profile:
|
|
logger.debug(f"Auth profile '{auth_profile['name']}' has no corresponding download profile.")
|
|
continue
|
|
|
|
auth_state = auth_profile.get('state')
|
|
download_state = download_profile.get('state')
|
|
|
|
# Sync states from auth to download
|
|
if enforce_auth_lead and auth_state in sync_states and download_state != auth_state:
|
|
auth_reason = auth_profile.get('reason', '')
|
|
# If auth profile is waiting for downloads, we must NOT sync the RESTING state to the download profile,
|
|
# as that would prevent it from processing the very downloads we are waiting for.
|
|
if auth_state == auth_manager.STATE_RESTING and auth_reason == 'waiting_downloads':
|
|
logger.debug(f"Auth profile '{auth_profile['name']}' is waiting for downloads. Skipping state sync to download profile to prevent deadlock.")
|
|
else:
|
|
logger.info(f"Syncing download profile '{download_profile_name}' to state '{auth_state}' (auth lead)")
|
|
if not dry_run:
|
|
reason_to_sync = auth_reason or 'Synced from auth'
|
|
download_manager.update_profile_state(download_profile_name, auth_state, f"Synced from auth: {reason_to_sync}")
|
|
if auth_state == auth_manager.STATE_RESTING:
|
|
auth_rest_until = auth_profile.get('rest_until')
|
|
if auth_rest_until:
|
|
download_manager.update_profile_field(download_profile_name, 'rest_until', str(auth_rest_until))
|
|
|
|
# Handle rotation sync
|
|
if sync_rotation:
|
|
auth_reason = auth_profile.get('rest_reason', '')
|
|
|
|
# If auth profile is waiting for downloads, we must NOT sync the RESTING state to the download profile,
|
|
# as that would prevent it from processing the very downloads we are waiting for.
|
|
if auth_reason == 'waiting_downloads':
|
|
logger.debug(f"Auth profile '{auth_profile['name']}' is waiting for downloads. Skipping rotation sync to download profile to prevent deadlock.")
|
|
elif auth_state == auth_manager.STATE_RESTING and 'rotate' in auth_reason.lower():
|
|
if download_state != download_manager.STATE_RESTING:
|
|
logger.info(f"Rotating download profile '{download_profile_name}' due to auth rotation")
|
|
if not dry_run:
|
|
download_manager.update_profile_state(download_profile_name, download_manager.STATE_RESTING, f"Rotated due to auth rotation: {auth_reason}")
|
|
auth_rest_until = auth_profile.get('rest_until')
|
|
if auth_rest_until:
|
|
download_manager.update_profile_field(download_profile_name, 'rest_until', str(auth_rest_until))
|
|
|
|
# --- Active Profile Sync ---
|
|
sync_active = sync_config.get('sync_active_profile', False)
|
|
sync_waiting_downloads = sync_config.get('sync_waiting_downloads', False)
|
|
|
|
if not (sync_active or sync_waiting_downloads):
|
|
return
|
|
|
|
logger.debug("Syncing active profiles from Auth to Download simulation...")
|
|
|
|
# Get profiles that should be active in the download simulation
|
|
target_active_download_profiles = set()
|
|
|
|
# 1. Add profiles that are active in auth simulation (if sync_active is enabled)
|
|
if sync_active:
|
|
active_auth_profiles = [p for p in all_auth_profiles.values() if p['state'] in [auth_manager.STATE_ACTIVE, auth_manager.STATE_LOCKED]]
|
|
for auth_profile in active_auth_profiles:
|
|
target_active_download_profiles.add(auth_profile['name'])
|
|
|
|
# 2. Add profiles that are waiting for downloads to complete (if sync_waiting_downloads is enabled)
|
|
if sync_waiting_downloads:
|
|
waiting_auth_profiles = [p for p in all_auth_profiles.values()
|
|
if p['state'] == auth_manager.STATE_RESTING
|
|
and p.get('rest_reason') == 'waiting_downloads']
|
|
for auth_profile in waiting_auth_profiles:
|
|
target_active_download_profiles.add(auth_profile['name'])
|
|
logger.debug(f"Auth profile '{auth_profile['name']}' is waiting for downloads. Ensuring matching download profile is active.")
|
|
|
|
if not target_active_download_profiles:
|
|
logger.debug("No auth profiles found that need active download profiles.")
|
|
return
|
|
|
|
# Get download profile group info from Redis
|
|
dl_group_state_keys = [k for k in download_manager.redis.scan_iter(f"{download_manager.key_prefix}profile_group_state:*")]
|
|
dl_group_names = [k.split(':')[-1] for k in dl_group_state_keys]
|
|
dl_group_states = download_manager.get_profile_group_states(dl_group_names)
|
|
|
|
dl_profile_to_group = {}
|
|
for name, state in dl_group_states.items():
|
|
prefix = state.get('prefix')
|
|
if prefix:
|
|
for p_name in all_download_profiles:
|
|
if p_name.startswith(prefix):
|
|
dl_profile_to_group[p_name] = {'name': name, 'max_active': state.get('max_active_profiles', 1)}
|
|
|
|
# Activate download profiles that should be active but aren't
|
|
for target_profile_name in target_active_download_profiles:
|
|
download_profile = all_download_profiles.get(target_profile_name)
|
|
if not download_profile:
|
|
logger.warning(f"Auth profile '{target_profile_name}' needs an active download profile, but no corresponding download profile found.")
|
|
continue
|
|
|
|
if download_profile['state'] not in [download_manager.STATE_ACTIVE, download_manager.STATE_LOCKED]:
|
|
is_from_cooldown = download_profile['state'] == download_manager.STATE_COOLDOWN
|
|
log_msg_suffix = " (from COOLDOWN)" if is_from_cooldown else ""
|
|
logger.info(f"Syncing active state: Activating download profile '{target_profile_name}' to match auth requirements{log_msg_suffix}.")
|
|
if not dry_run:
|
|
download_manager.update_profile_state(target_profile_name, download_manager.STATE_ACTIVE, "Synced from auth requirements")
|
|
# Only reset counters if it's coming from a long-term rest, not a short cooldown.
|
|
if not is_from_cooldown:
|
|
download_manager.reset_profile_counters(target_profile_name)
|
|
|
|
# Deactivate any download profiles that are active but shouldn't be
|
|
for dl_profile_name, dl_profile in all_download_profiles.items():
|
|
if dl_profile['state'] == download_manager.STATE_ACTIVE and dl_profile_name not in target_active_download_profiles:
|
|
group_info = dl_profile_to_group.get(dl_profile_name)
|
|
if group_info:
|
|
logger.info(f"Syncing active state: Resting download profile '{dl_profile_name}' as it is no longer the active profile in its group.")
|
|
if not dry_run:
|
|
download_manager.update_profile_state(dl_profile_name, download_manager.STATE_RESTING, "Synced rotation from auth")
|
|
download_manager.update_profile_field(dl_profile_name, 'rest_until', '0')
|
|
|
|
def main_policy_enforcer(args):
|
|
"""Main dispatcher for 'policy-enforcer' command."""
|
|
policy = {}
|
|
if args.policy_file:
|
|
if not yaml:
|
|
logger.error("Cannot load policy file because PyYAML is not installed.")
|
|
return 1
|
|
try:
|
|
with open(args.policy_file, 'r') 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
|
|
|
|
class Config:
|
|
def __init__(self, cli_args, policy_defaults, code_defaults):
|
|
for key, code_default in code_defaults.items():
|
|
cli_val = getattr(cli_args, key, None)
|
|
policy_val = policy_defaults.get(key)
|
|
if cli_val is not None:
|
|
setattr(self, key, cli_val)
|
|
elif policy_val is not None:
|
|
setattr(self, key, policy_val)
|
|
else:
|
|
setattr(self, key, code_default)
|
|
|
|
code_defaults = {
|
|
'max_failure_rate': 0.0, 'min_requests_for_rate': 20, 'ban_on_failures': 0,
|
|
'ban_on_failures_window_minutes': 5, 'rest_after_requests': 0,
|
|
'rest_duration_minutes': 15,
|
|
'rate_limit_requests': 0, 'rate_limit_window_minutes': 60, 'rate_limit_rest_duration_minutes': 5,
|
|
'proxy_work_minutes': 0,
|
|
'proxy_rest_duration_minutes': 0, 'proxy_ban_on_failures': 0,
|
|
'proxy_ban_window_minutes': 10,
|
|
'proxy_rate_limit_requests': 0, 'proxy_rate_limit_window_minutes': 60, 'proxy_rate_limit_rest_duration_minutes': 10,
|
|
'unlock_stale_locks_after_seconds': 120,
|
|
'unlock_cooldown_seconds': 0,
|
|
'max_global_proxy_active_minutes': 0, 'rest_duration_on_max_active': 10,
|
|
'interval_seconds': 60, 'proxy_groups': [], 'profile_groups': []
|
|
}
|
|
|
|
sim_params = policy.get('simulation_parameters', {})
|
|
env_file_from_policy = sim_params.get('env_file')
|
|
|
|
if load_dotenv:
|
|
env_file = args.env_file or env_file_from_policy
|
|
if not env_file and args.env and '.env' in args.env and os.path.exists(args.env):
|
|
print(f"WARNING: --env should be an environment name, not a file path. Treating '{args.env}' as --env-file.", file=sys.stderr)
|
|
env_file = args.env
|
|
if env_file and load_dotenv(env_file):
|
|
print(f"Loaded environment variables from {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
|
|
|
|
redis_host = args.redis_host or os.getenv('REDIS_HOST', os.getenv('MASTER_HOST_IP', 'localhost'))
|
|
redis_port = args.redis_port if args.redis_port is not None else int(os.getenv('REDIS_PORT', 6379))
|
|
redis_password = args.redis_password or os.getenv('REDIS_PASSWORD')
|
|
|
|
if args.verbose:
|
|
logging.getLogger().setLevel(logging.DEBUG)
|
|
|
|
signal.signal(signal.SIGINT, handle_shutdown)
|
|
signal.signal(signal.SIGTERM, handle_shutdown)
|
|
|
|
enforcer_setups = []
|
|
|
|
def setup_enforcer(sim_type, env_cli_arg, policy_config_key, env_policy_key):
|
|
policy_config = policy.get(policy_config_key)
|
|
# Fallback for single-enforcer policy files
|
|
if policy_config is None and sim_type == 'Auth':
|
|
policy_config = policy.get('policy_enforcer_config', {})
|
|
|
|
if policy_config is None:
|
|
logger.debug(f"No config block found for {sim_type} simulation ('{policy_config_key}'). Skipping.")
|
|
return None
|
|
|
|
logger.info(f"Setting up enforcer for {sim_type} simulation...")
|
|
config = Config(args, policy_config, code_defaults)
|
|
|
|
# Determine the effective environment name with correct precedence:
|
|
# 1. Specific CLI arg (e.g., --auth-env)
|
|
# 2. General CLI arg (--env)
|
|
# 3. Specific policy setting (e.g., simulation_parameters.auth_env)
|
|
# 4. General policy setting (simulation_parameters.env)
|
|
# 5. Hardcoded default ('dev')
|
|
policy_env = sim_params.get(env_policy_key)
|
|
default_policy_env = sim_params.get('env')
|
|
effective_env = env_cli_arg or args.env or policy_env or default_policy_env or 'dev'
|
|
|
|
logger.info(f"Using environment '{effective_env}' for {sim_type}.")
|
|
|
|
if args.key_prefix: key_prefix = args.key_prefix
|
|
elif args.legacy: key_prefix = 'profile_mgmt_'
|
|
else: key_prefix = f"{effective_env}_profile_mgmt_"
|
|
|
|
manager = ProfileManager(redis_host, redis_port, redis_password, key_prefix)
|
|
enforcer = PolicyEnforcer(manager, dry_run=args.dry_run)
|
|
|
|
# Write any relevant config to Redis for workers to use
|
|
cooldown = getattr(config, 'unlock_cooldown_seconds', None)
|
|
if cooldown is not None and not args.dry_run:
|
|
# If it's a list or int, convert to JSON string to store in Redis
|
|
manager.set_config('unlock_cooldown_seconds', json.dumps(cooldown))
|
|
|
|
proxy_work_minutes = getattr(config, 'proxy_work_minutes', None)
|
|
if proxy_work_minutes is not None and not args.dry_run:
|
|
manager.set_config('proxy_work_minutes', proxy_work_minutes)
|
|
|
|
proxy_rest_duration_minutes = getattr(config, 'proxy_rest_duration_minutes', None)
|
|
if proxy_rest_duration_minutes is not None and not args.dry_run:
|
|
manager.set_config('proxy_rest_duration_minutes', proxy_rest_duration_minutes)
|
|
|
|
return {'name': sim_type, 'enforcer': enforcer, 'config': config}
|
|
|
|
if not args.download_only:
|
|
auth_setup = setup_enforcer('Auth', args.auth_env, 'auth_policy_enforcer_config', 'auth_env')
|
|
if auth_setup: enforcer_setups.append(auth_setup)
|
|
|
|
if not args.auth_only:
|
|
download_setup = setup_enforcer('Download', args.download_env, 'download_policy_enforcer_config', 'download_env')
|
|
if download_setup: enforcer_setups.append(download_setup)
|
|
|
|
if not enforcer_setups:
|
|
logger.error("No policies to enforce. Check policy file and --auth-only/--download-only flags.")
|
|
return 1
|
|
|
|
# Determine interval. Precedence: CLI -> simulation_parameters -> per-setup config -> code default.
|
|
# The CLI arg is already handled by the Config objects, so we just need to check sim_params.
|
|
sim_params_interval = sim_params.get('interval_seconds')
|
|
if args.interval_seconds is None and sim_params_interval is not None:
|
|
interval = sim_params_interval
|
|
else:
|
|
interval = min(s['config'].interval_seconds for s in enforcer_setups)
|
|
|
|
# Get cross-simulation sync configuration
|
|
cross_sync_config = policy.get('cross_simulation_sync', {})
|
|
|
|
if not args.live:
|
|
for setup in enforcer_setups:
|
|
logger.info(f"--- Applying policies for {setup['name']} Simulation ---")
|
|
setup['enforcer'].apply_policies(setup['config'])
|
|
|
|
# Apply cross-simulation sync after all policies have been applied
|
|
if cross_sync_config and len(enforcer_setups) == 2:
|
|
# We need to identify which setup is auth and which is download
|
|
# Based on their names
|
|
auth_setup = next((s for s in enforcer_setups if s['name'] == 'Auth'), None)
|
|
download_setup = next((s for s in enforcer_setups if s['name'] == 'Download'), None)
|
|
if auth_setup and download_setup:
|
|
sync_cross_simulation(
|
|
auth_setup['enforcer'].manager,
|
|
download_setup['enforcer'].manager,
|
|
cross_sync_config,
|
|
dry_run=args.dry_run
|
|
)
|
|
return 0
|
|
|
|
logger.info(f"Running in live mode. Applying policies every {interval} seconds. Press Ctrl+C to stop.")
|
|
if not args.verbose:
|
|
print("Each '.' represents a check cycle with no actions taken.", file=sys.stderr)
|
|
|
|
while not shutdown_event:
|
|
had_action_in_cycle = False
|
|
for setup in enforcer_setups:
|
|
logger.debug(f"--- Applying policies for {setup['name']} Simulation ({setup['enforcer'].manager.key_prefix}) ---")
|
|
if setup['enforcer'].apply_policies(setup['config']):
|
|
had_action_in_cycle = True
|
|
|
|
# Apply cross-simulation sync after all policies have been applied in this cycle
|
|
if cross_sync_config and len(enforcer_setups) == 2:
|
|
auth_setup = next((s for s in enforcer_setups if s['name'] == 'Auth'), None)
|
|
download_setup = next((s for s in enforcer_setups if s['name'] == 'Download'), None)
|
|
if auth_setup and download_setup:
|
|
sync_cross_simulation(
|
|
auth_setup['enforcer'].manager,
|
|
download_setup['enforcer'].manager,
|
|
cross_sync_config,
|
|
dry_run=args.dry_run
|
|
)
|
|
# Note: sync_cross_simulation may take actions, but we don't track them for the dot indicator
|
|
# This is fine for now
|
|
|
|
if had_action_in_cycle:
|
|
if not args.verbose:
|
|
# Print a newline to separate the action logs from subsequent dots
|
|
print(file=sys.stderr)
|
|
else:
|
|
if not args.verbose:
|
|
print(".", end="", file=sys.stderr)
|
|
sys.stderr.flush()
|
|
|
|
sleep_end_time = time.time() + interval
|
|
while time.time() < sleep_end_time and not shutdown_event:
|
|
time.sleep(1)
|
|
|
|
logger.info("Policy enforcer stopped.")
|
|
return 0
|