2025-04-01 14:17:16 +02:00

235 lines
8.6 KiB
Python

from pathlib import Path
from typing import Any, Optional, Union
from aiohttp.typedefs import Handler
from aiohttp import web
from pyplayready import __version__, PSSH
from pyplayready.cdm import Cdm
from pyplayready.device import Device
from pyplayready.exceptions import (InvalidSession, TooManySessions, InvalidLicense, InvalidPssh)
routes = web.RouteTableDef()
async def _startup(app: web.Application) -> None:
app["cdms"] = {}
app["config"]["devices"] = {
path.stem: path
for x in app["config"]["devices"]
for path in [Path(x)]
}
for device in app["config"]["devices"].values():
if not device.is_file():
raise FileNotFoundError(f"Device file does not exist: {device}")
async def _cleanup(app: web.Application) -> None:
app["cdms"].clear()
del app["cdms"]
app["config"].clear()
del app["config"]
@routes.get("/")
async def ping(_: Any) -> web.Response:
return web.json_response({"message": "OK"})
@routes.get("/{device}/open")
async def open_(request: web.Request) -> web.Response:
secret_key = request.headers["X-Secret-Key"]
device_name = request.match_info["device"]
user = request.app["config"]["users"][secret_key]
if device_name not in user["devices"] or device_name not in request.app["config"]["devices"]:
# we don't want to be verbose with the error as to not reveal device names
# by trial and error to users that are not authorized to use them
return web.json_response({"message": f"Device '{device_name}' is not found or you are not authorized to use it."}, status=403)
cdm: Optional[Cdm] = request.app["cdms"].get((secret_key, device_name))
if not cdm:
device = Device.load(request.app["config"]["devices"][device_name])
cdm = request.app["cdms"][(secret_key, device_name)] = Cdm.from_device(device)
try:
session_id = cdm.open()
except TooManySessions as e:
return web.json_response({"message": str(e)}, status=400)
return web.json_response({
"message": "Success",
"data": {
"session_id": session_id.hex(),
"device": {
"security_level": cdm.security_level
}
}
})
@routes.get("/{device}/close/{session_id}")
async def close(request: web.Request) -> web.Response:
secret_key = request.headers["X-Secret-Key"]
device_name = request.match_info["device"]
session_id = bytes.fromhex(request.match_info["session_id"])
cdm: Optional[Cdm] = request.app["cdms"].get((secret_key, device_name))
if not cdm:
return web.json_response({"message": f"No Cdm session for {device_name} has been opened yet. No session to close."}, status=400)
try:
cdm.close(session_id)
except InvalidSession:
return web.json_response({"message": f"Invalid Session ID '{session_id.hex()}', it may have expired."}, status=400)
return web.json_response({"message": f"Successfully closed Session '{session_id.hex()}'."})
@routes.post("/{device}/get_license_challenge")
async def get_license_challenge(request: web.Request) -> web.Response:
secret_key = request.headers["X-Secret-Key"]
device_name = request.match_info["device"]
body = await request.json()
for required_field in ("init_data", "session_id"):
if not body.get(required_field):
return web.json_response({"message": f"Missing required field '{required_field}' in JSON body."}, status=400)
cdm: Optional[Cdm] = request.app["cdms"].get((secret_key, device_name))
if not cdm:
return web.json_response({"message": f"No Cdm session for {device_name} has been opened yet. No session to use."}, status=400)
session_id = bytes.fromhex(body["session_id"])
init_data = body["init_data"]
if not init_data.startswith("<WRMHEADER"):
try:
pssh = PSSH(init_data)
if pssh.wrm_headers:
init_data = pssh.wrm_headers[0]
except InvalidPssh as e:
return web.json_response({"message": f"Unable to parse base64 PSSH, {e}"}, status=500)
try:
license_request = cdm.get_license_challenge(
session_id=session_id,
wrm_header=init_data,
)
except InvalidSession:
return web.json_response({"message": f"Invalid Session ID '{session_id.hex()}', it may have expired."}, status=400)
except Exception as e:
return web.json_response({"message": f"Error, {e}"}, status=500)
return web.json_response({
"message": "Success",
"data": {
"challenge": license_request
}
}, status=200)
@routes.post("/{device}/parse_license")
async def parse_license(request: web.Request) -> web.Response:
secret_key = request.headers["X-Secret-Key"]
device_name = request.match_info["device"]
body = await request.json()
for required_field in ("license_message", "session_id"):
if not body.get(required_field):
return web.json_response({"message": f"Missing required field '{required_field}' in JSON body."}, status=400)
cdm: Optional[Cdm] = request.app["cdms"].get((secret_key, device_name))
if not cdm:
return web.json_response({"message": f"No Cdm session for {device_name} has been opened yet. No session to use."}, status=400)
session_id = bytes.fromhex(body["session_id"])
license_message = body["license_message"]
try:
cdm.parse_license(session_id, license_message)
except InvalidSession:
return web.json_response({"message": f"Invalid Session ID '{session_id.hex()}', it may have expired."}, status=400)
except InvalidLicense as e:
return web.json_response({"message": f"Invalid License, {e}"}, status=400)
except Exception as e:
return web.json_response({"message": f"Error, {e}"}, status=500)
return web.json_response({"message": "Successfully parsed and loaded the Keys from the License message."})
@routes.post("/{device}/get_keys")
async def get_keys(request: web.Request) -> web.Response:
secret_key = request.headers["X-Secret-Key"]
device_name = request.match_info["device"]
body = await request.json()
for required_field in ("session_id",):
if not body.get(required_field):
return web.json_response({"message": f"Missing required field '{required_field}' in JSON body."}, status=400)
session_id = bytes.fromhex(body["session_id"])
cdm = request.app["cdms"].get((secret_key, device_name))
if not cdm:
return web.json_response({"message": f"No Cdm session for {device_name} has been opened yet. No session to use."}, status=400)
try:
keys = cdm.get_keys(session_id)
except InvalidSession:
return web.json_response({"message": f"Invalid Session ID '{session_id.hex()}', it may have expired."}, status=400)
except Exception as e:
return web.json_response({"message": f"Error, {e}"}, status=500)
keys_json = [
{
"key_id": key.key_id.hex,
"key": key.key.hex(),
"type": key.key_type.value,
"cipher_type": key.cipher_type.value,
"key_length": key.key_length,
}
for key in keys
]
return web.json_response({
"message": "Success",
"data": {
"keys": keys_json
}
})
@web.middleware
async def authentication(request: web.Request, handler: Handler) -> web.Response:
secret_key = request.headers.get("X-Secret-Key")
if request.path != "/" and not secret_key:
request.app.logger.debug(f"{request.remote} did not provide authorization.")
response = web.json_response({"message": "Secret Key is Empty."}, status=401)
elif request.path != "/" and secret_key not in request.app["config"]["users"]:
request.app.logger.debug(f"{request.remote} failed authentication with '{secret_key}'.")
response = web.json_response({"message": "Secret Key is Invalid, the Key is case-sensitive."}, status=401)
else:
try:
response = await handler(request)
except web.HTTPException as e:
request.app.logger.error(f"An unexpected error has occurred, {e}")
response = web.json_response({"message": e.reason}, status=500)
response.headers.update({
"Server": f"https://github.com/ready-dl/pyplayready serve v{__version__}"
})
return response
def run(config: dict, host: Optional[Union[str, web.HostSequence]] = None, port: Optional[int] = None) -> None:
app = web.Application(middlewares=[authentication])
app.on_startup.append(_startup)
app.on_cleanup.append(_cleanup)
app.add_routes(routes)
app["config"] = config
web.run_app(app, host=host, port=port)