From d5b76327db2eb0c40d90c97b86573260b145b1d0 Mon Sep 17 00:00:00 2001 From: Spoked Date: Fri, 28 Jun 2024 20:06:28 -0500 Subject: [PATCH] fix: replaced os with pydantic for env getters --- comet/api/core.py | 6 +++--- comet/api/stream.py | 6 +++--- comet/utils/db.py | 2 +- comet/utils/general.py | 24 ++++++++++++++---------- comet/utils/models.py | 2 +- 5 files changed, 22 insertions(+), 18 deletions(-) diff --git a/comet/api/core.py b/comet/api/core.py index 9ebc31d..d5d4863 100644 --- a/comet/api/core.py +++ b/comet/api/core.py @@ -1,9 +1,9 @@ -import os - from fastapi import APIRouter, Request from fastapi.responses import RedirectResponse from fastapi.templating import Jinja2Templates +from comet.utils.models import settings + templates = Jinja2Templates("comet/templates") main = APIRouter() @@ -21,7 +21,7 @@ async def health(): @main.get("/configure") @main.get("/{b64config}/configure") async def configure(request: Request): - return templates.TemplateResponse("index.html", {"request": request, "CUSTOM_HEADER_HTML": os.getenv("CUSTOM_HEADER_HTML", "")}) + return templates.TemplateResponse("index.html", {"request": request, "CUSTOM_HEADER_HTML": settings.CUSTOM_HEADER_HTML}) @main.get("/manifest.json") diff --git a/comet/api/stream.py b/comet/api/stream.py index d89130f..ddf458f 100644 --- a/comet/api/stream.py +++ b/comet/api/stream.py @@ -14,7 +14,7 @@ from comet.utils.general import (bytesToSize, configChecking, generateDownloadLink, getIndexerManager, getTorrentHash, isVideo, translate) from comet.utils.logger import logger -from comet.utils.models import database, rtn +from comet.utils.models import database, rtn, settings streams = APIRouter() @@ -70,7 +70,7 @@ async def stream(request: Request, b64config: str, type: str, id: str): logger.info(f"Cache found for {name}") timestamp = await database.fetch_one(f"SELECT timestamp FROM cache WHERE cacheKey = '{cacheKey}'") - if timestamp[0] + int(os.getenv("CACHE_TTL", 86400)) < time.time(): + if timestamp[0] + settings.CACHE_TTL < time.time(): await database.execute(f"DELETE FROM cache WHERE cacheKey = '{cacheKey}'") logger.info(f"Cache expired for {name}") @@ -90,7 +90,7 @@ async def stream(request: Request, b64config: str, type: str, id: str): else: logger.info(f"No cache found for {name} with user configuration") - indexerManagerType = os.getenv("INDEXER_MANAGER_TYPE", "jackett") + indexerManagerType = settings.INDEXER_MANAGER_TYPE logger.info(f"Start of {indexerManagerType} search for {name} with indexers {config['indexers']}") diff --git a/comet/utils/db.py b/comet/utils/db.py index c6c0b86..f684e6b 100644 --- a/comet/utils/db.py +++ b/comet/utils/db.py @@ -36,7 +36,7 @@ def write_config(): if isinstance(indexers, str): indexers = indexers.split(",") elif not isinstance(indexers, list): - logger.warning("Invalid indexers") + logger.error(f"Invalid indexers: {indexers}") config_data = { "indexers": indexers, diff --git a/comet/utils/general.py b/comet/utils/general.py index 6b63c67..9f3336e 100644 --- a/comet/utils/general.py +++ b/comet/utils/general.py @@ -10,6 +10,7 @@ import bencodepy from RTN.patterns import language_code_mapping from comet.utils.logger import logger +from comet.utils.models import settings translationTable = { 'ā': 'a', 'ă': 'a', 'ą': 'a', 'ć': 'c', 'č': 'c', 'ç': 'c', @@ -75,19 +76,19 @@ def configChecking(b64config: str): async def getIndexerManager(session: aiohttp.ClientSession, indexerManagerType: str, indexers: list, query: str): try: - timeout = aiohttp.ClientTimeout(total=int(os.getenv("INDEXER_MANAGER_TIMEOUT", 30))) + timeout = aiohttp.ClientTimeout(total=settings.INDEXER_MANAGER_TIMEOUT) results = [] if indexerManagerType == "jackett": - response = await session.get(f"{os.getenv('INDEXER_MANAGER_URL', 'http://127.0.0.1:9117')}/api/v2.0/indexers/all/results?apikey={os.getenv('INDEXER_MANAGER_API_KEY')}&Query={query}&Tracker[]={'&Tracker[]='.join(indexer for indexer in indexers)}", timeout=timeout) + response = await session.get(f"{settings.INDEXER_MANAGER_URL}/api/v2.0/indexers/all/results?apikey={settings.INDEXER_MANAGER_API_KEY}&Query={query}&Tracker[]={'&Tracker[]='.join(indexer for indexer in indexers)}", timeout=timeout) response = await response.json() for result in response["Results"]: results.append(result) if indexerManagerType == "prowlarr": - getIndexers = await session.get(f"{os.getenv('INDEXER_MANAGER_URL', 'http://127.0.0.1:9696')}/api/v1/indexer", headers={ - "X-Api-Key": os.getenv("INDEXER_MANAGER_API_KEY") + getIndexers = await session.get(f"{settings.INDEXER_MANAGER_URL}/api/v1/indexer", headers={ + "X-Api-Key": settings.INDEXER_MANAGER_API_KEY }) getIndexers = await getIndexers.json() @@ -96,8 +97,8 @@ async def getIndexerManager(session: aiohttp.ClientSession, indexerManagerType: if indexer["definitionName"] in indexers: indexersId.append(indexer["id"]) - response = await session.get(f"{os.getenv('INDEXER_MANAGER_URL', 'http://127.0.0.1:9696')}/api/v1/search?query={query}&indexerIds={'&indexerIds='.join(str(indexerId) for indexerId in indexersId)}&type=search", headers={ - "X-Api-Key": os.getenv("INDEXER_MANAGER_API_KEY") + response = await session.get(f"{settings.INDEXER_MANAGER_URL}/api/v1/search?query={query}&indexerIds={'&indexerIds='.join(str(indexerId) for indexerId in indexersId)}&type=search", headers={ + "X-Api-Key": settings.INDEXER_MANAGER_API_KEY }) response = await response.json() @@ -119,7 +120,7 @@ async def getTorrentHash(session: aiohttp.ClientSession, indexerManagerType: str url = torrent["Link"] if indexerManagerType == "jackett" else torrent["downloadUrl"] try: - timeout = aiohttp.ClientTimeout(total=int(os.getenv("GET_TORRENT_TIMEOUT", 5))) + timeout = aiohttp.ClientTimeout(total=settings.GET_TORRENT_TIMEOUT) response = await session.get(url, allow_redirects=False, timeout=timeout) if response.status == 200: torrentData = await response.read() @@ -151,9 +152,12 @@ async def generateDownloadLink(debridApiKey: str, hash: str, index: str): proxy = None if "Your ISP or VPN provider IP address is currently blocked on our website" in checkBlacklisted: - proxy = os.getenv("DEBRID_PROXY_URL", "http://127.0.0.1:1080") - - logger.warning(f"Real-Debrid blacklisted server's IP. Switching to proxy {proxy} for {hash}|{index}") + proxy = settings.DEBRID_PROXY_URL + if not proxy: + logger.warning(f"Real-Debrid blacklisted server's IP. No proxy found.") + return "https://comet.fast" # TODO: This needs to be handled better + else: + logger.warning(f"Real-Debrid blacklisted server's IP. Switching to proxy {proxy} for {hash}|{index}") addMagnet = await session.post(f"https://api.real-debrid.com/rest/1.0/torrents/addMagnet", headers={ "Authorization": f"Bearer {debridApiKey}" diff --git a/comet/utils/models.py b/comet/utils/models.py index 8fa5bf8..49efbbc 100644 --- a/comet/utils/models.py +++ b/comet/utils/models.py @@ -15,7 +15,7 @@ class AppSettings(BaseSettings): FASTAPI_HOST: str = "0.0.0.0" FASTAPI_PORT: int = 8000 FASTAPI_WORKERS: int = 2 * (os.cpu_count() or 1) - DATABASE_PATH: str = "database.db" + DATABASE_PATH: str = "data/comet.db" CACHE_TTL: int = 86400 GET_TORRENT_TIMEOUT: int = 5 INDEXER_MANAGER_INDEXERS: List[str] = ["jackett", "qbittorrent"]