CADASTRO

import customtkinter as ctk
from tkinter import filedialog, messagebox
import threading
import os
import zipfile
import time

# =========================
# CONFIG VISUAL
# =========================
ctk.set_appearance_mode("dark")
ctk.set_default_color_theme("blue")

# =========================
# ASSINATURAS
# =========================
SIG_JPG = b"\xFF\xD8"
END_JPG = b"\xFF\xD9"

SIG_PNG = b"\x89PNG\r\n\x1a\n"
END_PNG = b"IEND\xaeB`\x82"

SIG_PDF = b"%PDF-"
END_PDF = b"%%EOF"

SIG_ZIP = b"PK\x03\x04"
SIG_OLE = b"\xD0\xCF\x11\xE0\xA1\xB1\x1A\xE1"

CHUNK = 1024 * 1024
OVERLAP = 128 * 1024
MAX_OFFICE = 100 * 1024 * 1024  # 100 MB

# =========================
# MOTOR DE RECUPERAÇÃO
# =========================
def identify_zip(path):
    try:
        with zipfile.ZipFile(path, "r") as z:
            names = z.namelist()
            if "word/document.xml" in names:
                return "docx"
            if "xl/workbook.xml" in names:
                return "xlsx"
    except Exception:
        pass
    return "zip"


def recover(src, out, log, progress_cb=None):
    out = os.path.abspath(out)
    os.makedirs(out, exist_ok=True)

    size_total = os.path.getsize(src)
    processed = 0

    tail = b""
    active = None
    counters = {}

    with open(src, "rb") as f:
        while True:
            chunk = f.read(CHUNK)
            if not chunk:
                break

            processed += len(chunk)
            if progress_cb:
                progress_cb(processed, size_total)

            data = tail + chunk

            # continua escrita se houver arquivo ativo
            if active:
                active["fh"].write(chunk)
                active["size"] += len(chunk)

                ext = active["ext"]
                if ext == "jpg" and END_JPG in data:
                    active["fh"].close()
                    log("✅ JPG recuperado")
                    active = None
                elif ext == "png" and END_PNG in data:
                    active["fh"].close()
                    log("✅ PNG recuperado")
                    active = None
                elif ext == "pdf" and END_PDF in data:
                    active["fh"].close()
                    log("✅ PDF recuperado")
                    active = None
                elif ext in ("zip", "doc") and active["size"] > MAX_OFFICE:
                    active["fh"].close()
                    log("⚠️ Office fechado por limite de tamanho")
                    active = None

            # se não houver ativo, procurar novas assinaturas
            if not active:
                found = []
                for sig, ext in [
                    (SIG_JPG, "jpg"),
                    (SIG_PNG, "png"),
                    (SIG_PDF, "pdf"),
                    (SIG_ZIP, "zip"),
                    (SIG_OLE, "doc"),
                ]:
                    p = data.find(sig)
                    if p != -1:
                        found.append((ext, p))

                if found:
                    ext, pos = min(found, key=lambda x: x[1])
                    counters.setdefault(ext, 0)
                    counters[ext] += 1
                    name = os.path.join(out, f"rec_{ext}_{counters[ext]:06d}.{ext}")
                    fh = open(name, "wb")
                    fh.write(data[pos:])
                    active = {"ext": ext, "fh": fh, "path": name, "size": len(data) - pos}
                    log(f"📄 Iniciando {ext.upper()}")

            tail = data[-OVERLAP:]

    # renomear ZIPs Office
    for z in os.listdir(out):
        if z.endswith(".zip"):
            zp = os.path.join(out, z)
            t = identify_zip(zp)
            if t in ("docx", "xlsx"):
                os.rename(zp, zp.replace(".zip", f".{t}"))

    log("🎉 Recuperação finalizada")


# =========================
# INTERFACE GRÁFICA
# =========================
class Conexus(ctk.CTk):
    def __init__(self):
        super().__init__()
        self.title("Conexus Recovery")
        self.geometry("720x560")
        self.resizable(False, False)

        self.src = ctk.StringVar()
        self.out = ctk.StringVar()

        self.ui()

    def ui(self):
        ctk.CTkLabel(self, text="Conexus Recovery", font=("Segoe UI", 26, "bold")).pack(pady=10)
        ctk.CTkLabel(self, text="Recuperação avançada de arquivos apagados").pack(pady=5)

        box = ctk.CTkFrame(self)
        box.pack(padx=20, pady=10, fill="x")

        ctk.CTkLabel(box, text="Imagem do disco (.dd / .img)").pack(anchor="w", padx=10)
        f1 = ctk.CTkFrame(box, fg_color="transparent")
        f1.pack(fill="x", padx=10, pady=5)
        ctk.CTkEntry(f1, textvariable=self.src).pack(side="left", fill="x", expand=True)
        ctk.CTkButton(f1, text="Procurar", command=self.pick_src).pack(side="right")

        ctk.CTkLabel(box, text="Pasta de saída (outro disco)").pack(anchor="w", padx=10)
        f2 = ctk.CTkFrame(box, fg_color="transparent")
        f2.pack(fill="x", padx=10, pady=5)
        ctk.CTkEntry(f2, textvariable=self.out).pack(side="left", fill="x", expand=True)
        ctk.CTkButton(f2, text="Procurar", command=self.pick_out).pack(side="right")

        self.btn = ctk.CTkButton(self, text="▶ Iniciar Recuperação", height=40, command=self.start)
        self.btn.pack(pady=10)

        self.progress = ctk.CTkProgressBar(self)
        self.progress.pack(padx=20, fill="x")
        self.progress.set(0)

        self.logbox = ctk.CTkTextbox(self, height=240)
        self.logbox.pack(padx=20, pady=10, fill="both", expand=True)
        self.logbox.insert("end", "🟢 Pronto para iniciar...\n")

    def log(self, msg):
        self.after(0, lambda: (
            self.logbox.insert("end", msg + "\n"),
            self.logbox.see("end")
        ))

    def update_progress(self, done, total):
        self.after(0, lambda: self.progress.set(done / total))

    def pick_src(self):
        p = filedialog.askopenfilename(filetypes=[("Imagem de Disco", "*.dd *.img")])
        if p:
            self.src.set(p)

    def pick_out(self):
        p = filedialog.askdirectory()
        if p:
            self.out.set(p)

    def start(self):
        if not self.src.get() or not self.out.get():
            messagebox.showerror("Erro", "Selecione origem e destino")
            return

        self.btn.configure(state="disabled")
        self.log("🚀 Iniciando recuperação...")

        threading.Thread(
            target=recover,
            args=(self.src.get(), self.out.get(), self.log, self.update_progress),
            daemon=True
        ).start()


# =========================
# START
# =========================
if __name__ == "__main__":
    Conexus().mainloop()