fixed issues and improved security

.gitignore

@@ -1 +1,59 @@
-*.log
+# Python cache and bytecode
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+
+# Log files - may contain sensitive paths/device info
+*.log
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual environments
+venv/
+env/
+ENV/
+.venv
+
+# IDE and editor files
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+.DS_Store
+
+# Testing
+.pytest_cache/
+.coverage
+htmlcov/
+
+# Claude docs
+.claude_docs
+
+# Sensitive files that should never be committed
+*.iso
+*.img
+*.key
+*.pem
+*.p12
+*.pfx
+test_*
+*.bak
+*.backup

README.md

@@ -2,31 +2,38 @@
 
 ## Introduction
 
-The **Covert SD Card Tool** is a Python script designed to automate the process of setting up a bootable USB drive with either Kali Linux or Tails OS. It includes options to create encrypted persistence partitions, encrypted documents partitions, and an unencrypted partition containing helpful scripts and instructions. This tool simplifies the complex steps involved in preparing a secure, portable operating system on a USB drive.
+The **Covert SD Card Tool** is a Python script designed to automate the process of setting up a bootable USB/SD card with either Kali Linux or Tails OS. It includes options to create encrypted persistence partitions, secure document storage, and user-friendly access scripts. This tool simplifies the complex steps involved in preparing a secure, portable operating system on a USB drive or SD card.
 
 ## Features
 
-- **Install Kali Linux or Tails OS** on a USB drive.
-- **Create an encrypted persistence partition** for Kali Linux.
-- **Create an encrypted documents partition** using VeraCrypt.
-- **Add an unencrypted partition** containing automount scripts and instructions.
-- **Customizable encryption options** with a fast setup mode.
-- **Automated dependency checking and installation**.
+- **Install Kali Linux or Tails OS** on a USB/SD card
+- **Create an encrypted persistence partition** for Kali Linux (LUKS encryption)
+- **Create a maximum-security encrypted documents partition** with triple-cascade encryption
+- **Secure drive wiping** using multi-pass shred for data sanitization
+- **User-friendly access scripts** for mounting and locking secure storage
+- **OPSEC-focused design** - generated scripts use generic terminology
+- **Automated dependency checking and installation**
 
 ## Prerequisites
 
-- **Operating System:** Linux (Debian-based distributions recommended).
-- **Python Version:** Python 3.x.
+- **Operating System:** Linux (Debian-based distributions recommended)
+- **Python Version:** Python 3.x
+- **Root Access:** Required for disk operations
 - **Dependencies:**
-  - `parted`
-  - `cryptsetup`
-  - `lsblk`
-  - `dd`
-  - `sgdisk`
-  - `wipefs`
-  - `bc`
-  - `fdisk`
-  - `veracrypt` (The script can install this if not present).
+  - `parted` - Partition management
+  - `cryptsetup` - LUKS encryption
+  - `lsblk` - Block device listing
+  - `dd` - Disk writing
+  - `sgdisk` - GPT partition manipulation
+  - `wipefs` - Filesystem signature removal
+  - `shred` - Secure data wiping
+  - `bc` - Calculator for partition math
+  - `fdisk` - Partition table manipulation
+  - `veracrypt` - Document partition encryption
+  - `lsof`, `fuser` - Process detection
+  - `udevadm` - Device management
+
+**Note:** The script will automatically detect missing dependencies and offer to install them.
 
 ## Installation
 
@@ -53,13 +60,14 @@ sudo ./covert_sd_card_tool.py [options]
 
 ### Command-Line Options
 
-- `-a`, `--all` : Set up both the OS bootable USB and the documents partition.
-- `-k`, `--kali` : Create a Kali bootable USB and persistence partition.
-- `-t`, `--tails` : Create a Tails bootable USB (no persistence).
-- `-d`, `--docs` : Create an encrypted documents partition.
-- `-i`, `--iso` : Path to the Kali or Tails ISO file.
-- `--fast` : Enable fast setup with less secure encryption.
-- `--debug` : Enable debug mode.
+- `-a`, `--all` : Set up both the OS bootable USB and the documents partition (defaults to Kali)
+- `-k`, `--kali` : Create a Kali bootable USB and persistence partition
+- `-t`, `--tails` : Create a Tails bootable USB (no persistence, mutually exclusive with `-a`)
+- `-c`, `--custom` : Create a custom ISO bootable USB (uses Kali-style partitioning with persistence)
+- `-d`, `--docs` : Create an encrypted documents partition
+- `-i`, `--iso` : Path to the ISO file (Kali, Tails, or custom)
+- `--fast` : Enable fast setup mode (Note: Documents partition always uses maximum security)
+- `--debug` : Enable debug mode with verbose logging
 
 ### Examples
 
@@ -69,20 +77,192 @@ sudo ./covert_sd_card_tool.py [options]
   sudo ./covert_sd_card_tool.py -a -i /path/to/kali.iso
   ```
 
-- **Install Tails and Encrypted Documents Partition:**
+- **Install Tails with Encrypted Documents Partition:**
 
   ```bash
-  sudo ./covert_sd_card_tool.py -a -t -i /path/to/tails.iso
+  sudo ./covert_sd_card_tool.py -t -d -i /path/to/tails.iso
   ```
 
-- **Install Tails Without Encrypted Documents Partition:**
+- **Install Tails Only (No Documents Partition):**
 
   ```bash
   sudo ./covert_sd_card_tool.py -t -i /path/to/tails.iso
   ```
 
-- **Install Encrypted Documents Partition Only:**
+- **Create Encrypted Documents Partition Only (No OS):**
 
   ```bash
   sudo ./covert_sd_card_tool.py -d
