yt-dlp-dags/get_info_json_client.py

#!/usr/bin/env python3
"""
Client script to get info.json from the Thrift service.

Usage:
  python get_info_json_client.py [URL] --host [HOST] --port [PORT] [options]

Options:
  --host HOST           Thrift server host
  --port PORT           Thrift server port
  --account-id ID       Account ID to use
  --output FILE         Output file path
  --verbose             Enable verbose output
"""

import argparse
import json
import os
import sys
import logging
from typing import Dict, Any, Optional

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('info_json_client')

# Import Thrift modules
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from thrift.transport import TTransport
from pangramia.yt.common.ttypes import TokenUpdateMode
from pangramia.yt.exceptions.ttypes import PBServiceException, PBUserException
from yt_ops_services.client_utils import get_thrift_client

def parse_args():
    """Parse command line arguments"""
    parser = argparse.ArgumentParser(description='Get info.json from Thrift service')
    parser.add_argument('url', help='YouTube URL or video ID')
    parser.add_argument('--host', default='127.0.0.1', help="Thrift server host. 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=9080, help='Thrift server port')
    parser.add_argument('--profile', default='default_profile', help='The profile name (accountId) to use for the request.')
    parser.add_argument('--client', help='Specific client to use (e.g., web, ios, android). Overrides server default.')
    parser.add_argument('--output', help='Output file path for the info.json. If not provided, prints to stdout.')
    parser.add_argument('--machine-id', help='Identifier for the client machine. Defaults to hostname.')
    parser.add_argument('--verbose', action='store_true', help='Enable verbose output')
    return parser.parse_args()

def main():
    """Main entry point"""
    args = parse_args()
    
    # Set log level
    if args.verbose:
        logger.setLevel(logging.DEBUG)
    
    transport = None
    try:
        # Create Thrift client
        client, transport = get_thrift_client(args.host, args.port)
        
        # Get token data, which includes the info.json
        logger.info(f"Requesting info.json for URL '{args.url}' using profile '{args.profile}'")
        
        # Prepare arguments for the Thrift call
        machine_id = args.machine_id
        if not machine_id:
            import socket
            machine_id = socket.gethostname()
            logger.info(f"No machine ID provided, using hostname: {machine_id}")

        thrift_args = {
            'accountId': args.profile,
            'updateType': TokenUpdateMode.AUTO,
            'url': args.url,
            'clients': args.client,
            'machineId': machine_id
        }
        if args.client:
            logger.info(f"Requesting to use specific client: {args.client}")
        else:
            logger.info("No specific client requested, server will use its default.")

        token_data = client.getOrRefreshToken(**thrift_args)
        
        if not token_data or not hasattr(token_data, 'infoJson') or not token_data.infoJson:
            logger.error("Server did not return valid info.json data.")
            print("Error: Server did not return valid info.json data.", file=sys.stderr)
            return 1
            
        info_json_str = token_data.infoJson
        
        # Check if the returned info.json is an error report
        try:
            info_data = json.loads(info_json_str)
            if isinstance(info_data, dict) and 'error' in info_data:
                error_code = info_data.get('errorCode', 'N/A')
                error_message = info_data.get('message', info_data.get('error', 'Unknown error'))
                logger.error(f"Server returned an error in info.json (Code: {error_code}): {error_message}")
                print(f"Error from server (Code: {error_code}): {error_message}", file=sys.stderr)
                # Optionally print the full error JSON
                if args.verbose:
                    print(json.dumps(info_data, indent=2), file=sys.stderr)
                return 1
        except json.JSONDecodeError:
            logger.error(f"Failed to parse info.json from server: {info_json_str[:200]}...")
            print("Error: Failed to parse the info.json response from the server.", file=sys.stderr)
            return 1

        logger.info(f"Successfully retrieved info.json ({len(info_json_str)} bytes)")

        # Write to output file if specified, otherwise print to stdout
        if args.output:
            try:
                with open(args.output, 'w', encoding='utf-8') as f:
                    # Pretty-print the JSON to the file
                    json.dump(info_data, f, indent=2)
                logger.info(f"Wrote info.json to {args.output}")
                print(f"Successfully saved info.json to {args.output}")
            except IOError as e:
                logger.error(f"Failed to write to output file {args.output}: {e}")
                print(f"Error: Failed to write to output file {args.output}: {e}", file=sys.stderr)
                return 1
        else:
            # Pretty-print the JSON to stdout
            print(json.dumps(info_data, indent=2))
        
        return 0
    except (PBServiceException, PBUserException) as e:
        logger.error(f"A Thrift error occurred: {e.message}", exc_info=args.verbose)
        print(f"Error: {e.message}", file=sys.stderr)
        if hasattr(e, 'context') and e.context:
            print(f"Context: {e.context}", file=sys.stderr)
        return 1
    except TTransport.TTransportException as e:
        logger.error(f"Connection to server failed: {e}", exc_info=args.verbose)
        print(f"Error: Connection to server at {args.host}:{args.port} failed.", file=sys.stderr)
        return 1
    except Exception as e:
        logger.exception(f"An unexpected error occurred: {e}")
        print(f"An unexpected error occurred: {e}", file=sys.stderr)
        return 1
    finally:
        if transport and transport.isOpen():
            transport.close()
            logger.info("Thrift connection closed.")

if __name__ == "__main__":
    sys.exit(main())