From 3b8b2d8bddd9c2fce6416ccd62fba685fd1c5f2d Mon Sep 17 00:00:00 2001 From: boocmp Date: Mon, 27 May 2024 20:44:10 +0700 Subject: [PATCH] Added Brave service key v2 check. --- bentofile.yaml | 2 + env/python/requirements.txt | 3 + src/requirements.txt | 5 +- src/service.py | 2 + src/stt_api.py | 33 ++++++++-- src/utils/config/config.py | 7 ++ src/utils/service_key/brave_service_key.py | 77 ++++++++++++++++++++++ 7 files changed, 121 insertions(+), 8 deletions(-) create mode 100644 src/utils/config/config.py create mode 100644 src/utils/service_key/brave_service_key.py diff --git a/bentofile.yaml b/bentofile.yaml index ee2597f..21dc8cd 100644 --- a/bentofile.yaml +++ b/bentofile.yaml @@ -7,6 +7,8 @@ include: - "utils/google_streaming/google_streaming_api_pb2.py" - "utils/npipe/__init__.py" - "utils/npipe/_posix.py" + - "utils/config/config.py" + - "utils/service_key/brave_service_key.py" - "configuration.yaml" python: requirements_txt: "requirements.txt" diff --git a/env/python/requirements.txt b/env/python/requirements.txt index c6f79d4..f1c3e18 100644 --- a/env/python/requirements.txt +++ b/env/python/requirements.txt @@ -3,3 +3,6 @@ faster_whisper fastapi aiofiles asyncio +pydantic +pydantic-settings +six diff --git a/src/requirements.txt b/src/requirements.txt index 9071a74..f1c3e18 100644 --- a/src/requirements.txt +++ b/src/requirements.txt @@ -2,4 +2,7 @@ ctranslate2 faster_whisper fastapi aiofiles -asyncio \ No newline at end of file +asyncio +pydantic +pydantic-settings +six diff --git a/src/service.py b/src/service.py index 59c0f77..70e35f3 100644 --- a/src/service.py +++ b/src/service.py @@ -1,6 +1,8 @@ import io + import bentoml from bentoml.io import JSON, File + from stt_api import app, runner_audio_transcriber diff --git a/src/stt_api.py b/src/stt_api.py index 961a6c8..9209bdf 100644 --- a/src/stt_api.py +++ b/src/stt_api.py @@ -1,12 +1,17 @@ -from fastapi import FastAPI, Request +import json +import io + +import bentoml +from runners.audio_transcriber import AudioTranscriber + +from fastapi import FastAPI, Request, Depends from fastapi.responses import StreamingResponse, JSONResponse from fastapi.encoders import jsonable_encoder + import utils.google_streaming.google_streaming_api_pb2 as speech -import bentoml from utils.npipe import AsyncChannelWriter, AsyncChannelReader -import io -from runners.audio_transcriber import AudioTranscriber -import json +from utils.service_key.brave_service_key import check_stt_request + runner_audio_transcriber = bentoml.Runner( AudioTranscriber, @@ -31,7 +36,14 @@ def to_bytes(self): app = FastAPI() @app.post("/up") -async def handleUpstream(pair: str, request: Request): +async def handleUpstream( + pair: str, + request: Request, + is_valid_brave_key = Depends(check_stt_request) +): + if not is_valid_brave_key: + return JSONResponse(content = jsonable_encoder({ "status" : "Invalid Brave Service Key" })) + try: mic_data = bytes() async with await AsyncChannelWriter.open(pair) as pipe: @@ -48,7 +60,14 @@ async def handleUpstream(pair: str, request: Request): return JSONResponse(content = jsonable_encoder({ "status" : "ok" })) @app.get("/down") -async def handleDownstream(pair: str, request: Request, output: str = "pb"): +async def handleDownstream( + pair: str, + output: str = "pb", + is_valid_brave_key = Depends(check_stt_request) +): + if not is_valid_brave_key: + return JSONResponse(content = jsonable_encoder({ "status" : "Invalid Brave Service Key" })) + async def handleStream(pair): try: async with await AsyncChannelReader.open(pair) as pipe: diff --git a/src/utils/config/config.py b/src/utils/config/config.py new file mode 100644 index 0000000..0db3e59 --- /dev/null +++ b/src/utils/config/config.py @@ -0,0 +1,7 @@ +from pydantic import Field +from pydantic_settings import BaseSettings + +class AppSettings(BaseSettings): + master_services_key_seed: str = Field('dummy') + +app_settings = AppSettings() \ No newline at end of file diff --git a/src/utils/service_key/brave_service_key.py b/src/utils/service_key/brave_service_key.py new file mode 100644 index 0000000..13f7a3d --- /dev/null +++ b/src/utils/service_key/brave_service_key.py @@ -0,0 +1,77 @@ +import base64 +import hmac +import re +from binascii import hexlify +from typing import Optional +from hashlib import sha256 +from fastapi import Header, Query +from ..config.config import app_settings + +# HKDF-SHA256 with L fixed to 32 +def hkdf_sha256_l_32(ikm, info, salt): + # HKDF-Extract(salt, IKM) -> PRK + # + # PRK = HMAC-Hash(salt, IKM) + prk = hmac.new(salt, ikm, sha256).digest() + + # HKDF-Expand(PRK, info, L) -> OKM + # + # N = ceil(L/HashLen) + # T = T(1) | T(2) | T(3) | ... | T(N) + # T(0) = empty string (zero length) + # T(1) = HMAC-Hash(PRK, T(0) | info | 0x01) + # T(2) = HMAC-Hash(PRK, T(1) | info | 0x02) + # ... + # OKM = first L octets of T + # + # L = 32 ( HashLen ) + # N = 1 + # OKM = T = T(1) = HMAC-Hash(PRK, T(0) | info | 0x01) + return hmac.new(prk, b"" + info + bytearray([1]), sha256).digest() + + +def derive_service_key(master_services_key_seed, key_id, service="stt"): + salt = sha256(service.encode("utf-8")).digest() + return hexlify( + hkdf_sha256_l_32( + master_services_key_seed.encode("utf-8"), key_id.encode("utf-8"), salt + ) + ) + +def parse_authorization_header( + header: str, +) -> (Optional[str], Optional[str], Optional[str], Optional[str]): + # Parses header values that look like: + pattern = ( + r'Signature keyId="(.+?)",algorithm="(.+?)",headers="(.+?)",signature="(.+?)"' + ) + result = re.search(pattern, header) + if result: + return result.group(1), result.group(2), result.group(3), result.group(4) + return None, None, None, None + +def check_stt_request( + pair: str = Query(), + authorization: Optional[str] = Header(None), + request_key: Optional[str] = Header(None) +): + if not authorization or not request_key or request_key != pair: + return False + + # Parse the keyId, algorithm, signature from the header + key_id, algorithm, headers, signature_b64 = parse_authorization_header(authorization) + if not key_id or not algorithm or not headers or not signature_b64 or algorithm != "hs2019" or headers != "request-key": + return False + + # Derive the service key, and expected signature and verify + service_key = derive_service_key(app_settings.master_services_key_seed, key_id) + expected_signing_string = f"request-key: {request_key}" + expected_signature = hmac.new(service_key, expected_signing_string.encode("utf-8"), sha256).digest() + expected_signature_b64 = base64.b64encode(expected_signature).decode("utf-8") + + if not hmac.compare_digest(expected_signature_b64, signature_b64): + return False + + return True + + \ No newline at end of file