basic sub/unsub function

requirements.txt

@@ -10,6 +10,7 @@ filetype==1.0.9
 flower==1.0.0
 psutil==5.9.0
 influxdb==5.3.1
+beautifulsoup4==4.10.0
 
 fakeredis
 supervisor

ytdlbot/db.py

@@ -166,6 +166,28 @@ class MySQL:
     );
             """
 
+    channel_sql = """
+    create table if not exists channel
+    (
+        link              varchar(256) null,
+        title             varchar(256) null,
+        description       text         null,
+        channel_id        varchar(256),
+        playlist          varchar(256) null,
+        latest_video varchar(256) null,
+        constraint channel_pk
+            primary key (channel_id)
+    );
+    """
+
+    subscribe_sql = """
+    create table if not exists subscribe
+    (
+        user_id    bigint       null,
+        channel_id varchar(256) null
+    );
+    """
+
     def __init__(self):
         if MYSQL_HOST:
             self.con = pymysql.connect(host=MYSQL_HOST, user=MYSQL_USER, passwd=MYSQL_PASS, db="vip", charset="utf8mb4")
@@ -178,6 +200,8 @@ class MySQL:
     def init_db(self):
         self.cur.execute(self.vip_sql)
         self.cur.execute(self.settings_sql)
+        self.cur.execute(self.channel_sql)
+        self.cur.execute(self.subscribe_sql)
         self.con.commit()
 
     def __del__(self):

ytdlbot/limit.py

@@ -10,12 +10,15 @@ __author__ = "Benny <benny.think@gmail.com>"
 import hashlib
 import logging
 import math
+import os
 import tempfile
 import time
 from unittest.mock import MagicMock
 
 import requests
 
+import requests
+from bs4 import BeautifulSoup
 from config import (AFD_TOKEN, AFD_USER_ID, COFFEE_TOKEN, ENABLE_VIP, EX,
                     MULTIPLY, OWNER, QUOTA, USD2CNY)
 from db import MySQL, Redis
@@ -79,6 +82,73 @@ class VIP(Redis, MySQL):
         else:
             self.r.set(user_id, user_quota - traffic, ex=EX)
 
+    def subscribe_channel(self, user_id: "int", share_link: "str"):
+        data = self.get_channel_info(share_link)
+        self.cur.execute(
+            "INSERT INTO channel values(%(link)s,%(title)s,%(description)s,%(channel_id)s,%(playlist)s,%(last_video)s)",
+            data)
+        self.cur.execute("INSERT INTO subscribe values(%s,%s)", (user_id, data["channel_id"]))
+        self.con.commit()
+        return data["title"]
+
+    def unsubscribe_channel(self, user_id: "int", channel_id: "str"):
+        affected_rows = self.cur.execute("DELETE FROM subscribe WHERE user_id=%s AND channel_id=%s",
+                                         (user_id, channel_id))
+        self.con.commit()
+        return affected_rows
+
+    @staticmethod
+    def get_channel_info(url: "str"):
+        api_key = os.getenv("GOOGLE_API_KEY")
+        html_doc = requests.get(url).text
+        soup = BeautifulSoup(html_doc, 'html.parser')
+        element = soup.find("link", rel="canonical")
+        channel_id = element['href'].split("https://www.youtube.com/channel/")[1]
+        channel_api = f"https://www.googleapis.com/youtube/v3/channels?part=snippet,contentDetails&" \
+                      f"id={channel_id}&key={api_key}"
+        data = requests.get(channel_api).json()
+        snippet = data['items'][0]['snippet']
+        title = snippet['title']
+        description = snippet['description']
+        playlist = data['items'][0]['contentDetails']['relatedPlaylists']['uploads']
+
+        return {
+            "link": url,
+            "title": title,
+            "description": description,
+            "channel_id": channel_id,
+            "playlist": playlist,
+            "last_video": VIP.get_latest_video(playlist)
+        }
+
+    @staticmethod
+    def get_latest_video(playlist_id: "str"):
+        api_key = os.getenv("GOOGLE_API_KEY")
+        video_api = f"https://www.googleapis.com/youtube/v3/playlistItems?part=snippet&maxResults=1&" \
+                    f"playlistId={playlist_id}&key={api_key}"
+        data = requests.get(video_api).json()
+        video_id = data['items'][0]['snippet']['resourceId']['videoId']
+        return f"https://www.youtube.com/watch?v={video_id}"
+
+    def has_newer_update(self, playlist_id: "str"):
+        self.cur.execute("SELECT last_video FROM channel WHERE playlist=%s", (playlist_id,))
+        old_video = self.cur.fetchone()[0]
+        newest_video = VIP.get_latest_video(playlist_id)
+        if old_video != newest_video:
+            return newest_video
+
+    def get_user_subscription(self, user_id: "int"):
+        self.cur.execute(
+            """
+               select title, link, channel.channel_id from channel, subscribe 
+               where subscribe.user_id = %s and channel.channel_id = subscribe.channel_id
+            """, (user_id,))
+        data = self.cur.fetchall()
+        text = ""
+        for item in data:
+            text += "[{}]({}) `{}\n`".format(*item)
+        return text
+
 
 class BuyMeACoffee:
     def __init__(self):

ytdlbot/ytdl_bot.py

@@ -11,7 +11,7 @@ import logging
 import os
 import re
 import typing
-
+from io import BytesIO, StringIO
 from apscheduler.schedulers.background import BackgroundScheduler
 from pyrogram import Client, filters, types
 from pyrogram.errors.exceptions.bad_request_400 import UserNotParticipant
@@ -23,7 +23,7 @@ from config import (AUTHORIZED_USER, ENABLE_CELERY, ENABLE_VIP, OWNER,
                     REQUIRED_MEMBERSHIP)
 from constant import BotText
 from db import InfluxDB, MySQL, Redis
-from limit import verify_payment
+from limit import verify_payment, VIP
 from tasks import app as celery_app
 from tasks import (audio_entrance, direct_download_entrance, hot_patch,
                    ytdl_download_entrance)
@@ -92,6 +92,38 @@ def help_handler(client: "Client", message: "types.Message"):
     client.send_message(chat_id, bot_text.help, disable_web_page_preview=True)
 
 
+@app.on_message(filters.command(["subscribe"]))
+def subscribe_handler(client: "Client", message: "types.Message"):
+    vip = VIP()
+    chat_id = message.chat.id
+    client.send_chat_action(chat_id, "typing")
+    if message.text == "/subscribe":
+        result = vip.get_user_subscription(chat_id)
+    else:
+        link = message.text.split(" ")[1]
+        title = vip.subscribe_channel(chat_id, link)
+        result = f"Subscribed to {title}"
+    client.send_message(chat_id, result, disable_web_page_preview=True)
+
+
+@app.on_message(filters.command(["unsubscribe"]))
+def unsubscribe_handler(client: "Client", message: "types.Message"):
+    vip = VIP()
+    chat_id = message.chat.id
+    client.send_chat_action(chat_id, "typing")
+    text = message.text.split(" ")
+    if len(text) == 1:
+        client.send_message(chat_id, "/unsubscribe channel_id", disable_web_page_preview=True)
+        return
+
+    rows = vip.unsubscribe_channel(chat_id, text[1])
+    if rows:
+        text = f"Unsubscribed from {text[1]}"
+    else:
+        text = "Unable to find the channel."
+    client.send_message(chat_id, text, disable_web_page_preview=True)
+
+
 @app.on_message(filters.command(["hot_patch"]))
 def help_handler(client: "Client", message: "types.Message"):
     username = message.from_user.username