Update main.py

main.py

@@ -1,4 +1,5 @@
 from fastapi import FastAPI
+from fastapi import HTTPException
 from contextlib import asynccontextmanager
 from app.enka_logic import EnkaApp
 from app.akasa_logic import AkasaApp
@@ -16,6 +17,11 @@ yt_info = YoutubeInfo()
 yt_dl = YoutubeDownload()
 keygen = KeyGenApp()
 
+BOT_API_KEYS = {
+    "fdd8ca0526ca46c9ba3a40101e530810",
+    "botkey456"
+}
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     await enka.start()
@@ -59,10 +65,26 @@ async def ytdl(url: str):
     return await yt_dl.get_download(url)
 
 @app.get("/keys")
-def get_keys(count: int = 10):
+def get_keys(
+    count: int = 10,
+    url: str | None = None,
+    apikey: str | None = None
+):
+    if apikey not in BOT_API_KEYS:
+        raise HTTPException(status_code=403, detail="Invalid API key")
+
+    if not url:
+        raise HTTPException(status_code=400, detail="Missing url")
+
+    if count < 1:
+        count = 1
     if count > 15:
         count = 15
-    return keygen.fetch_keys(count)
+
+    return keygen.fetch_keys(
+        target_url=url,
+        count=count
+    )
 
 @app.get("/")
 def home():