upgrade to cobalt 10

alexBot/cobalt.py

@@ -1,137 +1,118 @@
+import os
 from dataclasses import asdict, dataclass
-from typing import List, Literal, Optional, Dict
+from typing import Dict, List, Literal, Optional
 
-import aiohttp
+import urllib.request
+import logging
 
-ENDPOINT = "https://api.cobalt.tools"
+import aiohttp
 
+ENDPOINT = os.environ.get("COBALT_URL") or "http://cobalt-api:9000"
+FALLBACK_ENDPOINT = "https://api.cobalt.tools/"
 
 DEFAULT_HEADERS = {"Accept": "application/json", "Content-Type": "application/json", "User-Agent": "alexBot/1.0"}
 
-# comments are from https://github.com/wukko/cobalt/blob/current/docs/api.md
-
-
-# ## POST: `/api/json`
-# cobalt's main processing endpoint.
-
-# request body type: `application/json`
-# response body type: `application/json`
-
-# ```
-# ⚠️ you must include Accept and Content-Type headers with every POST /api/json request.
-
-# Accept: application/json
-# Content-Type: application/json
-# ```
-
-
-# ### request body variables
-# | key               | type      | variables                          | default   | description                                                                     |
-# |:------------------|:----------|:-----------------------------------|:----------|:--------------------------------------------------------------------------------|
-# | `url`             | `string`  | URL encoded as URI                 | `null`    | **must** be included in every request.                                          |
-# | `vCodec`          | `string`  | `h264 / av1 / vp9`                 | `h264`    | applies only to youtube downloads. `h264` is recommended for phones.            |
-# | `vQuality`        | `string`  | `144 / ... / 2160 / max`           | `720`     | `720` quality is recommended for phones.                                        |
-# | `aFormat`         | `string`  | `best / mp3 / ogg / wav / opus`    | `mp3`     |                                                                                 |
-# | `filenamePattern` | `string`  | `classic / pretty / basic / nerdy` | `classic` | changes the way files are named. previews can be seen in the web app.           |
-# | `isAudioOnly`     | `boolean` | `true / false`                     | `false`   |                                                                                 |
-# | `isTTFullAudio`   | `boolean` | `true / false`                     | `false`   | enables download of original sound used in a tiktok video.                      |
-# | `isAudioMuted`    | `boolean` | `true / false`                     | `false`   | disables audio track in video downloads.                                        |
-# | `dubLang`         | `boolean` | `true / false`                     | `false`   | backend uses Accept-Language header for youtube video audio tracks when `true`. |
-# | `disableMetadata` | `boolean` | `true / false`                     | `false`   | disables file metadata when set to `true`.                                      |
-# | `twitterGif`      | `boolean` | `true / false`                     | `false`   | changes whether twitter gifs are converted to .gif                              |
+log = logging.getLogger(__name__)
 
 
 @dataclass
 class RequestBody:
     url: str
-    vCodec: Literal["h264", "av1", "vp9"] = "h264"
-    vQuality: Literal["max", "4320", "2160", "1440", "1080", "720", "480", "360", "240", "144"] = "720"
-    aFormat: Literal["best", "mp3", "ogg", "wav", "opus"] = "mp3"
-    filenamePattern: Literal["classic", "pretty", "basic", "nerdy"] = "classic"
-    isAudioOnly: bool = False
-    isTTFullAudio: bool = False
-    isAudioMuted: bool = False
-    dubLang: bool = False
+    videoQuality: Literal["max", "4320", "2160", "1440", "1080", "720", "480", "360", "240", "144"] = "720"
+    audioFormat: Literal["best", "mp3", "ogg", "wav", "opus"] = "mp3"
+    filenameStyle: Literal["classic", "pretty", "basic", "nerdy"] = "classic"
+    downloadMode: Literal["auto", "audio", "mute"] = "auto"
+    youtubeVideoCodec: Literal["h264", "av1", "vp9"] = "h264"
+    youtubeDubLang: Optional[Literal["en", "ru", "cs", "ja"]] = None
+    youtubeDubBrowserLang: bool = False
+    alwaysProxy: bool = False
     disableMetadata: bool = False
-    twitterGif: bool = False
-
+    tiktokFullAudio: bool = False
+    tiktokH265: bool = False
+    twitterGif: bool = True
 
-# | key     | type     | variables                                               | description                            |
-# |:--------|:---------|:--------------------------------------------------------|:---------------------------------------|
-# | `type`  | `string` | `video`                                                 | used only if `pickerType`is `various`. |
-# | `url`   | `string` | direct link to a file or a link to cobalt's live render |                                        |
-# | `thumb` | `string` | item thumbnail that's displayed in the picker           | used only for `video` type.            |
+    def dict(self):
+        without_none = {k: v for k, v in asdict(self).items() if v is not None}
+        return without_none
 
 
 @dataclass
 class Picker:
     url: str
-    type: Optional[Literal["video"]] = None
+    type: Optional[Literal["video", "photo", "gif"]] = None
     thumb: Optional[str] = None
