#!/usr/bin/env python3 """ Client script to manage proxies in the YTTokenOpService. This script allows you to list, ban, unban, and reset proxies that are managed by a ytdlp-ops-server instance via Redis. """ import argparse import sys import os import logging from pathlib import Path import datetime # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) try: from thrift.transport import TTransport from pangramia.yt.exceptions.ttypes import PBServiceException, PBUserException from pangramia.yt.common.constants import ErrorCode from tabulate import tabulate from yt_ops_services.client_utils import get_thrift_client, format_timestamp except ImportError as e: print(f"Error importing required modules: {e}") print("Please ensure you have installed dependencies by running: pip install -e .") sys.exit(1) def main(): parser = argparse.ArgumentParser( description="Manage proxies for the YTDLP Operations Server.\n\n" "This script allows you to list, ban, unban, and reset proxies that are managed\n" "by a ytdlp-ops-server instance via Redis. It provides a command-line interface\n" "to interact with the proxy management features of the server.", epilog="Usage examples:\n" " # List statuses for a specific server identity\n" " python proxy_manager_client.py list --server-identity ytdlp-ops-airflow-service\n\n" " # Ban a proxy for a specific server\n" " python proxy_manager_client.py ban --server-identity ytdlp-ops-airflow-service --proxy-url socks5://proxy.example.com:1080\n\n" " # Unban a proxy\n" " python proxy_manager_client.py unban --server-identity ytdlp-ops-airflow-service --proxy-url socks5://proxy.example.com:1080\n\n" " # Reset all proxies for a server to ACTIVE\n" " python proxy_manager_client.py reset --server-identity ytdlp-ops-airflow-service", formatter_class=argparse.RawTextHelpFormatter ) parser.add_argument('--host', default=os.getenv('YTDLP_HOST', '127.0.0.1'), help="Server host (default: 127.0.0.1 or YTDLP_HOST env). Using 127.0.0.1 avoids harmless connection errors when the local Envoy proxy only listens on IPv4.") parser.add_argument('--port', type=int, default=int(os.getenv('YTDLP_PORT', '9080')), help='Server port (default: 9080 or YTDLP_PORT env)') subparsers = parser.add_subparsers(dest='command', required=True, help='Available commands') # List command list_parser = subparsers.add_parser( 'list', help='List proxy statuses for a given server identity.', description="List the status of all proxies associated with a specific server identity.\n" "The status includes:\n" "- Server: The server identity.\n" "- Proxy URL: The URL of the proxy.\n" "- Status: ACTIVE or BANNED.\n" "- Success: Count of successful uses.\n" "- Failures: Count of failed uses.\n" "- Last Success: Timestamp of the last successful use.\n" "- Last Failure: Timestamp of the last failed use.", formatter_class=argparse.RawTextHelpFormatter ) list_parser.add_argument('--server-identity', type=str, help='The identity of the server to query. If not provided, shows status for the connected server instance.') # Ban command ban_parser = subparsers.add_parser( 'ban', help='Ban a specific proxy for a server.', description="Manually set a proxy's status to BANNED for a specific server identity.\n" "A banned proxy will not be used for future requests by that server instance\n" "until it is unbanned or reset.", formatter_class=argparse.RawTextHelpFormatter ) ban_parser.add_argument('--server-identity', type=str, required=True, help='The identity of the server where the proxy should be banned.') ban_parser.add_argument('--proxy-url', type=str, required=True, help="The full URL of the proxy to ban (e.g., 'socks5://host:port').") # Unban command unban_parser = subparsers.add_parser( 'unban', help='Unban a specific proxy for a server.', description="Manually set a proxy's status to ACTIVE for a specific server identity.\n" "This will allow the server instance to use the proxy for future requests.", formatter_class=argparse.RawTextHelpFormatter ) unban_parser.add_argument('--server-identity', type=str, required=True, help='The identity of the server where the proxy should be unbanned.') unban_parser.add_argument('--proxy-url', type=str, required=True, help="The full URL of the proxy to unban (e.g., 'socks5://host:port').") # Reset command reset_parser = subparsers.add_parser( 'reset', help='Reset all proxy statuses for a server to ACTIVE.', description="Reset the status of all proxies associated with a specific server identity to ACTIVE.\n" "This is useful for clearing all bans and making all configured proxies available again.", formatter_class=argparse.RawTextHelpFormatter ) reset_parser.add_argument('--server-identity', type=str, required=True, help='The identity of the server whose proxies should be reset.') args = parser.parse_args() client, transport = None, None try: client, transport = get_thrift_client(args.host, args.port) if args.command == 'list': logger.info(f"Getting proxy statuses for server: {args.server_identity or 'local server'}") statuses = client.getProxyStatus(args.server_identity) if not statuses: print("\nThe server reported no proxy statuses.") print("This can happen if no proxies are configured, or if all configured proxies failed their initial health check on server startup.\n") else: # Determine which proxy is next in rotation for each server identity next_proxies = {s.serverIdentity: s.proxyUrl for s in statuses if '(next)' in s.status} status_list = [] for s in statuses: is_next = next_proxies.get(s.serverIdentity) == s.proxyUrl status_list.append({ "Server": s.serverIdentity, "Proxy URL": f"{s.proxyUrl} ->" if is_next else s.proxyUrl, "Status": s.status.replace(" (next)", ""), "Success": s.successCount, "Failures": s.failureCount, "Last Success": format_timestamp(s.lastSuccessTimestamp), "Last Failure": format_timestamp(s.lastFailureTimestamp), }) print("\n--- Proxy Statuses ---") print(tabulate(status_list, headers="keys", tablefmt="grid")) print("----------------------\n") elif args.command == 'ban': logger.info(f"Banning proxy '{args.proxy_url}' for server '{args.server_identity}'...") success = client.banProxy(args.proxy_url, args.server_identity) if success: print(f"Successfully banned proxy '{args.proxy_url}' for server '{args.server_identity}'.") else: print("Failed to ban proxy. Check server logs for details.") sys.exit(1) elif args.command == 'unban': logger.info(f"Unbanning proxy '{args.proxy_url}' for server '{args.server_identity}'...") success = client.unbanProxy(args.proxy_url, args.server_identity) if success: print(f"Successfully unbanned proxy '{args.proxy_url}' for server '{args.server_identity}'.") else: print("Failed to unban proxy. Check server logs for details.") sys.exit(1) elif args.command == 'reset': logger.info(f"Resetting all proxy statuses for server '{args.server_identity}'...") success = client.resetAllProxyStatuses(args.server_identity) if success: print(f"Successfully reset all proxy statuses for server '{args.server_identity}'.") else: print("Failed to reset all proxy statuses. Check server logs for details.") sys.exit(1) except (PBServiceException, PBUserException) as e: if hasattr(e, 'errorCode') and e.errorCode == ErrorCode.NOT_IMPLEMENTED: logger.error(f"Action '{args.command}' is not implemented by the server. It may be running in the wrong service mode.") print(f"Error: The server does not support the action '{args.command}'.") print("Please check that the server is running in 'all-in-one' or 'management' mode.") else: logger.error(f"Thrift error performing action '{args.command}': {e.message}", exc_info=True) print(f"Error: {e.message}") sys.exit(1) except TTransport.TTransportException as e: # The logger.error is not needed here because TSocket already logs connection errors. print(f"Error: Connection to server at {args.host}:{args.port} failed. Is the server running?") print(f"Details: {e}") sys.exit(1) except Exception as e: logger.error(f"An unexpected error occurred: {e}", exc_info=True) print(f"An unexpected error occurred: {e}") sys.exit(1) finally: if transport and transport.isOpen(): transport.close() logger.info("Thrift connection closed.") if __name__ == "__main__": main()