KeyDive/docs/server/server.py

import argparse
import json
import time
import logging

import requests

from pathlib import Path
from flask import Flask, Response, request, redirect

from keydive.__main__ import configure_logging

# Suppress urllib3 warnings
logging.getLogger('urllib3.connectionpool').setLevel(logging.ERROR)

# Initialize Flask application
app = Flask(__name__)

# Define paths, constants, and global flags
PARENT = Path(__file__).parent
VERSION = '1.0.0'
KEYBOX = False
DELAY = 10


@app.route('/', methods=['GET'])
def health_check() -> Response:
    """
    Health check endpoint to confirm the server is running.

    Returns:
        Response: A simple "pong" message with a 200 OK status.
    """
    return Response(response='pong', status=200, content_type='text/html; charset=utf-8')


@app.route('/shaka-demo-assets/angel-one-widevine/<path:file>', methods=['GET'])
def shaka_demo_assets(file) -> Response:
    """
    Serves cached assets for Widevine demo content. If the requested file is
    not available locally, it fetches it from a remote server and caches it.

    Args:
        file (str): File path requested by the client.

    Returns:
        Response: File content as a byte stream, or a 404 error if not found.
    """
    logger = logging.getLogger('Shaka')
    logger.info('%s %s', request.method, request.path)

    try:
        path = PARENT / '.assets' / file
        path.parent.mkdir(parents=True, exist_ok=True)

        if path.is_file():
            # Serve cached file content if available
            content = path.read_bytes()
        else:
            # Fetch the file from remote storage if not cached locally
            r = requests.get(
                url=f'https://storage.googleapis.com/shaka-demo-assets/angel-one-widevine/{file}',
                headers={
                    'Accept': '*/*',
                    'User-Agent': 'KalturaDeviceInfo/1.4.1 (Linux;Android 10) ExoPlayerLib/2.9.3'
                }
            )
            r.raise_for_status()
            path.write_bytes(r.content)  # Cache the downloaded content
            content = r.content
            logger.debug('Downloaded assets: %s', path)

        return Response(response=content, status=200, content_type='application/octet-stream')
    except Exception as e:
        return Response(response=str(e), status=404, content_type='text/html; charset=utf-8')


@app.route('/certificateprovisioning/v1/devicecertificates/create', methods=['POST'])
def certificate_provisioning() -> Response:
    """
    Handles device certificate provisioning requests by intercepting the request,
    saving it as a curl command, and then responding based on cached data
    or redirecting if no cached response is available.

    Returns:
        Response: JSON response if provisioning is complete, else a redirection.
    """
    global KEYBOX, DELAY
    logger = logging.getLogger('Google')
    logger.info('%s %s', request.method, request.path)

    if KEYBOX:
        logger.warning('Provisioning request aborted to prevent keybox spam')
        return Response(response='Internal Server Error', status=500, content_type='text/html; charset=utf-8')

    # Generate a curl command from the incoming request for debugging or testing
    user_agent = request.headers.get('User-Agent', 'Unknown')
    url = request.url.replace('http://', 'https://')
    prompt = [
        "curl --request POST",
        f"--url '{url}'",
        "--compressed",
        "--header 'accept-encoding: gzip'",
        "--header 'connection: Keep-Alive'",
        "--header 'content-type: application/x-www-form-urlencoded'",
        "--header 'host: www.googleapis.com'",
        f"--header 'user-agent: {user_agent}'"
    ]

    # Save the curl command for potential replay or inspection
    curl = PARENT / 'curl.txt'
    curl.write_text(' \\\n  '.join(prompt))
    logger.debug('Saved curl command to: %s', curl)

    # Wait for provisioning response data with retries
    logger.warning('Waiting for provisioning response...')
    provision = PARENT / 'provisioning.json'
    provision.unlink(missing_ok=True)
    provision.write_bytes(b'')  # Create empty file for manual input if needed

    # Poll for the presence of a response up to DELAY times with 1-second intervals
    for _ in range(DELAY):
        try:
            content = json.loads(provision.read_bytes())
            if content:
                # Cleanup after successful response
                curl.unlink(missing_ok=True)
                provision.unlink(missing_ok=True)
                return Response(response=content, status=200, content_type='application/json')
        except Exception as e:
            pass  # Continue waiting if file is empty or not yet ready
        time.sleep(1)

    # Redirect to the secure URL if response is not available
    logger.warning('Redirecting to avoid timeout')
    return redirect(url, code=302)


def main() -> None:
    """
    Main entry point for the application. Parses command-line arguments
    to set global parameters and configures logging, then starts the Flask server.
    """
    global VERSION, DELAY, KEYBOX
    parser = argparse.ArgumentParser(description='Local DRM provisioning video player.')

    # Global arguments for the application
    group_global = parser.add_argument_group('Global')
    group_global.add_argument('--host', required=False, type=str, default='127.0.0.1', metavar='<host>', help='Host address for the server to bind to.')
    group_global.add_argument('--port', required=False, type=int, default=9090, metavar='<port>', help='Port number for the server to listen on.')
    group_global.add_argument('-v', '--verbose', required=False, action='store_true', help='Enable verbose logging for detailed debug output.')
    group_global.add_argument('-l', '--log', required=False, type=Path, metavar='<dir>', help='Directory to store log files.')
    group_global.add_argument('--version', required=False, action='store_true', help='Display Server version information.')

    # Advanced options
    group_advanced = parser.add_argument_group('Advanced')
    group_advanced.add_argument('-d', '--delay', required=False, type=int, metavar='<delay>', default=10, help='Delay (in seconds) between successive checks for provisioning responses.')
    group_advanced.add_argument('-k', '--keybox', required=False, action='store_true', help='Enable keybox mode, which aborts provisioning requests to prevent spam.')

    args = parser.parse_args()

    if args.version:
        print(f'Server {VERSION}')
        exit(0)

    # Configure logging
    log_path = configure_logging(path=args.log, verbose=args.verbose)
    logger = logging.getLogger('Server')
    logger.info('Version: %s', VERSION)

    try:
        # Set global variables based on parsed arguments
        DELAY = args.delay
        KEYBOX = args.keybox

        # Start Flask app with specified host, port, and debug mode
        logging.getLogger('werkzeug').setLevel(logging.INFO if args.verbose else logging.ERROR)
        app.run(host=args.host, port=args.port, debug=False)
    except KeyboardInterrupt:
        pass
    except Exception as e:
        logger.critical(e, exc_info=args.verbose)

    # Final logging and exit
    if log_path:
        logger.info('Log file: %s' % log_path)
    logger.info('Exiting')


if __name__ == '__main__':
    main()