-  ```
+  ```
+
+- **Install Custom ISO (e.g., Parrot OS, BlackArch) with Persistence and Documents:**
+
+  ```bash
+  sudo ./covert_sd_card_tool.py -c -d -i /path/to/custom.iso
+  ```
+
+- **Install Custom ISO with Persistence Only (No Documents):**
+
+  ```bash
+  sudo ./covert_sd_card_tool.py -c -i /path/to/custom.iso
+  ```
+
+## Security Features
+
+### Documents Partition Encryption
+
+The documents partition uses **maximum security settings** regardless of the `--fast` flag:
+
+- **Triple Cascade Encryption:** AES-Twofish-Serpent (3 layers)
+- **Hash Algorithm:** SHA-512
+- **Key Derivation:** PIM 2000 (enforced for strong key stretching)
+- **Filesystem:** ext4
+- **Full Format:** Always performs full format (no quick format) to overwrite old data
+
+This configuration provides **military-grade security** suitable for sensitive documents like travel documents, credentials, and confidential files.
+
+### Persistence Partition Encryption (Kali)
+
+- **Normal Mode:** AES-XTS-PLAIN64, 512-bit keys, SHA-512, 5 second iteration time
+- **Fast Mode:** AES-CBC-ESSIV:SHA256, 256-bit keys, SHA-256, 1 second iteration time
+
+### Secure Drive Wiping
+
+When you choose to wipe the drive, the tool uses `shred`:
+- **3 passes** of random data overwriting
+- **Final pass** with zeros
+- Makes data recovery virtually impossible
+- Confirmation required (must type 'WIPE')
+
+### OPSEC (Operational Security)
+
+The generated helper scripts use **generic terminology** to avoid disclosing encryption methods:
+
+- Scripts renamed to `mount_storage.sh` and `lock_storage.sh` (instead of mentioning encryption types)
+- README uses terms like "secure storage" instead of specific encryption names
+- No algorithm disclosure in user-facing documentation on the device
+- Suitable for travel scenarios where device inspection may occur
+
+## Generated Helper Scripts
+
+The tool creates a small unencrypted partition (TOOLS) containing:
+
+### `mount_storage.sh`
+- Interactive script to mount the encrypted documents partition
+- Shows available devices and validates input
+- Mounts to `/mnt/secure_storage`
+- Clear error messages and success confirmations
+
+### `lock_storage.sh`
+- Safely dismounts and locks the encrypted storage
+- Checks for open files before locking
+- Shows warnings if applications are still using the storage
+- Syncs pending writes before dismount
+- Prevents data loss from improper ejection
+
+### `README.txt`
+- Simple instructions for non-technical users
+- Generic terminology (no encryption disclosure)
+- Step-by-step mount/lock procedures
+
+## Important Security Notes
+
+⚠️ **Password Strength:** Use strong passphrases (20+ characters, mixed case, numbers, symbols)
+
+⚠️ **No Password Recovery:** If you forget your password, your data is **permanently inaccessible**
+
+⚠️ **PIM Value:** The tool enforces PIM 2000 for documents partition - this adds 2-3 seconds to unlock time but massively increases security
+
+⚠️ **Always Lock Before Removal:** Use `lock_storage.sh` before removing the device to prevent data corruption
+
+⚠️ **Fast Mode:** While available for persistence partition, documents partition **always uses maximum security**
+
+## Partition Layout Examples
+
+### Kali + Documents (`-a`)
+1. **Partition 1:** Kali Live OS (bootable)
+2. **Partition 2:** LUKS encrypted persistence (configurable size, e.g., 4GB)
+3. **Partition 3:** VeraCrypt encrypted documents (remaining space minus 1GB)
+4. **Partition 4:** Unencrypted tools partition (1GB, FAT32, contains scripts)
+
+### Tails Only (`-t`)
+- **Entire Drive:** Tails Live OS (bootable, no additional partitions)
+
+### Tails + Documents (`-t -d`)
+1. **Partition 1:** Tails Live OS (bootable, 12MB EFI System)
+2. **Partition 2:** Tails system partition (~2-3GB depending on Tails version)
+3. **Partition 3:** VeraCrypt encrypted documents (remaining space minus 1GB)
+4. **Partition 4:** Unencrypted tools partition (1GB, FAT32, contains scripts)
+
+### Documents Only (`-d`)
+1. **Partition 1:** VeraCrypt encrypted documents (remaining space minus 1GB)
+2. **Partition 2:** Unencrypted tools partition (1GB, FAT32, contains scripts)
+
+### Custom ISO + Documents (`-c -d`)
+1. **Partition 1:** Custom Live OS (bootable)
+2. **Partition 2:** LUKS encrypted persistence (configurable size, e.g., 4GB)
+3. **Partition 3:** VeraCrypt encrypted documents (remaining space minus 1GB)
+4. **Partition 4:** Unencrypted tools partition (1GB, FAT32, contains scripts)
+
+**Note:** Custom ISO mode uses Kali-style partitioning. Works well with Debian-based live ISOs like Parrot OS, BlackArch, BackBox, etc.
+
+## Using Custom ISOs
+
+The `-c` (custom) flag allows you to use **any bootable ISO** and set it up with encrypted persistence and documents partitions. This is useful for:
+
+### Compatible ISOs
+- **Parrot Security OS** - Privacy-focused security distro
+- **BlackArch Linux** - Penetration testing distro
+- **BackBox** - Ubuntu-based penetration testing
+- **Pentoo** - Gentoo-based security distro
+- **Any Debian/Ubuntu-based live ISO**
+
+### How It Works
+Custom ISOs are treated like Kali Linux:
+1. ISO is flashed to the drive
+2. LUKS encrypted persistence partition is created (if you want settings to persist)
+3. VeraCrypt encrypted documents partition is added (if `-d` flag is used)
+4. Tools partition with mount/lock scripts
+
+### Example: Parrot OS with Docs
+```bash
+sudo ./covert_sd_card_tool.py -c -d -i ~/Downloads/parrot-security.iso
+```
+
+### Compatibility Notes
+- **Best for:** Debian/Ubuntu-based live ISOs
+- **May not work with:** Arch-based ISOs (different partition structure), Windows ISOs
+- **Persistence:** Depends on the ISO supporting LUKS persistence (Debian-based usually do)
+- If persistence doesn't work with your ISO, you can still use the documents partition
+
+## Troubleshooting
+
+### Device Busy Errors
+- The tool automatically unmounts partitions and kills processes using the drive
+- If problems persist, manually unmount: `sudo umount /dev/sdX*`
+- Check for processes: `sudo lsof /dev/sdX`
+
+### VeraCrypt Not Found
+- On Debian/Ubuntu: `sudo apt install veracrypt`
+- Or the script will offer to install it automatically
+
+### Permission Denied
+- Always run with `sudo`
+- Ensure your user has sudo privileges
+
+### Drive Not Detected
+- Check if drive is connected: `lsblk`
+- Verify drive path (e.g., `/dev/sdb` not `/dev/sdb1`)
+- Try unplugging and reconnecting the device
+
+## License
+
+This tool is provided as-is for educational and legitimate security purposes only. Use responsibly and in compliance with applicable laws.
+
+## Contributing
+
+Contributions, bug reports, and feature requests are welcome! Please open an issue or submit a pull request.
+
+## Disclaimer
+
+This tool performs destructive operations on storage devices. **Always verify you've selected the correct drive** before proceeding. The authors are not responsible for data loss.

