feat: restructured/refactored/reorganized

comet/api/core.py

@@ -0,0 +1,51 @@
+import os
+
+from fastapi import APIRouter, Request
+from fastapi.responses import RedirectResponse
+from fastapi.templating import Jinja2Templates
+
+templates = Jinja2Templates("comet/templates")
+main = APIRouter()
+
+
+@main.get("/", status_code=200)
+async def root():
+    return RedirectResponse("/configure")
+
+
+@main.get("/health", status_code=200)
+async def health():
+    return {"status": "ok"}
+
+
+@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", "")})
+
+
+@main.get("/manifest.json")
+@main.get("/{b64config}/manifest.json")
+async def manifest():
+    return {
+        "id": "stremio.comet.fast",
+        "version": "1.0.0",
+        "name": "Comet",
+        "description": "Stremio's fastest torrent/debrid search add-on.",
+        "logo": "https://i.imgur.com/jmVoVMu.jpeg",
+        "background": "https://i.imgur.com/WwnXB3k.jpeg",
+        "resources": [
+            "stream"
+        ],
+        "types": [
+            "movie",
+            "series"
+        ],
+        "idPrefixes": [
+            "tt"
+        ],
+        "catalogs": [],
+        "behaviorHints": {
+            "configurable": True
+        }
+    }

comet/api/stream.py

