mtz_extractor_v2_1.py

 avatar
unknown
python
10 months ago
9.2 kB
19
Indexable
import os
import sys
import zipfile
import shutil
import logging
import time
import platform
import psutil 
import threading
import random
from typing import Set, Optional
from pathlib import Path
import contextlib
from datetime import datetime


class ColorText:
    HEADER = "\033[95m"
    BLUE = "\033[94m"
    CYAN = "\033[96m"
    GREEN = "\033[92m"
    YELLOW = "\033[93m"
    RED = "\033[91m"
    END = "\033[0m"
    BOLD = "\033[1m"
    UNDERLINE = "\033[4m"

    @staticmethod
    def blue(text): return f"{ColorText.BLUE}{text}{ColorText.END}"
    @staticmethod
    def green(text): return f"{ColorText.GREEN}{text}{ColorText.END}"
    @staticmethod
    def red(text): return f"{ColorText.RED}{text}{ColorText.END}"
    @staticmethod
    def yellow(text): return f"{ColorText.YELLOW}{text}{ColorText.END}"
    @staticmethod
    def cyan(text): return f"{ColorText.CYAN}{text}{ColorText.END}"
    @staticmethod
    def bold(text): return f"{ColorText.BOLD}{text}{ColorText.END}"


class LoadingAnimation:
    def __init__(self, description="Processing"):
        self.description = description
        self.is_running = False
        self.animation_thread = None
        self.animation_chars = ["◜", "◠", "◝", "◞", "◡", "◟"]
        self.current_color = random.choice(
            [ColorText.BLUE, ColorText.CYAN, ColorText.GREEN, ColorText.YELLOW])

    def start(self):
        self.is_running = True
        self.animation_thread = threading.Thread(target=self._animate, daemon=True)
        self.animation_thread.start()

    def stop(self):
        self.is_running = False
        if self.animation_thread:
            self.animation_thread.join()
        sys.stdout.write("\r" + " " * (len(self.description) + 20) + "\r")
        sys.stdout.flush()

    def _animate(self):
        while self.is_running:
            for frame in self.animation_chars:
                if not self.is_running:
                    break
                sys.stdout.write(f"\r{self.current_color}{frame}{ColorText.END} {self.description} ")
                sys.stdout.flush()
                time.sleep(0.1)


@contextlib.contextmanager
def loading_animation(description="Processing"):
    spinner = LoadingAnimation(description)
    spinner.start()
    try:
        yield spinner
    finally:
        spinner.stop()


