From 4bc0edcca98a67cdb0c4caaf5a3753874fb68618 Mon Sep 17 00:00:00 2001
From: rlaphoenix <rlaphoenix@pm.me>
Date: Tue, 2 Aug 2022 00:18:08 +0100
Subject: [PATCH] serve: Set `Server` response header with pywidevine version

This allows clients to test with a HEAD request to / to see what version the API is running and test if it's actually a pywidevine serve API.
---
 pywidevine/serve.py | 54 ++++++++++++++++++++++++++-------------------
 1 file changed, 31 insertions(+), 23 deletions(-)

diff --git a/pywidevine/serve.py b/pywidevine/serve.py
index bf2bf65..2f775d6 100644
--- a/pywidevine/serve.py
+++ b/pywidevine/serve.py
@@ -3,8 +3,6 @@ import sys
 from pathlib import Path
 from typing import Optional, Union
 
-from pywidevine.exceptions import TooManySessions, InvalidSession
-
 try:
     from aiohttp import web
 except ImportError:
@@ -15,8 +13,10 @@ except ImportError:
     )
     sys.exit(1)
 
+from pywidevine import __version__
 from pywidevine.cdm import Cdm
 from pywidevine.device import Device
+from pywidevine.exceptions import TooManySessions, InvalidSession
 from pywidevine.license_protocol_pb2 import LicenseType, License
 
 routes = web.RouteTableDef()
@@ -256,29 +256,37 @@ async def keys(request: web.Request) -> web.Response:
 
 @web.middleware
 async def authentication(request: web.Request, handler) -> web.Response:
-    secret_key = request.headers.get("X-Secret-Key")
-    if not secret_key:
-        request.app.logger.debug(f"{request.remote} did not provide authorization.")
-        return web.json_response({
-            "status": "401",
-            "message": "Secret Key is Empty."
-        }, status=401)
+    response = None
+    if request.path != "/":
+        secret_key = request.headers.get("X-Secret-Key")
+        if not secret_key:
+            request.app.logger.debug(f"{request.remote} did not provide authorization.")
+            response = web.json_response({
+                "status": "401",
+                "message": "Secret Key is Empty."
+            }, status=401)
+        elif secret_key not in request.app["config"]["users"]:
+            request.app.logger.debug(f"{request.remote} failed authentication with '{secret_key}'.")
+            response = web.json_response({
+                "status": "401",
+                "message": "Secret Key is Invalid, the Key is case-sensitive."
+            }, status=401)
 
-    if secret_key not in request.app["config"]["users"]:
-        request.app.logger.debug(f"{request.remote} failed authentication with '{secret_key}'.")
-        return web.json_response({
-            "status": "401",
-            "message": "Secret Key is Invalid, the Key is case-sensitive."
-        }, status=401)
+    if response is None:
+        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({
+                "status": 500,
+                "message": e.reason
+            }, status=500)
 
-    try:
-        return await handler(request)
-    except web.HTTPException as e:
-        request.app.logger.error(f"An unexpected error has occurred, {e}")
-        return web.json_response({
-            "status": 500,
-            "message": e.reason
-        }, status=500)
+    response.headers.update({
+        "Server": f"https://github.com/rlaphoenix/pywidevine serve v{__version__}"
+    })
+
+    return response
 
 
 def run(config: dict, host: Optional[Union[str, web.HostSequence]] = None, port: Optional[int] = None):