@@ -0,0 +1,229 @@
+import asyncio
+import hashlib
+import json
+import os
+import time
+from typing import Dict
+
+import aiohttp
+from fastapi import APIRouter, Request
+from fastapi.responses import RedirectResponse
+from RTN import ParsedData, Torrent, parse, sort_torrents
+
+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
+
+streams = APIRouter()
+
+@streams.get("/stream/{type}/{id}.json")
+@streams.get("/{b64config}/stream/{type}/{id}.json")
+async def stream(request: Request, b64config: str, type: str, id: str):
+    config = configChecking(b64config)
+    if not config:
+            return {
+                "streams": [
+                    {
+                        "name": "[⚠️] Comet", 
+                        "title": "Invalid Comet config.",
+                        "url": "https://comet.fast"
+                    }
+                ]
+            }
+
+    async with aiohttp.ClientSession() as session:
+        checkDebrid = await session.get("https://api.real-debrid.com/rest/1.0/user", headers={
+            "Authorization": f"Bearer {config['debridApiKey']}"
+        })
+        checkDebrid = await checkDebrid.text()
+        if not '"type": "premium"' in checkDebrid:
+            return {
+                "streams": [
+                    {
+                        "name": "[⚠️] Comet", 
+                        "title": "Invalid Real-Debrid account.",
+                        "url": "https://comet.fast"
+                    }
+                ]
+            }
+
+        season = None
+        episode = None
+        if type == "series":
+            info = id.split(":")
+
+            id = info[0]
+            season = int(info[1])
+            episode = int(info[2])
+
+        getMetadata = await session.get(f"https://v3.sg.media-imdb.com/suggestion/a/{id}.json")
+        metadata = await getMetadata.json()
+
+        name = metadata["d"][0]["l"]
+        name = translate(name)
+
+        cacheKey = hashlib.md5(json.dumps({"debridService": config["debridService"], "name": name, "season": season, "episode": episode, "indexers": config["indexers"], "resolutions": config["resolutions"], "languages": config["languages"]}).encode("utf-8")).hexdigest()
+        cached = await database.fetch_one(f"SELECT EXISTS (SELECT 1 FROM cache WHERE cacheKey = '{cacheKey}')")
+        if cached[0] != 0:
+            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():
+                await database.execute(f"DELETE FROM cache WHERE cacheKey = '{cacheKey}'")
+
+                logger.info(f"Cache expired for {name}")
+            else:
+                sortedRankedFiles = await database.fetch_one(f"SELECT results FROM cache WHERE cacheKey = '{cacheKey}'")
+                sortedRankedFiles = json.loads(sortedRankedFiles[0])
+
+                results = []
+                for hash in sortedRankedFiles:
+                    results.append({
+                        "name": f"[RD⚡] Comet {sortedRankedFiles[hash]['data']['resolution'][0] if len(sortedRankedFiles[hash]['data']['resolution']) > 0 else 'Unknown'}",
+                        "title": f"{sortedRankedFiles[hash]['data']['title']}\n💾 {bytesToSize(sortedRankedFiles[hash]['data']['size'])}",
+                        "url": f"{request.url.scheme}://{request.url.netloc}/{b64config}/playback/{hash}/{sortedRankedFiles[hash]['data']['index']}"
+                    })
+
+                return {"streams": results}
+        else:
+            logger.info(f"No cache found for {name} with user configuration")
+
+        indexerManagerType = os.getenv("INDEXER_MANAGER_TYPE", "jackett")
+
+        logger.info(f"Start of {indexerManagerType} search for {name} with indexers {config['indexers']}")
+
+        tasks = []
+        tasks.append(getIndexerManager(session, indexerManagerType, config["indexers"], name))
+        if type == "series":
+            tasks.append(getIndexerManager(session, indexerManagerType, config["indexers"], f"{name} S0{season}E0{episode}"))
+        searchResponses = await asyncio.gather(*tasks)
+
+        torrents = []
+        for results in searchResponses:
+            if results == None:
+                continue
+
+            for result in results:
+                torrents.append(result)
+
+        logger.info(f"{len(torrents)} torrents found for {name}")
+
+        if len(torrents) == 0:
+            return {"streams": []}
+
+        tasks = []
+        filtered = 0
+        for torrent in torrents:
+            parsedTorrent: ParsedData = parse(torrent["Title"]) if indexerManagerType == "jackett" else parse(torrent["title"])
+            if not "All" in config["resolutions"] and len(parsedTorrent.resolution) > 0 and parsedTorrent.resolution[0] not in config["resolutions"]:
+                filtered += 1
+
+                continue
+            if not "All" in config["languages"] and not parsedTorrent.is_multi_audio and not any(language in parsedTorrent.language for language in config["languages"]):
+                filtered += 1
+
+                continue
+
+            tasks.append(getTorrentHash(session, indexerManagerType, torrent))
+
+        torrentHashes = await asyncio.gather(*tasks)
+        torrentHashes = list(set([hash for hash in torrentHashes if hash]))
+
+        logger.info(f"{len(torrentHashes)} info hashes found for {name}")
+
+        if len(torrentHashes) == 0:
+            return {"streams": []}
+
+        getAvailability = await session.get(f"https://api.real-debrid.com/rest/1.0/torrents/instantAvailability/{'/'.join(torrentHashes)}", headers={
+            "Authorization": f"Bearer {config['debridApiKey']}"
+        })
+
+        files = {}
+
+        availability = await getAvailability.json()
+        for hash, details in availability.items():
+            if not "rd" in details:
+                continue
+
+            if type == "series":
+                for variants in details["rd"]:
+                    for index, file in variants.items():
+                        filename = file["filename"].lower()
+
+                        if not isVideo(filename):
+                            continue
+
+                        filenameParsed = parse(file["filename"])
+                        if season in filenameParsed.season and episode in filenameParsed.episode:
+                            files[hash] = {
+                                "index": index,
+                                "title": file["filename"],
+                                "size": file["filesize"]
+                            }
+
+                continue
+
+            for variants in details["rd"]:
+                for index, file in variants.items():
+                    filename = file["filename"].lower()
+
+                    if not isVideo(filename):
+                        continue
+
+                    files[hash] = {
+                        "index": index,
+                        "title": file["filename"],
+                        "size": file["filesize"]
+                    }
+
+        rankedFiles = set()
+        for hash in files:
+            # try:
+            rankedFile = rtn.rank(files[hash]["title"], hash) # , remove_trash=True, correct_title=name - removed because it's not working great
+            rankedFiles.add(rankedFile)
+            # except:
+            #     continue
+
+        sortedRankedFiles: Dict[str, Torrent] = sort_torrents(rankedFiles)
+
+        logger.info(f"{len(sortedRankedFiles)} cached files found on Real-Debrid for {name}")
+
+        if len(sortedRankedFiles) == 0:
+            return {"streams": []}
+
+        sortedRankedFiles = {
+            key: (value.model_dump() if isinstance(value, Torrent) else value)
+            for key, value in sortedRankedFiles.items()
+        }
+        for hash in sortedRankedFiles: # needed for caching
+            sortedRankedFiles[hash]["data"]["title"] = files[hash]["title"]
+            sortedRankedFiles[hash]["data"]["size"] = files[hash]["size"]
+            sortedRankedFiles[hash]["data"]["index"] = files[hash]["index"]
+
+        jsonData = json.dumps(sortedRankedFiles).replace("'", "''")
+        await database.execute(f"INSERT INTO cache (cacheKey, results, timestamp) VALUES ('{cacheKey}', '{jsonData}', {time.time()})")
+        logger.info(f"Results have been cached for {name}")
+
+        results = []
+        for hash in sortedRankedFiles:
+            results.append({
+                "name": f"[RD⚡] Comet {sortedRankedFiles[hash]['data']['resolution'][0] if len(sortedRankedFiles[hash]['data']['resolution']) > 0 else 'Unknown'}",
+                "title": f"{sortedRankedFiles[hash]['data']['title']}\n💾 {bytesToSize(sortedRankedFiles[hash]['data']['size'])}",
+                "url": f"{request.url.scheme}://{request.url.netloc}/{b64config}/playback/{hash}/{sortedRankedFiles[hash]['data']['index']}"
+            })
+
+        return {
+            "streams": results
+        }
+
+@streams.route("/{b64config}/playback/{hash}/{index}", methods=["HEAD", "GET"])
+async def stream(b64config: str, hash: str, index: str):
+    config = configChecking(b64config)
+    if not config:
+        return
+
+    downloadLink = await generateDownloadLink(config["debridApiKey"], hash, index)
+
+    return RedirectResponse(downloadLink, status_code=302)