covert_sd_card_tool.py

@@ -17,30 +17,55 @@ LOG_FILE = f"covert_sd_setup_{TIMESTAMP}.log"
 CREATE_KALI = False
 CREATE_DOCS = False
 CREATE_TAILS = False
+CREATE_CUSTOM = False
 KALI_ISO = ""
 TAILS_ISO = ""
+CUSTOM_ISO = ""
 DRIVE = ""
 
 def log(message):
+    """Logs a message to both the log file and the terminal."""
     with open(LOG_FILE, "a") as log_file:
         log_file.write(message + "\n")
     print(message)
 
 def run_command(command, shell=False, interactive=False):
+    """
+    Runs a system command.
+    
+    Args:
+        command (list or str): The command to execute.
+        shell (bool): Whether to execute the command through the shell.
+        interactive (bool): If True, streams the command's output live.
+    """
     if DEBUG:
         log(f"Running command: {command}")
     try:
         if interactive:
-            subprocess.run(command, shell=shell, check=True)
+            # Use subprocess.run with no stdout/stderr capture for truly interactive commands
+            # This allows the command to directly interact with the terminal
+            result = subprocess.run(command, shell=shell, check=True)
+            return
         else:
-            result = subprocess.run(
-                command,
-                shell=shell,
-                check=True,
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
-                universal_newlines=True
-            )
+            # Handle non-interactive commands
+            if isinstance(command, list):
+                result = subprocess.run(
+                    command,
+                    shell=shell,
+                    check=True,
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                    universal_newlines=True
+                )
+            else:
+                result = subprocess.run(
+                    command,
+                    shell=shell,
+                    check=True,
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                    universal_newlines=True
+                )
             if result.stdout:
                 log(result.stdout.strip())
             if result.stderr:
@@ -50,7 +75,11 @@ def run_command(command, shell=False, interactive=False):
         sys.exit(1)
 
 def check_dependencies():
-    dependencies = ["parted", "cryptsetup", "lsblk", "dd", "sgdisk", "wipefs", "bc", "fdisk", "veracrypt", "lsof", "fuser", "mountpoint", "udevadm"]
+    """Checks and installs missing dependencies."""
+    dependencies = [
+        "parted", "cryptsetup", "lsblk", "dd", "sgdisk", "wipefs",
+        "bc", "fdisk", "veracrypt", "lsof", "fuser", "mountpoint", "udevadm"
+    ]
     missing = []
     for dep in dependencies:
         if not shutil.which(dep):
@@ -66,23 +95,45 @@ def check_dependencies():
             sys.exit(1)
 
 def list_drives():
+    """Lists all available drives."""
     log("Available drives:")
     result = subprocess.run(["lsblk", "-J", "-o", "NAME,SIZE,TYPE"], capture_output=True, text=True)
-    lsblk_output = json.loads(result.stdout)
-    for device in lsblk_output['blockdevices']:
-        if device['type'] == 'disk':
-            name = device['name']
-            size = device['size']
-            drive = f"/dev/{name} {size}"
-            log(drive)
+    try:
+        lsblk_output = json.loads(result.stdout)
+        for device in lsblk_output['blockdevices']:
+            if device['type'] == 'disk':
+                name = device['name']
+                size = device['size']
+                drive = f"/dev/{name} {size}"
+                log(drive)
+    except json.JSONDecodeError:
+        log("Error: Unable to parse lsblk output.")
+        sys.exit(1)
 
 def get_partition_name(drive, partition_number):
+    """
+    Generates the partition name based on the drive and partition number.
+    
+    Args:
+        drive (str): The drive path (e.g., /dev/sda).
+        partition_number (int): The partition number.
+        
+    Returns:
+        str: The full partition path (e.g., /dev/sda1 or /dev/nvme0n1p1).
+    """
     if 'nvme' in drive or 'mmcblk' in drive:
         return f"{drive}p{partition_number}"
     else:
         return f"{drive}{partition_number}"
 
 def prepare_drive(drive):
