diff --git a/covert_sd_card_tool.py b/covert_sd_card_tool.py
index 151ce76..42aa2fd 100755
--- a/covert_sd_card_tool.py
+++ b/covert_sd_card_tool.py
@@ -6,9 +6,6 @@ import sys
 import os
 import shutil
 from datetime import datetime
-import platform
-import urllib.request
-import tarfile
 import time
 import json
 
@@ -53,7 +50,7 @@ def run_command(command, shell=False, interactive=False):
         sys.exit(1)
 
 def check_dependencies():
-    dependencies = ["parted", "cryptsetup", "lsblk", "dd", "sgdisk", "wipefs", "bc", "fdisk"]
+    dependencies = ["parted", "cryptsetup", "lsblk", "dd", "sgdisk", "wipefs", "bc", "fdisk", "veracrypt", "lsof", "fuser", "mountpoint"]
     missing = []
     for dep in dependencies:
         if not shutil.which(dep):
@@ -107,7 +104,7 @@ def prepare_drive(drive):
         log(f"Processes using {drive}:\n{result.stdout}")
         kill = input(f"Do you want to kill these processes? (y/n) [Default: y]: ") or "y"
         if kill.lower() == "y":
-            run_command(f"sudo fuser -k {drive}", shell=True)
+            run_command(["sudo", "fuser", "-k", drive])
             log(f"Killed processes using {drive}.")
         else:
             log("Cannot proceed while processes are using the drive. Exiting.")
@@ -117,7 +114,7 @@ def prepare_drive(drive):
 
 def setup_usb():
     global DRIVE
-    log("Setting up bootable USB...")
+    log("Setting up bootable USB or preparing partitions...")
     list_drives()
     DRIVE = input("Enter the drive to use for USB (e.g., /dev/sda) [Default: /dev/sda]: ") or "/dev/sda"
 
@@ -127,60 +124,125 @@ def setup_usb():
         sys.exit(1)
 
     prepare_drive(DRIVE)
-    wipe = input(f"Do you want to wipe the drive {DRIVE} before starting? (y/n) [Default: n]: ") or "n"
-    if wipe.lower() == "y":
-        log(f"Wiping {DRIVE} and clearing any existing file system or encryption signatures...")
-        run_command(["sudo", "wipefs", "--all", DRIVE])
-        run_command(["sudo", "sgdisk", "--zap-all", DRIVE])
-        run_command(["sudo", "dd", "if=/dev/zero", f"of={DRIVE}", "bs=1M", "count=10"])
-        log(f"{DRIVE} wiped successfully.")
 
-    global KALI_ISO
-    if CREATE_KALI:
-        if not KALI_ISO:
-            KALI_ISO = input("Enter the path to the Kali ISO file: ")
-        if not os.path.isfile(KALI_ISO):
-            log(f"Error: Kali ISO file not found at {KALI_ISO}")
+    if CREATE_KALI or CREATE_TAILS or CREATE_DOCS:
+        wipe = input(f"Do you want to wipe the drive {DRIVE} before starting? (y/n) [Default: n]: ") or "n"
+        if wipe.lower() == "y":
+            log(f"Wiping {DRIVE} and clearing any existing file system or encryption signatures...")
+            run_command(["sudo", "wipefs", "--all", DRIVE])
+            run_command(["sudo", "sgdisk", "--zap-all", DRIVE])
+            run_command(["sudo", "dd", "if=/dev/zero", f"of={DRIVE}", "bs=1M", "count=10"])
+            log(f"{DRIVE} wiped successfully.")
+
+    if CREATE_KALI or CREATE_TAILS:
+        global KALI_ISO, TAILS_ISO
+        if CREATE_KALI:
+            if not KALI_ISO:
+                KALI_ISO = input("Enter the path to the Kali ISO file: ")
+            if not os.path.isfile(KALI_ISO):
+                log(f"Error: Kali ISO file not found at {KALI_ISO}")
+                sys.exit(1)
+            ISO_PATH = KALI_ISO
+        elif CREATE_TAILS:
+            if not TAILS_ISO:
+                TAILS_ISO = input("Enter the path to the Tails ISO file: ")
+            if not os.path.isfile(TAILS_ISO):
+                log(f"Error: Tails ISO file not found at {TAILS_ISO}")
+                sys.exit(1)
+            ISO_PATH = TAILS_ISO
+
+        log(f"Writing ISO to {DRIVE}...")
+        run_command(f"sudo dd if='{ISO_PATH}' of='{DRIVE}' bs=64M status=progress", shell=True, interactive=True)
+        log(f"ISO written to {DRIVE} successfully.")
+
+        if CREATE_KALI:
+            fix_partition_table()
+        elif CREATE_TAILS:
+            fix_partition_table_tails()
+    elif CREATE_DOCS:
+        fix_partition_table_docs_only()
+    else:
+        log("No valid setup option selected. Exiting.")
+        sys.exit(1)
+
+def fix_partition_table_docs_only():
+    log("Setting up partitions for documents only...")
+
+    # Clear existing partitions if any
+    run_command(f"sudo parted -a optimal -s {DRIVE} mklabel gpt", shell=True)
+    run_command(f"sudo partprobe {DRIVE}", shell=True)
+    run_command("sudo udevadm settle", shell=True)
+    time.sleep(5)  # Increased sleep to ensure partitions are recognized
+
+    # Get the total size of the drive in bytes
+    result = subprocess.run(["lsblk", "-b", "-n", "-o", "SIZE", DRIVE], capture_output=True, text=True)
+    total_size_bytes = int(result.stdout.strip())
+    total_size_mib = total_size_bytes / (1024 * 1024)  # Convert to MiB
+
+    # Ask for document partition size
+    size_docs = input("Enter size for documents partition in GB (leave blank to use remaining space minus 1GB): ")
+    start_docs_mib = 1  # Starting immediately after the first MiB
+    if size_docs:
+        try:
+            size_docs_gb = float(size_docs)
+            end_docs_mib = start_docs_mib + (size_docs_gb * 1024)
+            if end_docs_mib > (total_size_mib - 1024):
+                log("Error: Documents partition size exceeds available space when reserving 1GB for unencrypted partition.")
+                sys.exit(1)
+        except ValueError:
+            log("Invalid size entered for documents partition. Exiting.")
             sys.exit(1)
-        ISO_PATH = KALI_ISO
     else:
-        log("No OS selected for installation. Exiting.")
-        sys.exit(1)
+        end_docs_mib = total_size_mib - 1024  # Reserve 1GB for unencrypted partition
 
-    log(f"Writing ISO to {DRIVE}...")
-    run_command(f"sudo dd if='{ISO_PATH}' of='{DRIVE}' bs=64M status=progress", shell=True, interactive=True)
-    log(f"ISO written to {DRIVE} successfully.")
+    # Create documents partition
+    run_command(f"sudo parted -a optimal -s {DRIVE} mkpart primary {start_docs_mib}MiB {end_docs_mib}MiB", shell=True)
+    log("Created documents partition.")
 
-    if CREATE_KALI:
-        fix_partition_table()
-    else:
-        log("Only basic setup is performed. Exiting.")
-        sys.exit(1)
+    # Create unencrypted partition
+    start_unencrypted_mib = end_docs_mib
+    run_command(f"sudo parted -a optimal -s {DRIVE} mkpart primary {start_unencrypted_mib}MiB 100%", shell=True)
+    log("Created unencrypted partition for scripts/instructions.")
+
+    # Refresh partition table to recognize new partitions
+    run_command(f"sudo partprobe {DRIVE}", shell=True)
+    run_command("sudo udevadm settle", shell=True)
+    time.sleep(5)  # Increased sleep to ensure partitions are recognized
+
+    setup_unencrypted_partition()
+    setup_docs_partition()
 
 def fix_partition_table():
     log("Fixing partition table to reclaim remaining space...")
 
-    run_command(f"sudo parted -a optimal -s {DRIVE} rm 2", shell=True)
-    log("Deleted partition 2.")
+    # Attempt to delete partition 2 if it exists
+    try:
+        run_command(f"sudo parted -a optimal -s {DRIVE} rm 2", shell=True)
+        log("Deleted partition 2.")
+    except SystemExit:
+        log("No partition 2 to delete.")
 