+    data: Optional[bytes] = None
+
+    async def fetch(self, session: aiohttp.ClientSession) -> bytes:
+        async with session.get(self.url) as response:
+            self.data = await response.read()
+            return self.data
 
 
-# ### response body variables
-# | key          | type     | variables                                                   |
-# |:-------------|:---------|:------------------------------------------------------------|
-# | `status`     | `string` | `error / redirect / stream / success / rate-limit / picker` |
-# | `text`       | `string` | various text, mostly used for errors                        |
-# | `url`        | `string` | direct link to a file or a link to cobalt's live render     |
-# | `pickerType` | `string` | `various / images`                                          |
-# | `picker`     | `array`  | array of picker items                                       |
-# | `audio`      | `string` | direct link to a file or a link to cobalt's live render     |
 @dataclass
 class ResponceBody:
-    status: Literal["error", "redirect", "stream", "success", "rate-limit", "picker"]
+    status: Literal["error", "redirect", "tunnel", "picker"]
+    error: Optional[str] = None
+
+    # status: "redirect", "tunnel"
     url: Optional[str] = None
-    text: Optional[str] = None
-    pickerType: Optional[Literal["various", "images"]] = None
-    picker: Optional[List[Picker]] = None
-    audio: Optional[str] = None
-    raw_picker: Optional[List[dict]] = None
+    filename: Optional[str] = None
 
+    # status: "picker"
+    audio: Optional[str] = None
+    audioFilename: Optional[str] = None
+    picker: Optional[List[Picker]] = None
+    _picker: Optional[List[Dict]] = None
 
-# ## GET: `/api/serverInfo`
-# returns current basic server info.
-# response body type: `application/json`
 
-# ### response body variables
-# | key         | type     | variables         |
-# |:------------|:---------|:------------------|
-# | `version`   | `string` | cobalt version    |
-# | `commit`    | `string` | git commit        |
-# | `branch`    | `string` | git branch        |
-# | `name`      | `string` | server name       |
-# | `url`       | `string` | server url        |
-# | `cors`      | `int`    | cors status       |
-# | `startTime` | `string` | server start time |
+@dataclass
+class CobaltServerData:
+    version: str
+    url: str
+    startTime: str
+    durationLimit: int
+    services: List[str]
 
 
 @dataclass
-class ServerInfo:
-    version: str
+class GitServerData:
     commit: str
     branch: str
-    name: str
-    url: str
-    cors: int
-    startTime: str
+    remote: str
+
+
+@dataclass
+class ServerInfo:
+    cobalt: CobaltServerData
+    git: GitServerData
 
 
 class Cobalt:
-    HEADERS = DEFAULT_HEADERS
+    def __init__(self) -> None:
+        # make a request to ENDPOINT and check if it's up, if not, set to fallback server
+        try:
+            contents = urllib.request.urlopen(ENDPOINT).read()
+        except Exception:
+            contents = None
+        if contents is None:
+            self.endpoint = FALLBACK_ENDPOINT
+            log.warning(f"Cobalt API {ENDPOINT} is down, using fallback server ({FALLBACK_ENDPOINT})")
+        else:
+            self.endpoint = ENDPOINT
+        self.headers = DEFAULT_HEADERS
 
     async def get_server_info(self):
-        async with aiohttp.ClientSession(headers=self.HEADERS) as session:
-            async with session.get(ENDPOINT + "/api/serverInfo") as resp:
-                return ServerInfo(**await resp.json())
+        async with aiohttp.ClientSession(headers=self.headers) as session:
+            async with session.get(self.endpoint) as resp:
+                data = await resp.json()
+                return ServerInfo(cobalt=CobaltServerData(**data['cobalt']), git=GitServerData(**data['git']))
 
     async def process(self, request_body: RequestBody):
-        async with aiohttp.ClientSession(headers=self.HEADERS) as session:
-            async with session.post(ENDPOINT + "/api/json", json=asdict(request_body)) as resp:
+        async with aiohttp.ClientSession(headers=self.headers) as session:
+            async with session.post(self.endpoint, json=request_body.dict()) as resp:
                 rb = ResponceBody(**await resp.json())
-                if rb.picker:
-                    rb.raw_picker: List[Dict] = rb.picker  # type: ignore
-                    rb.picker = [
-                        Picker(picker['url'], picker.get("type"), picker.get("thumb")) for picker in rb.raw_picker
-                    ]
+                if rb.status == "picker":
+                    rb._picker = rb.picker
+                    rb.picker = [Picker(**p) for p in rb._picker]  # type: ignore
+
                 return rb

alexBot/cogs/video_dl.py

@@ -7,7 +7,7 @@
 import subprocess
 import traceback
 from functools import partial
-from typing import List
+from typing import List, Optional
 
 import aiohttp
 import discord
@@ -52,6 +52,29 @@ class NotAVideo(Exception):
 class Video_DL(Cog):
     encode_lock = asyncio.Lock()
     mirror_upload_lock = asyncio.Lock()