+    """
+    Unmounts any mounted partitions, disables swap, and kills processes using the drive.
+    
+    Args:
+        drive (str): The drive to prepare (e.g., /dev/sda).
+    """
+    # Unmount all mounted partitions
     result = subprocess.run(["lsblk", "-lnp", drive], capture_output=True, text=True)
     for line in result.stdout.strip().splitlines():
         parts = line.strip().split()
@@ -91,6 +142,7 @@ def prepare_drive(drive):
             log(f"Unmounting {part}...")
             run_command(["sudo", "umount", "-l", part])
 
+    # Disable swap if it's on the drive
     with open("/proc/swaps") as swaps_file:
         for line in swaps_file:
             if drive in line:
@@ -98,6 +150,7 @@ def prepare_drive(drive):
                 log(f"Disabling swap on {swap_part}...")
                 run_command(["sudo", "swapoff", swap_part])
 
+    # Kill any processes using the drive
     log(f"Checking for processes using {drive}...")
     result = subprocess.run(["sudo", "lsof", drive], capture_output=True, text=True)
     if result.stdout.strip():
@@ -113,6 +166,9 @@ def prepare_drive(drive):
         log(f"No processes are using {drive}.")
 
 def setup_usb():
+    """
+    Sets up the bootable USB (Tails or Kali) and creates additional partitions if required.
+    """
     global DRIVE
     log("Setting up bootable USB or preparing partitions...")
     list_drives()
@@ -129,16 +185,24 @@ def setup_usb():
         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.")
+            log("This will securely overwrite the entire drive with random data. This may take a while...")
+            confirm_wipe = input(f"Are you ABSOLUTELY sure you want to completely wipe {DRIVE}? Type 'WIPE' to confirm: ")
+            if confirm_wipe == "WIPE":
+                run_command(["sudo", "wipefs", "--all", DRIVE])
+                run_command(["sudo", "sgdisk", "--zap-all", DRIVE])
+                # Secure wipe using shred - 3 passes with random data, then zeros
+                run_command(["sudo", "shred", "-vfz", "-n", "3", DRIVE], interactive=True)
+                log(f"{DRIVE} securely wiped successfully.")
+            else:
+                log("Wipe canceled. Proceeding without full wipe (only clearing partition table)...")
+                run_command(["sudo", "wipefs", "--all", DRIVE])
+                run_command(["sudo", "sgdisk", "--zap-all", DRIVE])
 
     # Initialize variables to track if ISO has been written
     iso_written = False
 
-    if CREATE_KALI or CREATE_TAILS:
-        global KALI_ISO, TAILS_ISO
+    if CREATE_KALI or CREATE_TAILS or CREATE_CUSTOM:
+        global KALI_ISO, TAILS_ISO, CUSTOM_ISO
         if CREATE_KALI:
             if not KALI_ISO:
                 KALI_ISO = input("Enter the path to the Kali ISO file: ")
@@ -153,46 +217,68 @@ def setup_usb():
                 log(f"Error: Tails ISO file not found at {TAILS_ISO}")
                 sys.exit(1)
             ISO_PATH = TAILS_ISO
+        elif CREATE_CUSTOM:
+            if not CUSTOM_ISO:
+                CUSTOM_ISO = input("Enter the path to your custom ISO file: ")
+            if not os.path.isfile(CUSTOM_ISO):
+                log(f"Error: Custom ISO file not found at {CUSTOM_ISO}")
+                sys.exit(1)
+            ISO_PATH = CUSTOM_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.")
-        iso_written = True
+        if CREATE_TAILS:
+            # For Tails, flash the ISO to the entire drive without any partitioning
+            log(f"Flashing Tails ISO to {DRIVE}...")
+            run_command(f"sudo dd if='{ISO_PATH}' of='{DRIVE}' bs=64M status=progress conv=fdatasync", shell=True, interactive=True)
+            log(f"Tails ISO flashed to {DRIVE} successfully.")
+            iso_written = True
+        else:
+            # For Kali or Custom ISO, flash and proceed with partition setup
+            log(f"Flashing ISO to {DRIVE}...")
+            run_command(f"sudo dd if='{ISO_PATH}' of='{DRIVE}' bs=64M status=progress conv=fdatasync", shell=True, interactive=True)
+            log(f"ISO flashed to {DRIVE} successfully.")
+            iso_written = True
 
     # After writing ISO, set up partitions accordingly
     if iso_written:
-        if CREATE_KALI:
+        if CREATE_KALI or CREATE_CUSTOM:
+            # Both Kali and custom ISOs use the same partitioning approach
             fix_partition_table()
-        if CREATE_TAILS:
-            fix_partition_table_tails()
-
-    # If documents partition is requested, set it up
-    if CREATE_DOCS and not (CREATE_KALI or CREATE_TAILS):
-        # If not setting up OS bootable USB, proceed directly
-        fix_partition_table_docs_only()
-    elif CREATE_DOCS and (CREATE_KALI or CREATE_TAILS):
-        # If setting up OS bootable USB, documents partition is handled in fix_partition_table() or fix_partition_table_tails()
-        # Depending on the specific needs, you might want to call fix_partition_table_docs_only() here
-        # Let's call it to ensure documents partition is created
+        elif CREATE_TAILS:
+            if CREATE_DOCS:
+                # For Tails with docs, add partitions after Tails
+                fix_partition_table_tails()
+            else:
+                # Tails only, no additional partitions
+                log("Tails installation complete. No additional partitions requested.")
+    elif CREATE_DOCS:
+        # If no ISO was written, just set up documents partition
         fix_partition_table_docs_only()
 
 def fix_partition_table_docs_only():
+    """
+    Sets up partitions solely for encrypted documents without altering existing OS partitions.
+    """
     log("Setting up partitions for documents only...")
 
-    # Clear existing partitions if any
+    # Clear existing GPT label to start fresh
     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)
+    result = subprocess.run(["lsblk", "-b", "-d", "-n", "-o", "SIZE", DRIVE], capture_output=True, text=True)
     try:
         total_size_bytes = int(result.stdout.strip())
     except ValueError:
-        log("Error: Unable to determine drive size.")
+        log(f"Error: Unable to determine drive size. lsblk output: {result.stdout}")
         sys.exit(1)
     total_size_mib = total_size_bytes / (1024 * 1024)  # Convert to MiB
+    total_size_gb = total_size_mib / 1024  # Convert to GiB
+    available_gb = (total_size_mib - 1024) / 1024  # Minus 1GB reserved for scripts
+
+    log(f"Total drive size: {total_size_gb:.2f} GB ({total_size_mib:.0f} MiB)")
+    log(f"Available space for documents (minus 1GB for scripts): {available_gb:.2f} GB")
 
     # Ask for document partition size
     size_docs = input("Enter size for documents partition in GB (leave blank to use remaining space minus 1GB): ")
@@ -227,7 +313,103 @@ def fix_partition_table_docs_only():
     setup_unencrypted_partition()
     setup_docs_partition()
 
+def fix_partition_table_tails():
+    """
+    Adds encrypted documents partition after Tails installation without breaking boot.
+    Tails creates its own partition table, so we extend it carefully.
+    """
+    log("Adding encrypted documents partition after Tails installation...")
+
+    # Wait for Tails partitions to settle
+    run_command(f"sudo partprobe {DRIVE}", shell=True)
+    run_command("sudo udevadm settle", shell=True)
+    time.sleep(5)
+
+    # Get current partition information
+    result = subprocess.run(["sudo", "parted", "-s", DRIVE, "unit", "MiB", "print"], capture_output=True, text=True)
+    log("Current partition table after Tails installation:")
+    log(result.stdout)
+
+    # Find the end of the last Tails partition
+    last_partition_end = None
+    partition_lines = [line for line in result.stdout.strip().splitlines() if line.strip() and line.strip()[0].isdigit()]
+
+    if partition_lines:
+        # Get the last partition's end position
+        last_line = partition_lines[-1]
+        parts = last_line.strip().split()
+        if len(parts) >= 3:
+            last_partition_end = parts[2].replace('MiB', '')
+
+    if last_partition_end is None:
+        log("Error: Could not determine end of Tails partitions.")
+        sys.exit(1)
+
+    log(f"Last Tails partition ends at: {last_partition_end}MiB")
+
+    # Get the total size of the drive
+    result = subprocess.run(["lsblk", "-b", "-d", "-n", "-o", "SIZE", DRIVE], capture_output=True, text=True)
+    try:
+        total_size_bytes = int(result.stdout.strip())
+    except ValueError:
+        log(f"Error: Unable to determine drive size. lsblk output: {result.stdout}")
+        sys.exit(1)
+    total_size_mib = total_size_bytes / (1024 * 1024)
+    total_size_gb = total_size_mib / 1024
+    available_after_tails_mib = total_size_mib - float(last_partition_end)
+    available_after_tails_gb = available_after_tails_mib / 1024
+
+    log(f"Total drive size: {total_size_gb:.2f} GB ({total_size_mib:.0f} MiB)")
+    log(f"Available space after Tails: {available_after_tails_gb:.2f} GB")
+
+    if available_after_tails_gb < 2:
+        log(f"Warning: Only {available_after_tails_gb:.2f} GB available after Tails. Need at least 2GB for docs + tools partitions.")
+        proceed = input("Continue anyway? (y/n) [Default: n]: ") or "n"
+        if proceed.lower() != "y":
+            log("Partition setup canceled.")
+            return
+
+    # Set up documents partition
+    start_docs_mib = float(last_partition_end)
+    remaining_after_tails_gb = (total_size_mib - start_docs_mib) / 1024
+    log(f"Remaining space for documents: {remaining_after_tails_gb:.2f} GB (will reserve 1GB for tools partition)")
+
+    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_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 tools partition.")
+                sys.exit(1)
+        except ValueError:
+            log("Invalid size entered for documents partition. Exiting.")
+            sys.exit(1)
+    else:
+        end_docs_mib = total_size_mib - 1024  # Reserve 1GB for tools 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 after Tails.")
+
+    # Create unencrypted tools 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 tools partition for scripts/instructions.")
+
+    # Refresh partition table
+    run_command(f"sudo partprobe {DRIVE}", shell=True)
+    run_command("sudo udevadm settle", shell=True)
+    time.sleep(5)
+
+    setup_docs_partition()
+    setup_unencrypted_partition()
+
+    log("Tails + encrypted documents partition setup complete!")
+
 def fix_partition_table():
+    """
+    Fixes the partition table for Kali Linux by adding persistence and documents partitions.
+    """
     log("Fixing partition table to reclaim remaining space...")
 
     # Attempt to delete partition 2 if it exists
@@ -253,13 +435,18 @@ def fix_partition_table():
     log(f"End of partition 1: {end_of_p1}MiB")
 
     # Get the total size of the drive in bytes
-    result = subprocess.run(["lsblk", "-b", "-n", "-o", "SIZE", DRIVE], capture_output=True, text=True)
+    result = subprocess.run(["lsblk", "-b", "-d", "-n", "-o", "SIZE", DRIVE], capture_output=True, text=True)
     try:
         total_size_bytes = int(result.stdout.strip())
     except ValueError:
-        log("Error: Unable to determine drive size.")
+        log(f"Error: Unable to determine drive size. lsblk output: {result.stdout}")
         sys.exit(1)
     total_size_mib = total_size_bytes / (1024 * 1024)  # Convert to MiB
+    total_size_gb = total_size_mib / 1024  # Convert to GiB
+    available_after_p1_gb = (total_size_mib - float(end_of_p1)) / 1024
+
+    log(f"Total drive size: {total_size_gb:.2f} GB ({total_size_mib:.0f} MiB)")
+    log(f"Available space after partition 1: {available_after_p1_gb:.2f} GB")
 
     # Ask for persistence partition size
     size_persistence = input("Enter size for persistence partition in GB (e.g., 4): ") or "4"
@@ -277,11 +464,14 @@ def fix_partition_table():
         sys.exit(1)
 
     # Create persistence partition
-    run_command(f"sudo parted -a optimal -s {DRIVE} mkpart primary {start_persistence_mib}MiB {end_persistence_mib}MiB", shell=True)
+    run_command(f"sudo parted -a optimal -s {DRIVE} mkpart primary ext4 {start_persistence_mib}MiB {end_persistence_mib}MiB", shell=True)
     log("Created persistence partition.")
 
     # Set up documents partition
     start_docs_mib = end_persistence_mib
+    remaining_after_persistence_gb = (total_size_mib - end_persistence_mib) / 1024
+    log(f"Remaining space after persistence partition: {remaining_after_persistence_gb:.2f} GB (will reserve 1GB for scripts)")
+
     size_docs = input("Enter size for documents partition in GB (leave blank to use remaining space minus 1GB): ")
     if size_docs:
         try:
@@ -314,27 +504,10 @@ def fix_partition_table():
         setup_docs_partition()
     setup_unencrypted_partition()
 
-def fix_partition_table_tails():
-    log("Setting up partition table for Tails...")
-
-    # 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
-
-    # 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)
-    run_command("sudo udevadm settle", shell=True)
-    time.sleep(5)
-
-    setup_unencrypted_partition()
-
 def setup_kali_partition():
+    """
+    Configures the persistence partition for Kali Linux.
+    """
     global DRIVE
     PERSIST_PART = get_partition_name(DRIVE, 2)
 
@@ -347,6 +520,11 @@ def setup_kali_partition():
 
     log("Configuring encrypted persistence partition...")
 
+    log("You will be prompted to enter a passphrase for the persistence partition.")
+    log("Please choose a strong passphrase and remember it!")
+    log("When asked 'Are you sure? (Type YES in capital letters):', type: YES")
+    log("Then enter your passphrase twice when prompted.")
+
     if FAST_MODE:
         luks_format_cmd = (
             f"sudo cryptsetup luksFormat '{PERSIST_PART}' "
@@ -354,7 +532,8 @@ def setup_kali_partition():
             f"--cipher aes-cbc-essiv:sha256 "
             f"--key-size 256 "
             f"--hash sha256 "
-            f"--iter-time 1000"
+            f"--iter-time 1000 "
+            f"--verify-passphrase"
         )
         mkfs_cmd = f"sudo mkfs.ext3 -L persistence /dev/mapper/kali_USB"
     else:
@@ -363,7 +542,8 @@ def setup_kali_partition():
             f"--cipher aes-xts-plain64 "
             f"--key-size 512 "
             f"--hash sha512 "
-            f"--iter-time 5000"
+            f"--iter-time 5000 "
+            f"--verify-passphrase"
         )
         mkfs_cmd = f"sudo mkfs.ext4 -L persistence /dev/mapper/kali_USB"
 
@@ -381,6 +561,9 @@ def setup_kali_partition():
     log("Kali persistence setup complete.")
 
 def setup_docs_partition():
+    """
+    Configures the VeraCrypt encrypted documents partition.
+    """
     global DRIVE
     DOCS_PART = get_partition_name(DRIVE, get_last_partition_number() - 1)
 
@@ -389,32 +572,48 @@ def setup_docs_partition():
         log(f"Error: Partition {DOCS_PART} does not exist.")
         sys.exit(1)
 
-    run_command(["sudo", "wipefs", "--all", DOCS_PART])
+    # Force wipe existing filesystem signatures
+    run_command(["sudo", "wipefs", "--all", "--force", DOCS_PART])
+    log(f"Removed existing filesystem signatures from {DOCS_PART}.")
 
-    log("Configuring VeraCrypt encryption for documents partition...")
+    log("\n" + "="*70)
+    log("ENCRYPTED VOLUME SETUP")
+    log("="*70)
+    log("\nYou will now create an encrypted volume for sensitive documents.")
+    log("\nWhat you'll be asked:")
+    log("  1. PASSWORD: Enter a strong passphrase (you'll type it twice)")
+    log("     - Use at least 20 characters")
+    log("     - Mix uppercase, lowercase, numbers, and symbols")
+    log("     - Avoid dictionary words or personal information")
+    log("  2. PIM (Personal Iterations Multiplier):")
+    log("     - RECOMMENDED: Enter 2000 for maximum security")
+    log("     - Minimum acceptable: 1000")
+    log("     - Higher = more secure but slower unlock (2-3 seconds)")
+    log("  3. KEYFILES: Leave blank unless you know what you're doing")
+    log("\nIMPORTANT: Remember your password! There is NO password recovery!")
+    log("="*70)
+    input("\nPress ENTER to continue with encryption setup...")
 
-    if FAST_MODE:
-        veracrypt_create_cmd = (
-            f"veracrypt --text --create '{DOCS_PART}' "
-            f"--encryption AES "
-            f"--hash SHA-256 "
-            f"--filesystem exfat "
-            f"--volume-type normal "
-            f"--quick "
-        )
-    else:
-        veracrypt_create_cmd = (
-            f"veracrypt --text --create '{DOCS_PART}' "
-            f"--encryption AES-Twofish-Serpent "
-            f"--hash whirlpool "
-            f"--filesystem exfat "
-            f"--volume-type normal "
-        )
+    log("\nConfiguring VeraCrypt encryption for documents partition...")
+
+    # Always use maximum security for documents partition (ignore FAST_MODE)
+    # Triple cascade encryption with strongest hash
+    veracrypt_create_cmd = (
+        f"veracrypt --text --create '{DOCS_PART}' "
+        f"--encryption AES-Twofish-Serpent "
+        f"--hash SHA-512 "
+        f"--filesystem ext4 "
+        f"--volume-type normal "
+        f"--pim 2000 "  # Enforce strong PIM for maximum security
+    )
 
     run_command(veracrypt_create_cmd, shell=True, interactive=True)
     log("Encrypted documents partition setup complete.")
 
 def setup_unencrypted_partition():
