{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c_h4ZT-aEU_b"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive/')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "GlgvVJA7NXeP"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import shutil\n",
        "import io\n",
        "import hashlib\n",
        "from pathlib import Path\n",
        "from PIL import Image\n",
        "!pip install pynacl\n",
        "from nacl.secret import SecretBox\n",
        "from google.colab import files\n",
        "import zipfile\n",
        "from google.colab import drive\n",
        "#@markdown 🔑 Decrypt Results (either by widget uploading zip or by manual filepath to saved zip on drive)\n",
        "print(\"🔓 Decryption Notebook Ready\")\n",
        "\n",
        "# ------------------- INPUTS -------------------\n",
        "password = 'banana' # @param {type:\"string\"}\n",
        "use_manual_filepath = False  # @param {type:\"boolean\"}\n",
        "manual_filepath = '/content/drive/MyDrive/encrypted.zip' # @param {type:\"string\"}\n",
        "\n",
        "# ------------------- FILE INPUT LOGIC -------------------\n",
        "zip_filename = None\n",
        "\n",
        "if use_manual_filepath:\n",
        "    if not os.path.exists('/content/drive'):\n",
        "        print(\"🔄 Mounting Google Drive...\")\n",
        "        drive.mount('/content/drive')\n",
        "\n",
        "    if not manual_filepath:\n",
        "        raise ValueError(\"❌ Manual filepath cannot be empty when 'Use manual filepath' is enabled!\")\n",
        "\n",
        "    if not os.path.exists(manual_filepath):\n",
        "        raise FileNotFoundError(f\"❌ File not found at manual filepath: {manual_filepath}\")\n",
        "\n",
        "    zip_filename = manual_filepath\n",
        "    print(f\"✅ Using manual filepath: {zip_filename}\")\n",
        "else:\n",
        "    print(\"\\n📤 Upload your final_encrypted_outputs.zip file:\")\n",
        "    uploaded = files.upload()\n",
        "\n",
        "    for fn in uploaded.keys():\n",
        "        if fn.endswith('.zip'):\n",
        "            zip_filename = fn\n",
        "            break\n",
        "\n",
        "    if not zip_filename:\n",
        "        raise ValueError(\"❌ No .zip file was uploaded!\")\n",
        "\n",
        "    print(f\"✅ Uploaded: {zip_filename}\")\n",
        "\n",
        "# ------------------- DECRYPTION FUNCTION -------------------\n",
        "def decrypt_data(encrypted_data: bytes, box: SecretBox) -> bytes:\n",
        "    try:\n",
        "        return box.decrypt(encrypted_data)\n",
        "    except Exception as e:\n",
        "        raise ValueError(f\"Decryption failed (wrong password?): {e})\")\n",
        "\n",
        "# ------------------- MAIN DECRYPTION -------------------\n",
        "\n",
        "if not password:\n",
        "    raise ValueError(\"❌ Password cannot be empty!\")\n",
        "\n",
        "# Derive the same key as in the encryption code\n",
        "key = hashlib.sha256(password.encode('utf-8')).digest()\n",
        "box = SecretBox(key)\n",
        "\n",
        "# Create output folder\n",
        "output_dir = \"/content/decrypted_images\"\n",
        "os.makedirs(output_dir, exist_ok=True)\n",
        "\n",
        "print(f\"\\n📂 Extracting and decrypting to: {output_dir}\")\n",
        "\n",
        "# Extract zip\n",
        "with zipfile.ZipFile(zip_filename, 'r') as zip_ref:\n",
        "    zip_ref.extractall(\"/content/temp_extracted\")\n",
        "\n",
        "# Find all .enc files\n",
        "enc_files = list(Path(\"/content/temp_extracted\").rglob(\"*.enc\"))\n",
        "\n",
        "print(f\"🔍 Found {len(enc_files)} encrypted files\")\n",
        "\n",
        "success_count = 0\n",
        "\n",
        "for enc_path in enc_files:\n",
        "    try:\n",
        "        with open(enc_path, \"rb\") as f:\n",
        "            encrypted_bytes = f.read()\n",
        "\n",
        "        # Decrypt\n",
        "        decrypted_bytes = decrypt_data(encrypted_bytes, box)\n",
        "\n",
        "        # Convert to PIL Image and save as JPG\n",
        "        img = Image.open(io.BytesIO(decrypted_bytes)).convert(\"RGB\")\n",
        "\n",
        "        output_path = os.path.join(output_dir, enc_path.stem + \".jpg\")\n",
        "        img.save(output_path, format=\"JPEG\", quality=95)\n",
        "\n",
        "        success_count += 1\n",
        "        print(f\"✅ Decrypted: {enc_path.name} → {Path(output_path).name}\")\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"❌ Failed {enc_path.name}: {e}\")\n",
        "\n",
        "# Cleanup temp folder\n",
        "shutil.rmtree(\"/content/temp_extracted\", ignore_errors=True)\n",
        "\n",
        "print(f\"\\n🎉 Decryption complete! {success_count}/{len(enc_files)} files decrypted successfully.\")\n",
        "print(f\"📁 Decrypted images saved to: {output_dir}\")\n",
        "\n",
        "# Optional: Download all decrypted images as a new zip\n",
        "download_zip = True #input(\"\\nDo you want to download all decrypted images as a zip? (y/n): \").lower().strip() == 'y'\n",
        "\n",
        "if download_zip:\n",
        "    download_zip_path = \"/content/decrypted_images.zip\"\n",
        "    shutil.make_archive(\"/content/decrypted_images\", 'zip', output_dir)\n",
        "    files.download(download_zip_path)\n",
        "    print(\"📥 Download started...\")"
      ]
    }
  ]
}