-    result = subprocess.run(["sudo", "parted", "-s", DRIVE, "unit", "MB", "print"], capture_output=True, text=True)
+    # Get the end of partition 1
+    result = subprocess.run(["sudo", "parted", "-s", DRIVE, "unit", "MiB", "print"], capture_output=True, text=True)
     end_of_p1 = None
     for line in result.stdout.strip().splitlines():
         if line.strip().startswith("1"):
             parts = line.strip().split()
-            end_of_p1 = parts[2]
-            break
+            if len(parts) >= 3:
+                end_of_p1 = parts[2].replace('MiB', '')
+                break
     if end_of_p1 is None:
         log("Error: Could not find end of partition 1.")
         sys.exit(1)
 
-    log(f"End of partition 1: {end_of_p1}")
+    log(f"End of partition 1: {end_of_p1}MiB")
 
-    result = subprocess.run(["lsblk", "-bn", "-o", "SIZE", DRIVE], capture_output=True, text=True)
-    sizes = result.stdout.strip().splitlines()
-    total_size_bytes = int(sizes[0].strip())
-    total_size_mb = total_size_bytes / (1024 * 1024)
+    # Get the total size of the drive in bytes
+    result = subprocess.run(["lsblk", "-b", "-n", "-o", "SIZE", DRIVE], capture_output=True, text=True)
+    total_size_bytes = int(result.stdout.strip())
+    total_size_mib = total_size_bytes / (1024 * 1024)  # Convert to MiB
 
+    # Ask for persistence partition size
     size_persistence = input("Enter size for persistence partition in GB (e.g., 4): ") or "4"
     try:
         size_persistence_gb = float(size_persistence)
@@ -188,41 +250,45 @@ def fix_partition_table():
         log("Invalid size entered for persistence partition. Exiting.")
         sys.exit(1)
 
-    start_persistence_mb = float(end_of_p1.replace('MB', ''))
-    end_persistence_mb = start_persistence_mb + (size_persistence_gb * 1024)
+    start_persistence_mib = float(end_of_p1)
+    end_persistence_mib = start_persistence_mib + (size_persistence_gb * 1024)
 
-    if end_persistence_mb > total_size_mb:
+    if end_persistence_mib > total_size_mib:
         log("Error: Persistence partition size exceeds available space.")
         sys.exit(1)
 
-    run_command(f"sudo parted -a optimal -s {DRIVE} mkpart primary {start_persistence_mb}MB {end_persistence_mb}MB", shell=True)
+    # Create persistence partition
+    run_command(f"sudo parted -a optimal -s {DRIVE} mkpart primary {start_persistence_mib}MiB {end_persistence_mib}MiB", shell=True)
     log("Created persistence partition.")
 
-    start_docs_mb = end_persistence_mb
-
+    # Set up documents partition
+    start_docs_mib = end_persistence_mib
     size_docs = input("Enter size for documents partition in GB (leave blank to use remaining space minus 1GB): ")
     if size_docs:
         try:
             size_docs_gb = float(size_docs)
-            end_docs_mb = start_docs_mb + (size_docs_gb * 1024)
-            if end_docs_mb > total_size_mb - 1024:
+            end_docs_mib = start_docs_mib + (size_docs_gb * 1024)
+            if end_docs_mib > (total_size_mib - 1024):
                 log("Error: Documents partition size exceeds available space when reserving 1GB for unencrypted partition.")
                 sys.exit(1)
         except ValueError:
             log("Invalid size entered for documents partition. Exiting.")
             sys.exit(1)
     else:
-        end_docs_mb = total_size_mb - 1024
+        end_docs_mib = total_size_mib - 1024  # Reserve 1GB for unencrypted partition
 
-    run_command(f"sudo parted -a optimal -s {DRIVE} mkpart primary {start_docs_mb}MB {end_docs_mb}MB", shell=True)
+    run_command(f"sudo parted -a optimal -s {DRIVE} mkpart primary {start_docs_mib}MiB {end_docs_mib}MiB", shell=True)
     log("Created documents partition.")
 
