From a3a821bfae45a77c12b7212b3d5bc66e3851ff7b Mon Sep 17 00:00:00 2001
From: BennyThink <benny.think@gmail.com>
Date: Fri, 11 Feb 2022 19:07:28 +0800
Subject: [PATCH] add hidden 'Local' mode

---
 ytdlbot/db.py         |  1 +
 ytdlbot/migration.sql |  3 +++
 ytdlbot/tasks.py      | 21 +++++++++++----------
 ytdlbot/utils.py      |  2 +-
 ytdlbot/ytdl_bot.py   | 20 ++++++++++++++++++--
 5 files changed, 34 insertions(+), 13 deletions(-)
 create mode 100644 ytdlbot/migration.sql

diff --git a/ytdlbot/db.py b/ytdlbot/db.py
index 93ef1e5..517ff88 100644
--- a/ytdlbot/db.py
+++ b/ytdlbot/db.py
@@ -172,6 +172,7 @@ class MySQL:
         user_id    bigint          not null,
         resolution varchar(128) null,
         method     varchar(64)  null,
+        mode varchar(32) default 'Celery' null,
         constraint settings_pk
             primary key (user_id)
     );
diff --git a/ytdlbot/migration.sql b/ytdlbot/migration.sql
new file mode 100644
index 0000000..15fc385
--- /dev/null
+++ b/ytdlbot/migration.sql
@@ -0,0 +1,3 @@
+alter table settings
+    add mode varchar(32) default 'Celery' null;
+
diff --git a/ytdlbot/tasks.py b/ytdlbot/tasks.py
index aca24e8..ad5c1cd 100644
--- a/ytdlbot/tasks.py
+++ b/ytdlbot/tasks.py
@@ -24,6 +24,7 @@ from celery import Celery
 from celery.worker.control import Panel
 from pyrogram import idle
 from pyrogram.types import InlineKeyboardButton, InlineKeyboardMarkup
+from requests_toolbelt.multipart.encoder import MultipartEncoder
 
 from client_init import create_app
 from config import (ARCHIVE_ID, AUDIO_FORMAT, BROKER, ENABLE_CELERY,
@@ -124,8 +125,8 @@ def ytdl_download_entrance(bot_msg, client, url):
     chat_id = bot_msg.chat.id
     if forward_video(chat_id, url, client):
         return
-
-    if ENABLE_CELERY:
+    mode = get_user_settings(str(chat_id))[-1]
+    if ENABLE_CELERY and mode in [None, "Celery"]:
         ytdl_download_task.delay(chat_id, bot_msg.message_id, url)
     else:
         ytdl_normal_download(bot_msg, client, url)
@@ -231,14 +232,14 @@ def get_dl_source():
     return ""
 
 
-def upload_transfer_sh(video_paths) -> "str":
-    file = {}
-    for p in video_paths:
-        file[p.name] = p.open("rb")
+def upload_transfer_sh(paths: list) -> "str":
+    d = {p.name: (p.name, p.open("rb")) for p in paths}
+    m = MultipartEncoder(fields=d)
+    headers = {'Content-Type': m.content_type}
     try:
-        req = requests.post("https://transfer.sh", files=file)
-        return req.text
-    except requests.exceptions as e:
+        req = requests.post("https://transfer.sh", data=m, headers=headers)
+        return re.sub(r"https://", "\nhttps://", req.text)
+    except requests.exceptions.RequestException as e:
         return f"Upload failed!❌\n\n```{e}```"
 
 
@@ -264,7 +265,7 @@ def ytdl_normal_download(bot_msg, client, url):
         bot_msg.edit_text('Download complete. Sending now...')
         for video_path in video_paths:
             # normally there's only one video in that path...
-            filename = pathlib.Path(video_path).name
+            filename = video_path.name
             remain = bot_text.remaining_quota_caption(chat_id)
             st_size = os.stat(video_path).st_size
             size = sizeof_fmt(st_size)
diff --git a/ytdlbot/utils.py b/ytdlbot/utils.py
index 3451a15..a9130e2 100644
--- a/ytdlbot/utils.py
+++ b/ytdlbot/utils.py
@@ -46,7 +46,7 @@ def get_user_settings(user_id: "str") -> "tuple":
     cur.execute("SELECT * FROM settings WHERE user_id = %s", (user_id,))
     data = cur.fetchone()
     if data is None:
-        return 100, "high", "video"
+        return 100, "high", "video", "Celery"
     return data
 
 
diff --git a/ytdlbot/ytdl_bot.py b/ytdlbot/ytdl_bot.py
index 5d8702d..b2d5761 100644
--- a/ytdlbot/ytdl_bot.py
+++ b/ytdlbot/ytdl_bot.py
@@ -199,6 +199,15 @@ def direct_handler(client: "Client", message: "types.Message"):
 def settings_handler(client: "Client", message: "types.Message"):
     chat_id = message.chat.id
     client.send_chat_action(chat_id, "typing")
+    data = get_user_settings(str(chat_id))
+    set_mode = (data[-1])
+    text = {"Local": "Celery", "Celery": "Local"}.get(set_mode, "Local")
+    mode_text = f"Download mode: **{set_mode}**"
+    if message.chat.username == OWNER:
+        extra = [InlineKeyboardButton(f"Change download mode to {text}", callback_data=text)]
+    else:
+        extra = []
+
     markup = InlineKeyboardMarkup(
         [
             [  # First row
@@ -211,11 +220,11 @@ def settings_handler(client: "Client", message: "types.Message"):
                 InlineKeyboardButton("Medium Quality", callback_data="medium"),
                 InlineKeyboardButton("Low Quality", callback_data="low"),
             ],
+            extra
         ]
     )
 
-    data = get_user_settings(str(chat_id))
-    client.send_message(chat_id, bot_text.settings.format(data[1], data[2]), reply_markup=markup)
+    client.send_message(chat_id, bot_text.settings.format(data[1], data[2]) + mode_text, reply_markup=markup)
 
 
 @app.on_message(filters.command(["vip"]))
@@ -303,6 +312,13 @@ def audio_callback(client: "Client", callback_query: types.CallbackQuery):
     audio_entrance(msg, client)
 
 
+@app.on_callback_query(filters.regex(r"Local|Celery"))
+def owner_local_callback(client: "Client", callback_query: types.CallbackQuery):
+    chat_id = callback_query.message.chat.id
+    set_user_settings(chat_id, "mode", callback_query.data)
+    callback_query.answer(f"Download mode was changed to {callback_query.data}")
+
+
 def periodic_sub_check():
     vip = VIP()
     for cid, uids in vip.group_subscriber().items():