193 lines
9.7 KiB
Python
193 lines
9.7 KiB
Python
#!/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()
|