From fe22b1cbb31ae4ea4702e1bc61ecf5775a2f811c Mon Sep 17 00:00:00 2001
From: Hafitz Setya <71178188+breakdowns@users.noreply.github.com>
Date: Wed, 17 Mar 2021 11:11:11 +0700
Subject: [PATCH] v4.2 (#27)

- Added Extract Password Protected Archive Files
- Many minor bug fixed
---
 Dockerfile            |   3 +-
 README.md             |   2 +
 bot/modules/mirror.py |  22 ++++-
 bot/modules/watch.py  |   4 +-
 pextract              | 195 ++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 219 insertions(+), 7 deletions(-)
 create mode 100644 pextract

diff --git a/Dockerfile b/Dockerfile
index 626df92..578b229 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -13,7 +13,8 @@ RUN apt-get -qq update && \
 
 COPY requirements.txt .
 COPY extract /usr/local/bin
-RUN chmod +x /usr/local/bin/extract
+COPY pextract /usr/local/bin
+RUN chmod +x /usr/local/bin/extract && chmod +x /usr/local/bin/pextract
 RUN pip3 install --no-cache-dir -r requirements.txt
 RUN locale-gen en_US.UTF-8
 ENV LANG en_US.UTF-8
diff --git a/README.md b/README.md
index 565ee53..6d1eff5 100644
--- a/README.md
+++ b/README.md
@@ -52,6 +52,8 @@ Fork this repo, than upload credentials.json and token.pickle to your forks
 - Custom Buttons
 - Custom Filename (Only for url, telegram files and ytdl. Not for mega links and magnet/torrents)
 - Speedtest with picture results
+- Extracting password protected files and using custom filename see these examples:-
+> https://telegra.ph/Magneto-Python-Aria---Custom-Filename-Examples-01-20
 - Extract these filetypes and uploads to google drive
 > ZIP, RAR, TAR, 7z, ISO, WIM, CAB, GZIP, BZIP2, 
 > APM, ARJ, CHM, CPIO, CramFS, DEB, DMG, FAT, 
diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py
index 5d27e88..c30ba7f 100644
--- a/bot/modules/mirror.py
+++ b/bot/modules/mirror.py
@@ -24,16 +24,18 @@ import pathlib
 import os
 import subprocess
 import threading
+import re
 
 ariaDlManager = AriaDownloadHelper()
 ariaDlManager.start_listener()
 
 class MirrorListener(listeners.MirrorListeners):
-    def __init__(self, bot, update, isTar=False, tag=None, extract=False):
+    def __init__(self, bot, update, pswd, isTar=False, tag=None, extract=False):
         super().__init__(bot, update)
         self.isTar = isTar
         self.tag = tag
         self.extract = extract
+        self.pswd = pswd
 
     def onDownloadStarted(self):
         pass
@@ -76,7 +78,11 @@ class MirrorListener(listeners.MirrorListeners):
                 )
                 with download_dict_lock:
                     download_dict[self.uid] = ExtractStatus(name, m_path, size)
-                archive_result = subprocess.run(["extract", m_path])
+                pswd = self.pswd
+                if pswd is not None:
+                    archive_result = subprocess.run(["pextract", m_path, pswd])
+                else:
+                    archive_result = subprocess.run(["extract", m_path])
                 if archive_result.returncode == 0:
                     threading.Thread(target=os.remove, args=(m_path,)).start()
                     LOGGER.info(f"Deleting archive : {m_path}")
@@ -200,12 +206,21 @@ def _mirror(bot, update, isTar=False, extract=False):
     name_args = update.message.text.split('|')
     try:
         link = message_args[1]
+        if link.startswith("|") or link.startswith("pswd: "):
+            link = ''
     except IndexError:
         link = ''
     try:
         name = name_args[1]
+        name = name.strip()
+        if name.startswith("pswd: "):
+            name = ''
     except IndexError:
         name = ''
+    pswd = re.search('(?<=pswd: )(.*)', update.message.text)
+    if pswd is not None:
+      pswd = pswd.groups()
+      pswd = " ".join(pswd)
     LOGGER.info(link)
     link = link.strip()
     reply_to = update.message.reply_to_message
@@ -240,8 +255,7 @@ def _mirror(bot, update, isTar=False, extract=False):
         link = direct_link_generator(link)
     except DirectDownloadLinkException as e:
         LOGGER.info(f'{link}: {e}')
-
-    listener = MirrorListener(bot, update, isTar, tag, extract)
+    listener = MirrorListener(bot, update, pswd, isTar, tag, extract)
     if bot_utils.is_mega_link(link):
         link_type = get_mega_link_type(link)
         if link_type == "folder" and BLOCK_MEGA_FOLDER:
diff --git a/bot/modules/watch.py b/bot/modules/watch.py
index 93edcd4..64338a7 100644
--- a/bot/modules/watch.py
+++ b/bot/modules/watch.py
@@ -45,8 +45,8 @@ def _watch(bot: Bot, update, isTar=False):
         tag = reply_to.from_user.username
     else:
         tag = None
-
-    listener = MirrorListener(bot, update, isTar, tag)
+    pswd = ""
+    listener = MirrorListener(bot, update, pswd, isTar, tag)
     ydl = YoutubeDLHelper(listener)
     threading.Thread(target=ydl.add_download,args=(link, f'{DOWNLOAD_DIR}{listener.uid}', qual, name)).start()
     sendStatusMessage(update, bot)
diff --git a/pextract b/pextract
new file mode 100644
index 0000000..4394968
--- /dev/null
+++ b/pextract
@@ -0,0 +1,195 @@
+#!/bin/bash
+
+if [ $# -lt 1 ]; then
+  echo "Usage: $(basename $0) FILES"
+  exit 1
+fi
+
+extract() {
+  arg="$1"
+  pswd="$2"
+  cd "$(dirname "$arg")" || exit
+  case "$arg" in
+  *.tar.bz2)
+    tar xjf "$arg" --one-top-level
+    local code=$?
+    ;;
+  *.tar.gz)
+    tar xzf "$arg" --one-top-level
+    local code=$?
+    ;;
+  *.bz2)
+    bunzip2 "$arg"
+    local code=$?
+    ;;
+  *.gz)
+    gunzip "$arg"
+    local code=$?
+    ;;
+  *.tar)
+    tar xf "$arg" --one-top-level
+    local code=$?
+    ;;
+  *.tbz2)
+    (tar xjf "$arg" --one-top-level)
+    local code=$?
+    ;;
+  *.tgz)
+    tar xzf "$arg" --one-top-level
+    local code=$?
+    ;;
+  *.zip)
+    a_dir=$(expr "$arg" : '\(.*\).zip')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.7z)
+    a_dir=$(expr "$arg" : '\(.*\).7z')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.Z)
+    uncompress "$arg"
+    local code=$?
+    ;;
+  *.rar)
+    a_dir=$(expr "$arg" : '\(.*\).rar')
+    mkdir "$a_dir"
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.iso)
+    a_dir=$(expr "$arg" : '\(.*\).iso')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.wim)
+    a_dir=$(expr "$arg" : '\(.*\).wim')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.cab)
+    a_dir=$(expr "$arg" : '\(.*\).cab')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.apm)
+    a_dir=$(expr "$arg" : '\(.*\).apm')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.arj)
+    a_dir=$(expr "$arg" : '\(.*\).arj')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.chm)
+    a_dir=$(expr "$arg" : '\(.*\).chm')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.cpio)
+    a_dir=$(expr "$arg" : '\(.*\).cpio')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.cramfs)
+    a_dir=$(expr "$arg" : '\(.*\).cramfs')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.deb)
+    a_dir=$(expr "$arg" : '\(.*\).deb')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.dmg)
+    a_dir=$(expr "$arg" : '\(.*\).dmg')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.fat)
+    a_dir=$(expr "$arg" : '\(.*\).fat')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.hfs)
+    a_dir=$(expr "$arg" : '\(.*\).hfs')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.lzh)
+    a_dir=$(expr "$arg" : '\(.*\).lzh')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.lzma)
+    a_dir=$(expr "$arg" : '\(.*\).lzma')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.lzma2)
+    a_dir=$(expr "$arg" : '\(.*\).lzma2')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.mbr)
+    a_dir=$(expr "$arg" : '\(.*\).mbr')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.msi)
+    a_dir=$(expr "$arg" : '\(.*\).msi')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.mslz)
+    a_dir=$(expr "$arg" : '\(.*\).mslz')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.nsis)
+    a_dir=$(expr "$arg" : '\(.*\).nsis')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.ntfs)
+    a_dir=$(expr "$arg" : '\(.*\).ntfs')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.rpm)
+    a_dir=$(expr "$arg" : '\(.*\).rpm')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.squashfs)
+    a_dir=$(expr "$arg" : '\(.*\).squashfs')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.udf)
+    a_dir=$(expr "$arg" : '\(.*\).udf')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.vhd)
+    a_dir=$(expr "$arg" : '\(.*\).vhd')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *.xar)
+    a_dir=$(expr "$arg" : '\(.*\).xar')
+    7z x "$arg" -o"$a_dir" -p"$pswd"
+    local code=$?
+    ;;
+  *)
+    echo "'$arg' cannot be extracted via extract()" 1>&2
+    exit 1
+    ;;
+  esac
+  cd - || exit $?
+  exit $code
+}
+
+extract "$1" "$2"