+    """
+    Sets up the unencrypted partition for scripts and instructions.
+    """
     global DRIVE
     UNENCRYPTED_PART = get_partition_name(DRIVE, get_last_partition_number())
 
@@ -431,14 +630,19 @@ def setup_unencrypted_partition():
     run_command(f"sudo mount {UNENCRYPTED_PART} /mnt/unencrypted", shell=True)
 
     instructions = f"""
-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.
+INSTRUCTIONS FOR ACCESSING SECURE STORAGE
 
-**Automount Script:**
-Run: sudo ./mount_encrypted_partitions.sh
+To access your secure partition:
+  1. Run: sudo ./mount_storage.sh
+  2. Enter your password when prompted
+  3. Your files will be available at /mnt/secure_storage
 
-**Cleanup Script:**
-Run: sudo ./cleanup_encrypted_partitions.sh
+To safely lock your storage:
+  1. Close all files and applications using the storage
+  2. Run: sudo ./lock_storage.sh
+  3. Your data is now secured
+
+IMPORTANT: Always lock your storage before removing the device.
 """
 
     with open("/tmp/README.txt", "w") as readme_file:
@@ -448,60 +652,114 @@ Run: sudo ./cleanup_encrypted_partitions.sh
     log("Created README.txt with mounting instructions.")
 
     mount_script = """#!/bin/bash
-# Script to mount encrypted documents partition
+# Script to mount secure storage 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 "=== Secure Storage Mount ==="
+echo ""
+echo "Available storage devices:"
+lsblk -o NAME,SIZE,TYPE,LABEL | grep -E 'disk|part'
+echo ""
+read -p "Enter the partition path (e.g., /dev/sdb3): " DRIVE_PATH
 
-if [ -b "$DRIVE_PATH" ]; then
-    echo "Mounting VeraCrypt documents partition at $DRIVE_PATH..."
-    sudo mkdir -p /mnt/veracrypt_docs
-    sudo veracrypt --text --mount "$DRIVE_PATH" /mnt/veracrypt_docs
-    echo "Documents partition mounted at /mnt/veracrypt_docs."
+if [ -z "$DRIVE_PATH" ]; then
+    echo "Error: No partition specified."
+    exit 1
+fi
+
+if [ ! -b "$DRIVE_PATH" ]; then
+    echo "Error: $DRIVE_PATH is not a valid block device."
+    exit 1
+fi
+
+echo ""
+echo "Mounting secure storage from $DRIVE_PATH..."
+sudo mkdir -p /mnt/secure_storage
+
+# Mount the encrypted volume
+if sudo veracrypt --text --mount "$DRIVE_PATH" /mnt/secure_storage; then
+    echo ""
+    echo "SUCCESS: Secure storage is now accessible at /mnt/secure_storage"
+    echo ""
+    echo "Remember to run ./lock_storage.sh when finished!"
 else
-    echo "Error: Partition $DRIVE_PATH not found."
+    echo ""
+    echo "ERROR: Failed to mount storage. Check your password and try again."
+    sudo rmdir /mnt/secure_storage 2>/dev/null
+    exit 1
 fi
 """
 
-    with open("/tmp/mount_encrypted_partitions.sh", "w") as script_file:
+    with open("/tmp/mount_storage.sh", "w") as script_file:
         script_file.write(mount_script)
-    run_command("sudo cp /tmp/mount_encrypted_partitions.sh /mnt/unencrypted/mount_encrypted_partitions.sh", shell=True)
-    run_command("sudo chmod +x /mnt/unencrypted/mount_encrypted_partitions.sh", shell=True)
-    run_command("sudo rm /tmp/mount_encrypted_partitions.sh", shell=True)
-    log("Created mount_encrypted_partitions.sh script.")
+    run_command("sudo cp /tmp/mount_storage.sh /mnt/unencrypted/mount_storage.sh", shell=True)
+    run_command("sudo chmod +x /mnt/unencrypted/mount_storage.sh", shell=True)
+    run_command("sudo rm /tmp/mount_storage.sh", shell=True)
+    log("Created mount_storage.sh script.")
 
     cleanup_script = """#!/bin/bash
-# Script to unmount and lock encrypted documents partition
+# Script to safely lock secure storage
 
-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 "=== Secure Storage Lock ==="
+echo ""
 
-echo "Unmounting and locking encrypted documents partition at $DRIVE_PATH..."
+# Check if storage is mounted
+if mountpoint -q /mnt/secure_storage; then
+    echo "Checking for open files..."
 
-if mountpoint -q /mnt/veracrypt_docs; then
-    sudo veracrypt --text --dismount /mnt/veracrypt_docs
-    echo "VeraCrypt documents partition unmounted."
+    # Check if any processes are using the mount
+    if sudo lsof /mnt/secure_storage 2>/dev/null | grep -q /mnt/secure_storage; then
+        echo ""
+        echo "WARNING: Files are still open in the secure storage!"
+        echo "Open files/processes:"
+        sudo lsof /mnt/secure_storage 2>/dev/null
+        echo ""
+        read -p "Force close and lock anyway? (y/N): " FORCE
+        if [ "$FORCE" != "y" ] && [ "$FORCE" != "Y" ]; then
+            echo "Lock cancelled. Please close all files and try again."
+            exit 1
+        fi
+    fi
+
+    echo ""
+    echo "Locking secure storage..."
+
+    # Sync any pending writes
+    sync
+
+    # Dismount the encrypted volume
+    if sudo veracrypt --text --dismount /mnt/secure_storage; then
+        sudo rmdir /mnt/secure_storage 2>/dev/null
+        echo ""
+        echo "SUCCESS: Secure storage is now locked."
+        echo "Your data is safe."
+    else
+        echo ""
+        echo "ERROR: Failed to lock storage. Try closing all applications and run again."
+        exit 1
+    fi
 else
-    echo "Documents partition is not currently mounted."
+    echo "Secure storage is not currently mounted."
+    echo "Nothing to lock."
 fi
 """
 
