alldebrid: Add print links as HTML support

alldebrid.py

@@ -1,27 +1,51 @@
+#!/usr/bin/env -S uv run
 # /// script
-# requirements = ["httpx", "pydantic"]
+# dependencies = [
+#     "httpx",
+# ]
 # ///
+from concurrent.futures import ThreadPoolExecutor
 import dataclasses
+import itertools
 import os
 from pathlib import Path
 import httpx
 import argparse
 import time
-import pydantic
+import urllib.parse
 
 
-class MagnetStatus(pydantic.BaseModel):
-    id: int
-    filename: str
-    size: int
-    status: str
-    downloaded: int
-
-
-@dataclasses.dataclass()
 class AlldebridError(Exception):
-    code: str
+    def __init__(self, code: str, message: str):
-    message: str
+        self.code = code
+        self.message = message
+        super().__init__(f"{code}: {message}")
+
+
+@dataclasses.dataclass
+class Link:
+    filename: str
+    url: str
+    size: int
+    download_url: str | None = None
+
+    @property
+    def size_mb(self) -> float:
+        return self.size / 1_048_576
+
+
+def flatten_tree(data: dict | list) -> list[dict]:
+    if isinstance(data, list):
+        return itertools.chain.from_iterable(flatten_tree(item) for item in data)
+    elif isinstance(data, dict):
+        if "e" in data:  # It's a directory
+            return flatten_tree(data["e"])
+        elif "l" in data and "s" in data:  # It's a file
+            return [data]
+    return []
+
+
+pool = ThreadPoolExecutor(max_workers=10)
 
 
 class AlldebridDownloader:
@@ -41,7 +65,7 @@ class AlldebridDownloader:
             error = data["error"]
             raise AlldebridError(code=error["code"], message=error["message"])
 
-    def unrestrict_url(self, url: str) -> str | None:
+    def unrestrict_url(self, url: str) -> Link | None:
         res = self.client.get(
             "https://api.alldebrid.com/v4/link/unlock",
             params={
@@ -51,7 +75,16 @@ class AlldebridDownloader:
         self._check_error(res)
         data = res.json()
         if data["status"] == "success":
-            return data["data"].get("link")
+            return Link(
+                filename=data["data"]["filename"],
+                url=url,
+                size=data["data"]["filesize"],
+                download_url=data["data"]["link"],
+            )
+
+    def _unrestrict_link(self, link: Link) -> Link | None:
+        link.download_url = self.unrestrict_url(link.url).download_url
+        return link
 
     def upload_torrent(self, torrent_path: Path) -> int:
         res = self.client.post(
@@ -91,21 +124,21 @@ class AlldebridDownloader:
             magnets = data["data"]["magnets"]
             if not magnets:
                 return
-            files = magnets[0]["files"][0]["e"]
+            files = magnets[0]["files"]
             if not files:
-                raise ValueError("torrent is empty")
+                return
-            large_files = [m for m in files if m["s"] > 1_000_000]
+            files = [Link(url=it["l"], filename=it["n"], size=it["s"]) for it in flatten_tree(files)]
+            large_files = [f for f in files if f.size_mb > 5]
             if not large_files:
                 raise ValueError("torrent only contains small files")
-            links = [it["l"] for it in large_files]
+            return large_files
-            return links
 
-    def wait_for_download_link(self, magnet_id, timeout=600):
+    def wait_for_download_links(self, magnet_id, timeout=600) -> list[Link]:
         start = time.time()
         while time.time() - start < timeout:
             links = self.get_torrent_links(magnet_id)
             if links:
-                return self.unrestrict_url(links[0])
+                return list(pool.map(self._unrestrict_link, links))
             time.sleep(5)
         return None
 
@@ -114,34 +147,95 @@ def parse_args():
     parser = argparse.ArgumentParser(description="Alldebrid Downloader")
     parser.add_argument("input", help="URL, magnet link, or torrent file")
     parser.add_argument("--token", default=os.getenv("ALLDEBRID_TOKEN"), help="Alldebrid API token")
+    parser.add_argument("--html", action="store_true", help="Print links as HTML")
     return parser.parse_args()
 
 
+def print_links(links: list[Link]):
+    if not links:
+        return
+    for link in links:
+        print(link.download_url)
+
+
+def print_links_as_html(links: list[Link]):
+    if not links:
+        return
+    css = """
+        body {
+            margin: 0;
+            padding: 1rem;
+            font-family: consolas, menlo, monospace;
+            font-size: 14px;
+        }
+        table {
+            border-collapse: collapse;
+            width: 100%;
+        }
+        th, td {
+            border: 1px solid black;
+            padding: 8px;
+            text-align: left;
+        }
+        td.numeric {
+            text-align: right;
+        }
+    """
+    row_htmls = []
+    for it in links:
+        encoded_link = urllib.parse.quote(it.download_url)
+        alfred_link = f"alfred://runtrigger/piracy/direct_link/?argument={encoded_link}"
+        row_html = f"""
+            <tr>
+                <td><a href='{it.download_url}'>{it.filename}</a></td>
+                <td class='numeric'>{it.size_mb:.1f}</td>
+                <td><a href='{alfred_link}'>Alfred</a></td>
+            </tr>
+        """
+        row_htmls.append(row_html)
+
+    table_html = f"""
+        <style>{css}</style>
+        <script>window.app.setFloating(false)</script>
+        <table>
+            <thead><tr>
+                <th>Filename</th>
+                <th>Size MB</th>
+                <th>Action</th>
+            </tr></thead>
+            <tbody>
+            {"".join(row_htmls)}
+            </tbody>
+        </table>
+    """
+    print(table_html)
+
+
 def main():
     args = parse_args()
+    if not args.token:
+        raise ValueError("No token provided")
 
     downloader = AlldebridDownloader(args.token)
 
+    printer = print_links_as_html if args.html else print_links
+
     input_as_path = Path(args.input)
-    if input_as_path.is_file() and input_as_path.suffix == ".torrent":
+    if args.input.startswith("magnet:"):
+        magnet_id = downloader.upload_magnet(args.input)
+        if not magnet_id:
+            return
+        links = downloader.wait_for_download_links(magnet_id)
+        printer(links)
+    elif args.input.startswith("http:") or args.input.startswith("https:"):
+        link = downloader.unrestrict_url(args.input)
+        printer([link])
+    elif input_as_path.is_file() and input_as_path.suffix == ".torrent":
         torrent_id = downloader.upload_torrent(input_as_path)
         if not torrent_id:
             return
-        result = downloader.wait_for_download_link(torrent_id)
+        links = downloader.wait_for_download_links(torrent_id)
-        if result:
+        printer(links)
-            print(result)
-    elif args.input.startswith("magnet:"):
-        magnet_data = downloader.upload_magnet(args.input)
-        if not magnet_data:
-            return
-        result = downloader.wait_for_download_link(magnet_data["id"])
-        if result:
-            print(result)
-    elif args.input.startswith("http:") or args.input.startswith("https:"):
-        result = downloader.unrestrict_url(args.input)
-        if not result:
-            return
-        print(result)
     else:
         raise ValueError("Invalid input")