Merge pull request #7 from captaincolonelfox/dev

Release 1.5.0
captaincolonelfox · Apr 29, 2021 · 5b0e932 · 5b0e932
2 parents 9e53ecf + 329ec3a
commit 5b0e932
Show file tree

Hide file tree

Showing 8 changed files with 91 additions and 78 deletions.
diff --git a/README.md b/README.md
@@ -2,6 +2,14 @@
 
 # [TeleTok](https://t.me/TeleTockerBot): Telegram bot for TikTok (and Likee)
 
+## Description
+
+This bot will send you a video from a TikTok. Pretty simple.
+
+Just share a link to the chat (no need to mention the bot)
+
+## Thanks to
+
 Built on top of [aiogram](https://github.com/aiogram/aiogram)
 
 # Installation

diff --git a/bot/__init__.py b/bot/__init__.py
@@ -1,9 +1,11 @@
 from aiogram import Bot, Dispatcher
-
 from bot.api.likee import LikeeAPI
 from bot.api.tiktok import TikTokAPI
 from settings import API_TOKEN
 
 
 bot = Bot(token=API_TOKEN)
-dp = Dispatcher(bot=bot)
+dp = Dispatcher(bot=bot)
+
+
+from . import handlers
diff --git a/bot/api/__init__.py b/bot/api/__init__.py
@@ -1,65 +1,3 @@
-import asyncio
-import re
-from abc import ABC, abstractmethod
-from typing import Any, List, Optional
-
-import httpx
-import sentry_sdk
-from aiogram.types import Message
-from bs4 import BeautifulSoup
-from httpx import HTTPStatusError
-
-from bot.data import VideoData
-
-
-class API(ABC):
-    client = httpx.AsyncClient(
-        headers={
-            "User-Agent": 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_10_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/39.0.2171.95 Safari/537.36'
-        }
-    )
-
-    @property
-    def headers(self) -> dict[str, Any]:
-        return {}
-
-    @property
-    @abstractmethod
-    def links(self) -> List[str]:
-        return ['platform.com']
-
-    @property
-    @abstractmethod
-    def regexp_key(self) -> str:
-        return 'key'
-
-    async def handle_message(self, message: Message) -> List[VideoData]:
-        urls = []
-        for e in message.entities:
-            for link in self.links:
-                if link in (url := message.text[e.offset:e.offset + e.length]):
-                    urls.append(url)
-        if not urls:
-            return []
-        try:
-            return [await self.download_video(url) for url in urls]
-        except (KeyError, HTTPStatusError) as ex:
-            sentry_sdk.capture_exception(ex)
-            return []
-
-    async def download_video(self, url: str, retries: int = 2) -> VideoData:
-        for _ in range(retries):
-            page = await self.client.get(url)
-            soup = BeautifulSoup(page.content, 'html.parser')
-            if data := soup(text=re.compile(self.regexp_key)):
-                for script in data:
-                    if link := self._parse_data(script):
-                        if video := await self.client.get(link, headers=self.headers):
-                            video.raise_for_status()
-                            return VideoData(link, video.content)
-            await asyncio.sleep(0.5)
-        return VideoData()
-
-    @abstractmethod
-    def _parse_data(self, script: str) -> Optional[str]:
-        pass
+from .base import API
+from .likee import LikeeAPI
+from .tiktok import TikTokAPI
diff --git a/bot/api/base.py b/bot/api/base.py
@@ -0,0 +1,62 @@
+import asyncio
+import re
+from abc import ABC, abstractmethod
+from typing import Any, List, Optional
+
+import httpx
+import sentry_sdk
+from aiogram.types import Message
+from bs4 import BeautifulSoup
+from httpcore import TimeoutException
+from httpx import HTTPStatusError
+from bot.data import VideoData
+
+
+class API(ABC):
+
+    @property
+    def headers(self) -> dict[str, Any]:
+        return {}
+
+    @property
+    @abstractmethod
+    def links(self) -> List[str]:
+        return ['platform.com']
+
+    @property
+    @abstractmethod
+    def regexp_key(self) -> str:
+        return 'key'
+
+    async def handle_message(self, message: Message) -> List[VideoData]:
+        urls = []
+        for e in message.entities:
+            for link in self.links:
+                if link in (url := message.text[e.offset:e.offset + e.length]):
+                    urls.append(url if url.startswith('http') else f'https://{url}')
+        try:
+            return [await self.download_video(url) for url in urls]
+        except (KeyError, HTTPStatusError) as ex:
+            sentry_sdk.capture_exception(ex)
+        return []
+
+    async def download_video(self, url: str, retries: int = 2) -> VideoData:
+        for _ in range(retries):
+            async with httpx.AsyncClient(headers=self.headers) as client:
+                try:
+                    page = await client.get(url)
+                    soup = BeautifulSoup(page.content, 'html.parser')
+                    if data := soup(text=re.compile(self.regexp_key)):
+                        for script in data:
+                            if link := self._parse_data(script):
+                                if video := await client.get(link):
+                                    video.raise_for_status()
+                                    return VideoData(link, video.content)
+                except TimeoutException:
+                    pass
+                await asyncio.sleep(0.5)
+        return VideoData()
+
+    @abstractmethod
+    def _parse_data(self, script: str) -> Optional[str]:
+        pass
diff --git a/bot/api/likee.py b/bot/api/likee.py
@@ -8,7 +8,11 @@ class LikeeAPI(API):
 
     @property
     def headers(self) -> dict[str, Any]:
-        return {'Referer': 'https://www.likee.video/'}
+        return {
+            "Referer": "https://www.likee.video/",
+            "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_10_1) AppleWebKit/537.36"
+                          "(KHTML, like Gecko) Chrome/39.0.2171.95 Safari/537.36",
+        }
 
     @property
     def links(self):

diff --git a/bot/api/tiktok.py b/bot/api/tiktok.py
@@ -8,7 +8,11 @@ class TikTokAPI(API):
 
     @property
     def headers(self) -> dict[str, Any]:
-        return {'Referer': 'https://www.tiktok.com/'}
+        return {
+            "Referer": "https://www.tiktok.com/",
+            "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_10_1) AppleWebKit/537.36"
+                          "(KHTML, like Gecko) Chrome/39.0.2171.95 Safari/537.36",
+        }
 
     @property
     def links(self):

diff --git a/bot/handlers/messages.py b/bot/handlers/messages.py
@@ -1,15 +1,12 @@
 import sentry_sdk
 from aiogram.types import Message
 from bot import dp, bot
-from bot.api.likee import LikeeAPI
-from bot.api.tiktok import TikTokAPI
+from bot.api import LikeeAPI, TikTokAPI
 from bot.exception import HandleException
 from settings import DOWNLOAD_ERROR
 
 
-tiktok = TikTokAPI()
-likee = LikeeAPI()
-platforms = [tiktok, likee]
+platforms = [TikTokAPI(), LikeeAPI()]
 
 
 @dp.message_handler()

diff --git a/main.py b/main.py
@@ -4,12 +4,11 @@
 import sentry_sdk
 from sentry_sdk.integrations.aiohttp import AioHttpIntegration
 
-from bot import bot, dp, handlers
-from bot.handlers.messages import platforms
+from bot import bot, dp
 from settings import ENVIRONMENT, SENTRY_DSN
 
 sentry_sdk.init(
-    SENTRY_DSN,
+    dsn=SENTRY_DSN,
     environment=ENVIRONMENT,
     integrations=[AioHttpIntegration()]
 )
@@ -22,7 +21,6 @@ async def main():
     finally:
         logging.info('Exited')
         await bot.close()
-        [await platform.client.aclose() for platform in platforms]
 
 
 if __name__ == '__main__':