-    start_unencrypted_mb = end_docs_mb
-    run_command(f"sudo parted -a optimal -s {DRIVE} mkpart primary {start_unencrypted_mb}MB 100%", shell=True)
+    # Create unencrypted partition
+    start_unencrypted_mib = end_docs_mib
+    run_command(f"sudo parted -a optimal -s {DRIVE} mkpart primary {start_unencrypted_mib}MiB 100%", shell=True)
     log("Created unencrypted partition for scripts/instructions.")
 
+    # Refresh partition table to recognize new partitions
     run_command(f"sudo partprobe {DRIVE}", shell=True)
-    time.sleep(2)
+    run_command("sudo udevadm settle", shell=True)
+    time.sleep(5)  # Increased sleep to ensure partitions are recognized
 
     setup_kali_partition()
     if CREATE_DOCS:
@@ -230,66 +296,29 @@ def fix_partition_table():
     setup_unencrypted_partition()
 
 def fix_partition_table_tails():
-    log("Fixing partition table to reclaim remaining space...")
+    log("Setting up partition table for Tails...")
 
-    run_command(f"sudo parted -a optimal -s {DRIVE} rm 2", shell=True)
-    log("Deleted partition 2.")
+    # Clear existing partitions if any
+    run_command(f"sudo parted -a optimal -s {DRIVE} mklabel gpt", shell=True)
+    run_command(f"sudo partprobe {DRIVE}", shell=True)
+    run_command("sudo udevadm settle", shell=True)
+    time.sleep(5)  # Increased sleep to ensure partitions are recognized
 
-    result = subprocess.run(["sudo", "parted", "-s", DRIVE, "unit", "MB", "print"], capture_output=True, text=True)
-    end_of_p1 = None
-    for line in result.stdout.strip().splitlines():
-        if line.strip().startswith("1"):
-            parts = line.strip().split()
-            end_of_p1 = parts[2]
-            break
-    if end_of_p1 is None:
-        log("Error: Could not find end of partition 1.")
-        sys.exit(1)
-
-    log(f"End of partition 1: {end_of_p1}")
-
-    result = subprocess.run(["lsblk", "-bn", "-o", "SIZE", DRIVE], capture_output=True, text=True)
-    sizes = result.stdout.strip().splitlines()
-    total_size_bytes = int(sizes[0].strip())
-    total_size_mb = total_size_bytes / (1024 * 1024)
-
-    start_docs_mb = float(end_of_p1.replace('MB', ''))
-    end_docs_mb = start_docs_mb + (size_docs_gb * 1024)
-
-    size_docs = input("Enter size for documents partition in GB (leave blank to use remaining space minus 1GB): ")
-    if size_docs:
-        try:
-            size_docs_gb = float(size_docs)
-            end_docs_mb = start_docs_mb + (size_docs_gb * 1024)
-            if end_docs_mb > total_size_mb - 1024:
-                log("Error: Documents partition size exceeds available space when reserving 1GB for unencrypted partition.")
-                sys.exit(1)
-        except ValueError:
-            log("Invalid size entered for documents partition. Exiting.")
-            sys.exit(1)
-    else:
-        end_docs_mb = total_size_mb - 1024
-
-    run_command(f"sudo parted -a optimal -s {DRIVE} mkpart primary {start_docs_mb}MB {end_docs_mb}MB", shell=True)
-    log("Created documents partition.")
-
-    start_unencrypted_mb = end_docs_mb
-    run_command(f"sudo parted -a optimal -s {DRIVE} mkpart primary {start_unencrypted_mb}MB 100%", shell=True)
+    # Tails typically does not require a persistence partition, but we'll create an unencrypted partition for scripts/instructions
+    run_command(f"sudo parted -a optimal -s {DRIVE} mkpart primary 1MiB 100%", shell=True)
     log("Created unencrypted partition for scripts/instructions.")
 