+    _cobalt: Optional[Cobalt] = None
+
+    @staticmethod
+    async def fetch_image(url: str, session: aiohttp.ClientSession, extra: any) -> bytes:
+        async with session.get(url) as response:
+            return await response.read()
+
+    async def get_cobalt_instace(self) -> Cobalt:
+        if self._cobalt:
+            return self._cobalt
+        self._cobalt = Cobalt()
+
+        async def test_cobalt():
+            while self._cobalt:
+                await asyncio.sleep(60 * 15)
+                try:
+                    await self._cobalt.get_server_info()
+                except Exception as e:
+                    log.error("Error testing cobalt", e)
+                    self._cobalt = None
+
+        self.bot.loop.create_task(test_cobalt())
+        return self._cobalt
 
     @Cog.listener()
     async def on_message(self, message: discord.Message, override=False, new_deleter=None):
@@ -94,14 +117,14 @@ async def on_message(self, message: discord.Message, override=False, new_deleter
             log.debug("Typing indicator started")
             stuff = None
 
-            cobalt = Cobalt()
+            cobalt = await self.get_cobalt_instace()
             rq = RequestBody(url=match.group(1))
             res = await cobalt.process(rq)
-            async with aiohttp.ClientSession(headers=cobalt.HEADERS) as session:
+            async with aiohttp.ClientSession(headers=cobalt.headers) as session:
                 match res.status:
-                    case "stream" | "redirect":
+                    case "tunnel" | "redirect":
                         # download the stream to reupload to discord
-                        log.debug("Status is stream or redirect. Downloading the stream.")
+                        log.debug("Status is stream or tunnel. Downloading the stream.")
                         async with session.get(res.url) as response:
                             stuff = await response.read()
                             if not response.content_disposition:
@@ -129,18 +152,16 @@ async def on_message(self, message: discord.Message, override=False, new_deleter
                         for m, group in enumerate(grouper(images, 10)):
                             for n, image in enumerate(group):
                                 stuff = await image.read()
-                                if not image.content_disposition:
-                                    filename = f"{n+(m*10)}.{image.url.suffix}"
-                                else:
-                                    filename = image.content_disposition.filename
+                                filename = f"{n+(m*10)}.{image.url.suffix}"
+
                                 attachments.append(discord.File(io.BytesIO(stuff), filename=filename))
                             try:
 
                                 uploaded = await message.reply(mention_author=False, files=attachments)
                             except DiscordException as e:
                                 log.error("Error uploading images", e)
                     case "error":
-                        log.error(f"Error in cobalt with url {rq.url}: {res.text}")
+                        log.error(f"Error in cobalt with url {rq.url}: {res.error}")
         log.debug("on_message function ended")
 
         if uploaded:

config.py

@@ -21,6 +21,7 @@
 
 prefix = os.environ.get('BOT_PREFIX')
 
+COBALT_URL = os.environ.get('COBALT_URL')
 
 # setthe DATABASE_URL env var
 db_full_url = os.environ.get("DATABASE_URL")

docker-compose.debug.yml

@@ -10,6 +10,7 @@ services:
       - .env
     environment:
       DATABASE_URL: postgresql://alexbot:alexbot@db/alexbot
+      COBALT_URL: "http://cobalt-api:9000"
     command:
       [
         "sh",
@@ -20,6 +21,7 @@ services:
       - 5678:5678
     depends_on:
       - db
+      - cobalt-api
   # firefox:
   #   image: selenium/standalone-firefox-debug
   #   shm_size: 2gb
@@ -35,5 +37,37 @@ services:
       - 5432:5432
     volumes:
       - alex-bot-db:/var/lib/postgresql/data
+  cobalt-api:
+    image: ghcr.io/imputnet/cobalt:10
+    restart: unless-stopped
+    container_name: cobalt-api
+
+    init: true
+
+    ports:
+        - 9000:9000/tcp
+        # if you're using a reverse proxy, uncomment the next line and remove the one above (9000:9000/tcp):
+        #- 127.0.0.1:9000:9000
+
+    environment:
+        # replace https://api.cobalt.tools/ with your instance's target url in same format
+        API_URL: "http://cobalt-api:9000/"
+        # if you want to use cookies when fetching data from services, uncomment the next line and the lines under volume
+        # COOKIE_PATH: "/cookies.json"
+        # see docs/run-an-instance.md for more information
+    labels:
+        - com.centurylinklabs.watchtower.scope=cobalt
+
+    # if you want to use cookies when fetching data from services, uncomment volumes and next line
+    #volumes:
+        #- ./cookies.json:/cookies.json
+
+# update the cobalt image automatically with watchtower
+  watchtower:
+    image: ghcr.io/containrrr/watchtower
+    restart: unless-stopped
+    command: --cleanup --scope cobalt --interval 900 --include-restarting
+    volumes:
+        - /var/run/docker.sock:/var/run/docker.sock
 volumes:
   alex-bot-db: