# Server related imports
from http.server import BaseHTTPRequestHandler, HTTPServer
import json
from urllib.parse import urlparse, parse_qs
from urllib.request import Request, urlopen
import ssl
import sys, traceback


LA_URLS = {
    'widevine': 'https://widevine.keyos.com/api/v4/getLicense',
    'playready': 'https://playready.keyos.com/api/v4/getLicense',
    'fairplay': 'https://fairplay.keyos.com/api/v4/getLicense',
}

CERT_URLS = {
    'widevine': 'https://widevine.keyos.com/api/v4/getCertificate',
    'fairplay': 'https://fairplay.keyos.com/api/v4/getCertificate',
}

RESPONSE_TYPES = {
    'widevine': 'application/octet-stream',
    'playready': 'text/xml; charset=utf-8',
    'fairplay': 'text/plain; charset=utf-8',
}




class LAProxy_RequestHandler(BaseHTTPRequestHandler):
    # The SHA1 or SHA256 hash of your FPS certificate, provided by KeyOS Support.
    _fps_cert_hash = '5448EC23A7492E3CB26596ACEA7B0D6B2A91206A'

    # SSL context that trusts the local authxml_generator's self-signed cert
    _ssl_ctx = ssl.create_default_context()
    _ssl_ctx.check_hostname = False
    _ssl_ctx.verify_mode = ssl.CERT_NONE

    def _fetch_auth_xml(self):
        with open("config.json", "r") as f:
            config_data = json.load(f)
            
        auth_xml_port = config_data.get("AUTHXML_PORT")
        keyos_key_file = config_data.get("KEYOS_KEY_FILE")
        keyos_key = keyos_key_file.rsplit('/', 1)[-1].split('.')[0]
        domain = config_data.get("DOMAIN", "localhost")
        authxml_url = f'https://{domain}:{auth_xml_port}/drm?getAuthXml&key={keyos_key}'
        req = Request(authxml_url)
        with urlopen(req, context=self._ssl_ctx) as resp:
            return resp.read().decode('utf-8')

    def _send_cors_headers(self):
        self.send_header('Access-Control-Allow-Origin', '*')
        self.send_header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS')
        self.send_header('Access-Control-Allow-Headers', 'Content-Type')

    def _respond(self, data, error_code, content_type='text/html; charset=utf-8'):
        if type(data) is str:
            data = bytes(data, 'UTF-8')

        self.send_response(error_code)
        self.send_header('Content-Type', content_type)
        self.send_header('Content-Length', len(data))
        self._send_cors_headers()
        self.end_headers()
        self.wfile.write(data)

    def do_OPTIONS(self):
        self.send_response(204)
        self._send_cors_headers()
        self.end_headers()

    def _forward_to_keyos(self, url, headers, payload=None, method='POST'):
        """Sends a request to KeyOS and returns (body, status_code)."""
        req = Request(url, data=payload, headers=headers, method=method)
        try:
            with urlopen(req) as resp:
                return resp.read(), resp.status
        except Exception as e:
            if hasattr(e, 'read'):
                return e.read(), e.code
            raise

    def do_GET(self):
        """Handles FairPlay certificate request.
        The playback client fetches the FPS certificate before it can generate the license challenge (SPC)."""
        try:
            path = urlparse(self.path).path
            if path != '/fps/certificate':
                return self._respond('Not found.', 404)

            cert_url = CERT_URLS['fairplay'] + '?certHash={}'.format(self._fps_cert_hash)
            headers = {
                'x-keyos-authorization': self._fetch_auth_xml(),
            }

            body, status = self._forward_to_keyos(cert_url, headers, method='GET')
            return self._respond(body, status, 'text/plain; charset=utf-8')

        except Exception as e:
            print('-' * 60)
            traceback.print_exc(file=sys.stdout)
            print('-' * 60)

            return self._respond('Sorry, there was an error. {}'.format(e), 403)

    def do_POST(self):
        try:
            content_length = int(self.headers['Content-Length'])
            payload = self.rfile.read(content_length)

            if not payload:
                return self._respond('Can\'t read the license challenge.', 403)

            # Get the type of the DRM to get the license for.
            query = urlparse(self.path).query
            query_components = parse_qs(query)

            if 'drm-type' not in query_components:
                return self._respond('\'drm-type\' query param not set.', 403)

            drm_type = query_components['drm-type'][0].lower()

            la_url = LA_URLS.get(drm_type)
            if not la_url:
                return self._respond('Can\'t determine the license acquisition URL.', 403)

            # Build per-request headers.
            headers = {'Content-Type': 'text/xml; charset=utf-8'}

            # The Widevine SERVICE_CERTIFICATE_REQUEST is a 2 byte payload: 0x08 0x04 (base64: CAQ=).
            # Forward it as a POST to the Widevine getCertificate endpoint.
            if drm_type == 'widevine' and payload == b'\x08\x04':
                headers['x-keyos-authorization'] = self._fetch_auth_xml()
                body, status = self._forward_to_keyos(CERT_URLS['widevine'], headers, payload)
                return self._respond(body, status, 'application/octet-stream')

            # Fetch authentication XML from the authxml_generator service.
            headers['x-keyos-authorization'] = self._fetch_auth_xml()

            # PlayReady license requests require a specific SOAP action header.
            if drm_type == 'playready':
                headers['soapaction'] = 'http://schemas.microsoft.com/DRM/2007/03/protocols/AcquireLicense'

            body, status = self._forward_to_keyos(la_url, headers, payload)
            return self._respond(body, status, RESPONSE_TYPES[drm_type])

        except Exception as e:
            print('-' * 60)
            traceback.print_exc(file=sys.stdout)
            print('-' * 60)

            return self._respond('Sorry, there was an error. {}'.format(e), 403)