+    # Refresh partition table to recognize new partitions
     run_command(f"sudo partprobe {DRIVE}", shell=True)
-    time.sleep(2)
+    run_command("sudo udevadm settle", shell=True)
+    time.sleep(5)
 
-    setup_kali_partition()
-    if CREATE_DOCS:
-        setup_docs_partition()
     setup_unencrypted_partition()
 
 def setup_kali_partition():
     global DRIVE
     PERSIST_PART = get_partition_name(DRIVE, 2)
 
-    # Wipe existing signatures on the partition
     run_command(["sudo", "wipefs", "--all", PERSIST_PART])
 
     log("Configuring encrypted persistence partition...")
@@ -329,9 +358,13 @@ def setup_kali_partition():
 
 def setup_docs_partition():
     global DRIVE
-    DOCS_PART = get_partition_name(DRIVE, 3)
+    DOCS_PART = get_partition_name(DRIVE, get_last_partition_number() - 1)
+
+    # Check if the partition exists
+    if not os.path.exists(DOCS_PART):
+        log(f"Error: Partition {DOCS_PART} does not exist.")
+        sys.exit(1)
 
-    # Wipe existing signatures on the partition
     run_command(["sudo", "wipefs", "--all", DOCS_PART])
 
     log("Configuring VeraCrypt encryption for documents partition...")
@@ -359,9 +392,14 @@ def setup_docs_partition():
 
 def setup_unencrypted_partition():
     global DRIVE
-    UNENCRYPTED_PART = get_partition_name(DRIVE, 4)
+    UNENCRYPTED_PART = get_partition_name(DRIVE, get_last_partition_number())
 
-    # Format the partition with FAT32
+    # Check if the partition exists
+    if not os.path.exists(UNENCRYPTED_PART):
+        log(f"Error: Partition {UNENCRYPTED_PART} does not exist.")
+        sys.exit(1)
+
+    log(f"Attempting to format partition {UNENCRYPTED_PART} with FAT32 filesystem.")
     run_command(f"sudo mkfs.vfat -n 'TOOLS' {UNENCRYPTED_PART}", shell=True)
     log("Formatted unencrypted partition with FAT32 filesystem.")
 
@@ -369,40 +407,37 @@ def setup_unencrypted_partition():
     run_command(f"sudo mount {UNENCRYPTED_PART} /mnt/unencrypted", shell=True)
 
     instructions = f"""
-To mount the encrypted documents partition:
-
-1. Open a terminal.
-2. Run: sudo veracrypt --text --mount {get_partition_name(DRIVE, 3)} /mnt/veracrypt_docs
+To mount the encrypted documents partition, use the provided 'mount_encrypted_partitions.sh' script.
+To unmount and lock it, use the 'cleanup_encrypted_partitions.sh' script.
 
 **Automount Script:**
-You can use the provided script 'mount_encrypted_partitions.sh' to automate this process.
-
-Usage:
-sudo ./mount_encrypted_partitions.sh
+Run: sudo ./mount_encrypted_partitions.sh
 
 **Cleanup Script:**
-To unmount and lock the encrypted documents partition, run:
-sudo ./cleanup_encrypted_partitions.sh
+Run: sudo ./cleanup_encrypted_partitions.sh
 """
 
-    # Write README.txt
     with open("/tmp/README.txt", "w") as readme_file:
         readme_file.write(instructions)
     run_command("sudo cp /tmp/README.txt /mnt/unencrypted/README.txt", shell=True)
     run_command("sudo rm /tmp/README.txt", shell=True)
     log("Created README.txt with mounting instructions.")
 
-    # Write mount script for documents partition
-    mount_script = f"""#!/bin/bash
+    mount_script = """#!/bin/bash
 # Script to mount encrypted documents partition
 
-if [ -b "{get_partition_name(DRIVE, 3)}" ]; then
-    echo "Mounting VeraCrypt documents partition..."
+echo "Available drives:"
+lsblk -o NAME,SIZE,TYPE | grep disk
+read -p "Enter the drive path for the documents partition (default: /dev/sdc1): " DRIVE_PATH
+DRIVE_PATH=${DRIVE_PATH:-/dev/sdc1}
+
+if [ -b "$DRIVE_PATH" ]; then
+    echo "Mounting VeraCrypt documents partition at $DRIVE_PATH..."
     sudo mkdir -p /mnt/veracrypt_docs