comet/main.py

@@ -129,4 +129,3 @@ def start_log():
         logger.exception(traceback.format_exc())
     finally:
         logger.log("COMET", "Server Shutdown")
-

comet/utils/db.py

@@ -0,0 +1,47 @@
+import json
+import os
+from comet.utils.logger import logger
+from comet.utils.models import database, settings
+from comet.utils.general import lang_code_map
+
+
+async def setup_database():
+    """Setup the database by ensuring the directory and file exist, and creating the necessary tables."""
+    try:
+        # Ensure the database directory exists
+        os.makedirs(os.path.dirname(settings.DATABASE_PATH), exist_ok=True)
+
+        # Ensure the database file exists
+        if not os.path.exists(settings.DATABASE_PATH):
+            open(settings.DATABASE_PATH, 'a').close()
+
+        await database.connect()
+        await database.execute("CREATE TABLE IF NOT EXISTS cache (cacheKey BLOB PRIMARY KEY, timestamp INTEGER, results TEXT)")
+    except Exception as e:
+        logger.error(f"Error setting up the database: {e}")
+
+async def teardown_database():
+    """Teardown the database by disconnecting."""
+    try:
+        await database.disconnect()
+    except Exception as e:
+        # Log the exception or handle it as needed
+        print(f"Error tearing down the database: {e}")
+
+def write_config():
+    """Write the config file."""
+    indexers = settings.INDEXER_MANAGER_INDEXERS
+    if indexers:
+        if isinstance(indexers, str):
+            indexers = indexers.split(",")
+        elif not isinstance(indexers, list):
+            logger.warning("Invalid indexers")
+
+    config_data = {
+        "indexers": indexers,
+        "languages": lang_code_map,
+        "resolutions": ["480p", "720p", "1080p", "1440p", "2160p", "2880p", "4320p"]
+    }
+
+    with open("comet/templates/config.json", "w", encoding="utf-8") as config_file:
+        json.dump(config_data, config_file, indent=4)

comet/utils/models.py

@@ -0,0 +1,56 @@
+import os
+from typing import List, Optional
+
+from databases import Database
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from RTN import RTN, BaseRankingModel, SettingsModel
+
+
+class AppSettings(BaseSettings):
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        env_file_encoding="utf-8",
+    )
+
+    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"
+    CACHE_TTL: int = 86400
+    GET_TORRENT_TIMEOUT: int = 5
+    INDEXER_MANAGER_INDEXERS: List[str] = ["jackett", "qbittorrent"]
+    INDEXER_MANAGER_TYPE: str = "jackett"
+    INDEXER_MANAGER_URL: str = "http://127.0.0.1:9117"
+    INDEXER_MANAGER_API_KEY: str = ""
+    INDEXER_MANAGER_TIMEOUT: int = 30
+    DEBRID_PROXY_URL: Optional[str] = None
+    CUSTOM_HEADER_HTML: Optional[str] = None
+
+
+class BestOverallRanking(BaseRankingModel):
+    uhd: int = 100
+    fhd: int = 90
+    hd: int = 80
+    sd: int = 70
+    dolby_video: int = 100
+    hdr: int = 80
+    hdr10: int = 90
+    dts_x: int = 100
+    dts_hd: int = 80
+    dts_hd_ma: int = 90
+    atmos: int = 90
+    truehd: int = 60
+    ddplus: int = 40
+    aac: int = 30
+    ac3: int = 20
+    remux: int = 150
+    bluray: int = 120
+    webdl: int = 90
+
+rtn_settings: SettingsModel = SettingsModel()
+rtn_ranking: BestOverallRanking = BestOverallRanking()
+
+# For use anywhere
+rtn: RTN = RTN(settings=rtn_settings, ranking_model=rtn_ranking)
+settings: AppSettings = AppSettings()
+database = Database(f"sqlite:///{settings.DATABASE_PATH}")