diff --git a/requirements.txt b/requirements.txt index 27fa70f..317fe62 100644 --- a/requirements.txt +++ b/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 diff --git a/ytdlbot/db.py b/ytdlbot/db.py index fe054a5..c87f693 100644 --- a/ytdlbot/db.py +++ b/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): diff --git a/ytdlbot/limit.py b/ytdlbot/limit.py index 7252e13..f3ade07 100644 --- a/ytdlbot/limit.py +++ b/ytdlbot/limit.py @@ -10,12 +10,15 @@ __author__ = "Benny " 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): diff --git a/ytdlbot/ytdl_bot.py b/ytdlbot/ytdl_bot.py index 87d7815..5a8a8d7 100644 --- a/ytdlbot/ytdl_bot.py +++ b/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