class MTZExtractor:
    def __init__(self, allowed_extensions: Set[str] = None):
        self.allowed_extensions = allowed_extensions or {
            ".java", ".kt", ".so", ".aar", ".jar", ".mp3", ".wav", ".mp4", ".3gp",
            ".txt", ".json", ".xml", ".html", ".css", ".js", ".ttf", ".otf", ".png",
            ".jpg", ".jpeg", ".gif", ".webp", ".pdf", ".gradle", ".properties", ".MF", ".SF", ".RSA"
        }
        self.setup_logging()
        self.stats = {"start_time": None, "total_files": 0, "total_size": 0, "extracted_size": 0}

    def setup_logging(self):
        log_folder = Path("logs")
        log_folder.mkdir(exist_ok=True)
        logging.basicConfig(level=logging.INFO)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        log_file = log_folder / f"compression_{timestamp}.log"
        self.logger = logging.getLogger(f"MTZExtractor_{timestamp}")
        file_handler = logging.FileHandler(log_file, mode="w", encoding="utf-8")
        formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
        file_handler.setFormatter(formatter)
        self.logger.addHandler(file_handler)
        self.logger.propagate = False
        self.logger.info("MTZ Extractor initialized")

    def print_banner(self):
        banner = (
            f"\n"
            f"{ColorText.cyan('╔═════════════════════════════════════════╗')}\n"
            f"{ColorText.cyan('║')}        {ColorText.bold(' MTZ Extractor v2.1   ')}           {ColorText.cyan('║')}\n"
            f"{ColorText.cyan('║')}      {ColorText.yellow('     BY Ryhshall & Dsertécn    ')}    {ColorText.cyan('║')}\n"
            f"{ColorText.cyan('╚═════════════════════════════════════════╝')}\n"
        )
        print(banner)

    def validate_mtz_file(self, file_path):
        if not os.path.exists(file_path):
            print(f"\n{ColorText.red('❌ File not found!')}")
            return False
        if not file_path.endswith(".mtz"):
            print(f"\n{ColorText.red('❌ Not an MTZ file!')}")
            return False
        return True

    def create_extract_folder(self, file_path):
        try:
            base = Path("./extracted")
            name = Path(file_path).stem
            folder = base / name
            counter = 1
            while folder.exists():
                folder = base / f"{name}_copy{counter}"
                counter += 1
            folder.mkdir(parents=True, exist_ok=True)
            return str(folder)
        except Exception as e:
            print(f"\n{ColorText.red(f'❌ Failed to create folder: {str(e)}')}")
            return None

    def extract_mtz(self, file_path, extract_folder):
        try:
            self.stats["start_time"] = time.time()
            self.stats["total_size"] = os.path.getsize(file_path)
            with loading_animation(f"Extracting {ColorText.yellow(os.path.basename(file_path))}"):
                with zipfile.ZipFile(file_path, "r") as zip_ref:
                    zip_ref.extractall(extract_folder)
                    self.stats["total_files"] = len(zip_ref.namelist())
            return True
        except Exception as e:
            print(f"\n{ColorText.red(f'❌ Extraction failed: {str(e)}')}")
            return False

    def process_files(self, folder):
        with loading_animation(f"Processing {ColorText.yellow(os.path.basename(folder))}"):
            for file in Path(folder).rglob('*'):
                if file.is_file() and file.suffix not in self.allowed_extensions:
                    new = file.with_suffix(file.suffix + '.zip')
                    file.rename(new)
            for file in Path(folder).rglob('*.zip'):
                if file.is_file():
                    out_dir = file.with_suffix('')
                    out_dir.mkdir(exist_ok=True)
                    try:
                        with zipfile.ZipFile(file, 'r') as zf:
                            zf.extractall(out_dir)
                        file.unlink()
                    except Exception:
                        continue

    def show_completion(self, folder):
        elapsed = time.time() - self.stats["start_time"]
        print(f"\n{ColorText.green('✅ Extraction complete!')}")
        print(f"📁 Location: {ColorText.yellow(folder)}")
        print(f"⏱️ Time: {elapsed:.2f}s\n")

    def convert_to_bak(self, file_path: str) -> Optional[str]:
        try:
            file_name = Path(file_path).stem
            bak_folder = Path("./backup")
            bak_folder.mkdir(exist_ok=True)
            bak_file = bak_folder / f"{file_name}.bak"
            with loading_animation(f"Converting {ColorText.yellow(os.path.basename(file_path))} to BAK"):
                shutil.make_archive(str(bak_file.with_suffix('')), 'zip', root_dir=Path(file_path).parent, base_dir=Path(file_path).name)
                zip_path = bak_file.with_suffix('.zip')
                if zip_path.exists():
                    zip_path.rename(bak_file)
            print(f"\n{ColorText.green('✅ File successfully converted to BAK!')}")
            print(f"📦 Location: {ColorText.yellow(bak_file)}\n")
            return str(bak_file)
        except Exception as e:
            print(f"\n{ColorText.red(f'❌ Conversion failed: {str(e)}')}")
            return None


def get_user_input():
    return input(f"{ColorText.cyan('📂')} Enter MTZ file location: ").strip()


def main():
    os.system("cls" if os.name == "nt" else "clear")
    extractor = MTZExtractor()
    extractor.print_banner()
    try:
        file_path = get_user_input()
        if not extractor.validate_mtz_file(file_path):
            sys.exit(1)

        print(f"\n{ColorText.cyan('🌀 Choose an action:')}")
        print(f"{ColorText.yellow('[1]')} Extract MTZ")
        print(f"{ColorText.yellow('[2]')} Convert to BAK\n")
        choice = input(f"{ColorText.cyan('👉 Select (1/2): ')}").strip()

        if choice == "1":
            extract_folder = extractor.create_extract_folder(file_path)
            if not extract_folder:
                sys.exit(1)
            print(f"\n{ColorText.cyan('⏳ Starting extraction...')}\n")
            if not extractor.extract_mtz(file_path, extract_folder):
                sys.exit(1)
            extractor.process_files(extract_folder)
            extractor.show_completion(extract_folder)

        elif choice == "2":
            extractor.convert_to_bak(file_path)
        else:
            print(f"{ColorText.red('❌ Invalid choice!')}")
    except KeyboardInterrupt:
        print(f"\n{ColorText.yellow('⚠️ Cancelled!')}\n")
        sys.exit(0)
    except Exception as e:
        print(f"\n{ColorText.red('❌ Error:')} {str(e)}\n")
        sys.exit(1)


if __name__ == "__main__":
    main()
Editor is loading...
Leave a Comment