-    with open("/tmp/cleanup_encrypted_partitions.sh", "w") as script_file:
+    with open("/tmp/lock_storage.sh", "w") as script_file:
         script_file.write(cleanup_script)
-    run_command("sudo cp /tmp/cleanup_encrypted_partitions.sh /mnt/unencrypted/cleanup_encrypted_partitions.sh", shell=True)
-    run_command("sudo chmod +x /mnt/unencrypted/cleanup_encrypted_partitions.sh", shell=True)
-    run_command("sudo rm /tmp/cleanup_encrypted_partitions.sh", shell=True)
-    log("Created cleanup_encrypted_partitions.sh script.")
+    run_command("sudo cp /tmp/lock_storage.sh /mnt/unencrypted/lock_storage.sh", shell=True)
+    run_command("sudo chmod +x /mnt/unencrypted/lock_storage.sh", shell=True)
+    run_command("sudo rm /tmp/lock_storage.sh", shell=True)
+    log("Created lock_storage.sh script.")
 
     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."""
+    """
+    Returns the highest partition number on the DRIVE.
+    
+    Returns:
+        int: The highest partition number.
+    """
     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 = []
@@ -524,20 +782,35 @@ def get_last_partition_number():
     return max(partition_numbers)
 
 def main():
-    global DEBUG, FAST_MODE, CREATE_KALI, CREATE_DOCS, CREATE_TAILS, KALI_ISO, TAILS_ISO, DRIVE
+    """
+    The main function that parses arguments and orchestrates the setup process.
+    """
+    global DEBUG, FAST_MODE, CREATE_KALI, CREATE_DOCS, CREATE_TAILS, KALI_ISO, TAILS_ISO, DRIVE, CREATE_CUSTOM, CUSTOM_ISO
+
+    CREATE_CUSTOM = False
+    CUSTOM_ISO = ""
 
     parser = argparse.ArgumentParser(description="Covert SD Card Tool")
-    parser.add_argument("-a", "--all", action="store_true", help="Set up both OS bootable USB and documents partition")
+    
+    # Creating a mutually exclusive group for -a and -t to prevent their simultaneous use
+    group = parser.add_mutually_exclusive_group()
+    group.add_argument("-a", "--all", action="store_true", help="Set up both OS bootable USB and documents partition")
+    group.add_argument("-t", "--tails", action="store_true", help="Create Tails bootable USB (no persistence)")
+    
+    # Other arguments remain outside the mutually exclusive group
     parser.add_argument("-k", "--kali", action="store_true", help="Create Kali bootable USB and persistence partition")
     parser.add_argument("-d", "--docs", action="store_true", help="Create encrypted documents partition")
-    parser.add_argument("-t", "--tails", action="store_true", help="Create Tails bootable USB (no persistence)")
-    parser.add_argument("-i", "--iso", help="Path to the Kali or Tails ISO file")
+    parser.add_argument("-c", "--custom", action="store_true", help="Create custom ISO bootable USB (like Kali, with persistence)")
+    parser.add_argument("-i", "--iso", help="Path to the ISO file (Kali, Tails, or custom)")
     parser.add_argument("--fast", action="store_true", help="Enable fast setup with less secure encryption")
     parser.add_argument("--debug", action="store_true", help="Enable debug mode")
 
     args = parser.parse_args()
 
-    if not any([args.all, args.kali, args.docs, args.tails]):
+    # Prevent using -a with -t; argparse handles this due to mutually exclusive group
+    # However, if you want to provide a custom error message or additional handling, you can add it here
+
+    if not any([args.all, args.kali, args.docs, args.tails, args.custom]):
         parser.print_help()
         sys.exit(1)
 
@@ -551,35 +824,24 @@ def main():
 
     if args.all:
         CREATE_DOCS = True
-        if args.tails:
-            CREATE_TAILS = True
-        else:
-            CREATE_KALI = True
+        # When -a is used without -t, default to creating Kali
+        CREATE_KALI = True
     else:
         CREATE_KALI = args.kali
         CREATE_DOCS = args.docs
         CREATE_TAILS = args.tails
+        CREATE_CUSTOM = args.custom
 
     if args.iso:
         if CREATE_KALI:
             KALI_ISO = args.iso
         elif CREATE_TAILS:
             TAILS_ISO = args.iso
+        elif CREATE_CUSTOM:
+            CUSTOM_ISO = args.iso
 
     check_dependencies()
-
-    if CREATE_KALI or CREATE_TAILS or CREATE_DOCS:
-        setup_usb()
-    else:
-        list_drives()
-        DRIVE = input("Enter the drive to use (e.g., /dev/sda) [Default: /dev/sda]: ") or "/dev/sda"
-        prepare_drive(DRIVE)
-
-        if CREATE_DOCS:
-            fix_partition_table_docs_only()
-        else:
-            setup_unencrypted_partition()
-
+    setup_usb()
     log("Partition setup complete.")
 
 if __name__ == "__main__":