zipasta/zipasta.py

"""
Copyright © 2022 Clóvis Fabrício Costa - All Rights Reserved

This file is part of zipasta.

zipasta is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public
License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later
version.

zipasta is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty
of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License along with zipasta. If not,
see https://www.gnu.org/licenses/
"""

import base64
import itertools
import random
import shutil
import string
import sys
import zlib
import pathlib

import tkinter as tk
from tkinter import filedialog as fdlg


MAX_SIZE = 230

_order_digits = tuple(itertools.chain(string.digits, string.ascii_uppercase, string.ascii_lowercase))


def _random_chars(number=1):
    return ''.join(random.choice(_order_digits) for _ign in range(number))


class App:
    def __init__(self, parent, default_path: pathlib.Path = None):
        self.parent = parent
        self.vars = {}
        self.labels = {}
        self.buttons = {}
        self.arquivos_prefixos = {}
        self.path = default_path
        self._build_gui()
        self.btn_caminho()
        if self.path is None:
            sys.exit(1)
        self._read_folder()

    def _get_or_create_var(self, varname):
        if varname not in self.vars:
            self.vars[varname] = tk.StringVar()
        return self.vars[varname]

    def __setitem__(self, key, value):
        self._get_or_create_var(key).set(str(value))

    def __getitem__(self, varname):
        return self.vars[varname].get()

    def _make_label(self, parent, name, initial_text='', **kwds):
        _var = self._get_or_create_var(name)
        _lbl = self.labels[name] = tk.Label(parent, textvariable=_var, **kwds)
        _var.set(initial_text)
        return _lbl

    def _make_button(self, parent, name, initial_text='', **kwds):
        def _command():
            getattr(self, f'btn_{name}')(**kwds)
        _btn = self.buttons[name] = tk.Button(parent, text=initial_text, command=_command)
        return _btn

    @staticmethod
    def _make_frame(parent, main_row=0, main_col=0):
        _frame = tk.Frame(parent)
        _frame.rowconfigure(main_row, weight=1)
        _frame.columnconfigure(main_col, weight=1)
        return _frame

    def _build_gui(self):
        main_frame = self._make_frame(self.parent, main_row=1)

        self._build_frame_caminho(main_frame).grid(row=0, column=0)
        self._build_frame_arquivos(main_frame).grid(row=1, column=0)

        main_frame.grid(row=0, column=0)
        return main_frame

    def _build_frame_caminho(self, parent):
        _frame = self._make_frame(parent)
        self._make_label(_frame, 'caminho').grid(row=0, column=0)
        if self.path is not None:
            self['caminho'] = self.path
        self._make_button(_frame, 'caminho', '…').grid(row=0, column=1)
        return _frame

    def _build_frame_arquivos(self, parent):
        _frame = self._make_frame(parent, main_row=1)
        self._make_label(_frame, 'arquivos', 'Arquivos:').grid(row=0, column=0, columnspan=3)
        self._arquivos = tk.Listbox(_frame)
        self._arquivos.grid(row=1, column=0, columnspan=3)
        self._make_button(_frame, 'adicionar', '➕').grid(row=2, column=0)
        self._make_button(_frame, 'extrair', '⤋').grid(row=2, column=1)
        self._make_button(_frame, 'deletar', '🗑').grid(row=2, column=2)
        return _frame

    def btn_caminho(self):
        _path = fdlg.askdirectory(parent=self.parent, title='Pasta destino', mustexist=True)
        if _path:
            _path = pathlib.Path(_path).absolute()
            if _path.is_dir():
                self['caminho'] = _path
                self.path = _path
                self._read_folder()
                return

    def _read_folder(self):
        self._arquivos.delete(0, tk.END)
        _mainpath = self.path / '_z'
        if _mainpath.is_dir():
            for filename in _mainpath.iterdir():
                if filename.is_dir() and filename.name != '_':
                    self.arquivos_prefixos[filename.name] = (filename, next(filename.iterdir()).name)
                    self._arquivos.insert(tk.END, filename.name)

    def btn_adicionar(self):
        _paths = fdlg.askopenfilenames(parent=self.parent, title='Selecione arquivos a codificar')
        pastas_criar = []
        for _path in _paths:
            _path = pathlib.Path(_path)
            with _path.open('rb') as f:
                data = f.read()
                original_size = len(data)
                data = zlib.compress(data, level=8)
                data = base64.urlsafe_b64encode(data)
            _mainpath = self.path / '_z'
            while True:
                prefix = _random_chars(3)
                _datapath = _mainpath / '_' / prefix
                if not _datapath.exists():
                    break
            # _datapath.mkdir(parents=True, exist_ok=False)
            print(f'Salvando {_path.name}({original_size // 1024}k) como {prefix}')

            pastas_criar.append(_datapath)
            pastas_criar.append(_mainpath / str(_path.name) / prefix)
            tamanho_restante = MAX_SIZE - (len(str(_datapath.absolute())) + 1)
            digitos_usar = 1
            while True:
                bytes_por_pasta = tamanho_restante - (digitos_usar + (digitos_usar // 2) + 1)
                num_pastas = len(data) // bytes_por_pasta
                if num_pastas < len(_order_digits) ** digitos_usar:
                    break
                digitos_usar += 1
            print(f'{digitos_usar} digitos para {num_pastas} pastas')
            for pos, prefixos in enumerate(itertools.product(_order_digits, repeat=digitos_usar)):
                pedaco = data[pos * bytes_por_pasta:(pos + 1) * bytes_por_pasta].decode("ascii")
                if not pedaco:
                    break
                prefixos = list(prefixos)
                pasta_criar = _datapath
                while len(prefixos) > 2:
                    pasta_criar = pasta_criar / ''.join(prefixos[:2])
                    prefixos = prefixos[2:]
                nome_pasta = f'{"".join(prefixos)}.{pedaco}'
                pasta_criar = pasta_criar / nome_pasta
                pastas_criar.append(pasta_criar)
        for pasta in pastas_criar:
            pasta.mkdir(parents=True)
        self._read_folder()

    def btn_deletar(self):
        _datapath = self.path / '_z' / '_'
        for i in self._arquivos.curselection():
            pastas_deletar = []
            _filename = self._arquivos.get(i)
            _path, _prefix = self.arquivos_prefixos[_filename]
            pastas_deletar.append(_path / _prefix)
            pastas_deletar.append(_path)
            for pasta in (_datapath / _prefix).rglob('*'):
                pastas_deletar.append(pasta)
            print(f'Deletando {len(pastas_deletar)} pastas...')
            # pastas_deletar.sort(key=lambda x: len(str(x)), reverse=True)
            # for pasta in pastas_deletar:
            #     pasta.rmdir()
            shutil.rmtree(_datapath / _prefix)
            shutil.rmtree(_path)
        self._read_folder()

    def btn_extrair(self):
        # Traverse the tuple returned by
        # curselection method and print
        # corresponding value(s) in the listbox
        _sel = self._arquivos.curselection()
        if len(_sel) == 1:
            _filename = self._arquivos.get(_sel[0])
            _new_filename = fdlg.asksaveasfilename(parent=self.parent, title='Salvar como...', initialfile=_filename)
            if not _new_filename:
                return
            _new_filename = pathlib.Path(_new_filename).absolute()
            self._extract_file(_filename, _new_filename)
        else:
            pasta_salvar = fdlg.askdirectory(parent=self.parent, mustexist=True, title='Pasta a salvar arquivos')
            if not pasta_salvar:
                return
            pasta_salvar = pathlib.Path(pasta_salvar)
            for i in self._arquivos.curselection():
                _filename = self._arquivos.get(i)
                _new_filename = pasta_salvar / _filename
                self._extract_file(_filename, _new_filename)

    def _extract_file(self, filename, new_filename):
        _fullname, prefix = self.arquivos_prefixos[filename]
        _datapath = self.path / '_z' / '_' / prefix
        data = bytearray()
        for folder in sorted(_datapath.rglob('*.*'), key=str):
            data += folder.name.rpartition('.')[-1].encode('ascii')
        data = base64.urlsafe_b64decode(data)
        data = zlib.decompress(data)
        with new_filename.open('wb') as f:
            f.write(data)


def main(argv=None):
    root = tk.Tk()
    root.title('Zipasta 0.1 © Clóvis Fabrício Costa')

    app = App(root)
    root.mainloop()


main(sys.argv)