-    sudo veracrypt --text --mount {get_partition_name(DRIVE, 3)} /mnt/veracrypt_docs
+    sudo veracrypt --text --mount "$DRIVE_PATH" /mnt/veracrypt_docs
     echo "Documents partition mounted at /mnt/veracrypt_docs."
 else
-    echo "Documents partition not found."
+    echo "Error: Partition $DRIVE_PATH not found."
 fi
 """
 
@@ -413,21 +448,22 @@ fi
     run_command("sudo rm /tmp/mount_encrypted_partitions.sh", shell=True)
     log("Created mount_encrypted_partitions.sh script.")
 
-    # Write cleanup script for documents partition
     cleanup_script = """#!/bin/bash
 # Script to unmount and lock encrypted documents partition
 
-echo "Unmounting and locking encrypted documents partition..."
+echo "Available drives:"
+lsblk -o NAME,SIZE,TYPE | grep disk
+read -p "Enter the drive path for the documents partition (default: /dev/sdc1): " DRIVE_PATH
+DRIVE_PATH=${DRIVE_PATH:-/dev/sdc1}
+
+echo "Unmounting and locking encrypted documents partition at $DRIVE_PATH..."
 
 if mountpoint -q /mnt/veracrypt_docs; then
-    echo "Unmounting VeraCrypt documents partition..."
     sudo veracrypt --text --dismount /mnt/veracrypt_docs
     echo "VeraCrypt documents partition unmounted."
 else
-    echo "Documents partition is not mounted."
+    echo "Documents partition is not currently mounted."
 fi
-
-echo "Cleanup complete. Encrypted documents partition is locked and unmounted."
 """
 
     with open("/tmp/cleanup_encrypted_partitions.sh", "w") as script_file:
@@ -440,6 +476,29 @@ echo "Cleanup complete. Encrypted documents partition is locked and unmounted."
     run_command("sudo umount /mnt/unencrypted", shell=True)
     log("Unencrypted partition setup complete.")
 
+def get_last_partition_number():
+    """Returns the highest partition number on the DRIVE."""
+    result = subprocess.run(["lsblk", "-ln", "-o", "NAME", DRIVE], capture_output=True, text=True)
+    partitions = [line.strip() for line in result.stdout.strip().splitlines() if line.strip()]
+    partition_numbers = []
+    for part in partitions:
+        if 'p' in part:
+            # Handle /dev/nvme0n1p1 format
+            if part.startswith(os.path.basename(DRIVE)):
+                num = part.replace(os.path.basename(DRIVE), '').replace('p', '')
+                if num.isdigit():
+                    partition_numbers.append(int(num))
+        else:
+            # Handle /dev/sda1 format
+            if part.startswith(os.path.basename(DRIVE)):
+                num = part.replace(os.path.basename(DRIVE), '')
+                if num.isdigit():
+                    partition_numbers.append(int(num))
+    if not partition_numbers:
+        log(f"No partitions found on {DRIVE}.")
+        sys.exit(1)
+    return max(partition_numbers)
+
 def main():
     global DEBUG, FAST_MODE, CREATE_KALI, CREATE_DOCS, CREATE_TAILS, KALI_ISO, TAILS_ISO, DRIVE
 
@@ -485,7 +544,7 @@ def main():
 
     check_dependencies()
 
-    if CREATE_KALI or CREATE_TAILS:
+    if CREATE_KALI or CREATE_TAILS or CREATE_DOCS:
         setup_usb()
     else:
         list_drives()
@@ -493,11 +552,7 @@ def main():
         prepare_drive(DRIVE)
 
         if CREATE_DOCS:
-            if CREATE_TAILS:
-                fix_partition_table_tails()
-            else:
-                fix_partition_table()
-                setup_docs_partition()
+            fix_partition_table_docs_only()
         else:
             setup_unencrypted_partition()