security.py

import os
import hashlib
from Crypto.Cipher import AES
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Random import get_random_bytes
from Crypto.Util.Padding import pad, unpad
from Crypto.Hash import SHA256
import getpass

# NEW: imports for optional GUI browse and config-based watch dir check
import tkinter as tk  # NEW
from tkinter import filedialog  # NEW
import configparser  # NEW


def encrypt_file(file_path, password):
    """
    Encrypts a file using AES-GCM with PBKDF2 key derivation.

    Args:
        file_path (str): Path to the file to encrypt
        password (str): Password for encryption

    Returns:
        str: Path to the encrypted file (.enc extension)

    Raises:
        FileNotFoundError: If the input file doesn't exist
        Exception: For other encryption errors
    """
    try:
        # Check if file exists
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"File not found: {file_path}")

        # Read the original file in binary mode
        with open(file_path, 'rb') as f:
            plaintext = f.read()

        # Generate a random salt (16 bytes)
        salt = get_random_bytes(16)

        # Derive key using PBKDF2
        # Using 100,000 iterations for strong security
        key = PBKDF2(password, salt, dkLen=32, count=100000, hmac_hash_module=SHA256)

        # Create AES-GCM cipher
        cipher = AES.new(key, AES.MODE_GCM)
        nonce = cipher.nonce  # GCM automatically generates a nonce

        # Encrypt the data and get authentication tag
        ciphertext, tag = cipher.encrypt_and_digest(plaintext)

        # Create encrypted file path
        encrypted_file_path = file_path + '.enc'

        # Write encrypted data to file
        # Format: salt (16 bytes) + nonce (16 bytes) + tag (16 bytes) + ciphertext
        with open(encrypted_file_path, 'wb') as f:
            f.write(salt)
            f.write(nonce)
            f.write(tag)
            f.write(ciphertext)

        os.remove(file_path)

        print(f"File encrypted successfully: {encrypted_file_path}")
        return encrypted_file_path

    except FileNotFoundError:
        raise
    except Exception as e:
        raise Exception(f"Encryption failed: {str(e)}")


def decrypt_file(encrypted_file_path, password, delete_encrypted=True):
    """
    Decrypts a file that was encrypted using encrypt_file function.

    Args:
        encrypted_file_path (str): Path to the encrypted file (.enc)
        password (str): Password for decryption
        delete_encrypted (bool): Whether to delete the encrypted file after decryption

    Returns:
        str: Path to the decrypted file

    Raises:
        FileNotFoundError: If the encrypted file doesn't exist
        ValueError: If password is incorrect or file is corrupted
        Exception: For other decryption errors
    """
    try:
        # Check if encrypted file exists
        if not os.path.exists(encrypted_file_path):
            raise FileNotFoundError(f"Encrypted file not found: {encrypted_file_path}")

        # Read the encrypted file
        with open(encrypted_file_path, 'rb') as f:
            encrypted_data = f.read()

        # Verify minimum file size (salt + nonce + tag = 48 bytes minimum)
        if len(encrypted_data) < 48:
            raise ValueError("Invalid encrypted file format")

        # Extract components in the correct order
        salt = encrypted_data[:16]  # First 16 bytes: salt
        nonce = encrypted_data[16:32]  # Next 16 bytes: nonce
        tag = encrypted_data[32:48]  # Next 16 bytes: authentication tag
        ciphertext = encrypted_data[48:]  # Remaining bytes: encrypted data

        # Re-derive the key using the same parameters
        from Crypto.Hash import SHA256
        key = PBKDF2(password, salt, dkLen=32, count=100000, hmac_hash_module=SHA256)

        # Create AES-GCM cipher for decryption
        cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)

        # Attempt to decrypt and verify authenticity
        try:
            plaintext = cipher.decrypt_and_verify(ciphertext, tag)
        except ValueError as e:
            raise ValueError("Decryption failed: Wrong password or corrupted file")

        # Verify decryption result
        if plaintext is None:
            raise ValueError("Decryption returned None")
        if len(plaintext) == 0:
            raise ValueError("Decryption returned empty data")

        # Determine output file path (remove .enc extension)
        if encrypted_file_path.endswith('.enc'):
            decrypted_file_path = encrypted_file_path[:-4]  # Remove .enc
        else:
            decrypted_file_path = encrypted_file_path + '.decrypted'

        # Ensure target directory exists and is writable
        target_dir = os.path.dirname(os.path.abspath(decrypted_file_path))
        if not os.access(target_dir, os.W_OK):
            raise Exception(f"No write permission to directory: {target_dir}")

        # Write decrypted data to file with verification
        try:
            with open(decrypted_file_path, 'wb') as f:
                bytes_written = f.write(plaintext)
                f.flush()  # Force buffer to disk
                os.fsync(f.fileno())  # Force OS to write to disk

            # Immediate verification
            if not os.path.exists(decrypted_file_path):
                raise Exception("File creation failed - file does not exist after write")

            file_size = os.path.getsize(decrypted_file_path)
            if file_size != len(plaintext):
                raise Exception(f"File size mismatch: expected {len(plaintext)}, got {file_size}")

            if file_size == 0:
                raise Exception("Created file is empty")

        except Exception as write_err:
            raise Exception(f"File write operation failed: {write_err}")

        # Only delete encrypted file after successful verification
        if delete_encrypted:
            try:
                os.remove(encrypted_file_path)
            except Exception as del_err:
                print(f"Warning: Could not delete encrypted file: {del_err}")

        print(f"File decrypted successfully: {decrypted_file_path}")
        return decrypted_file_path

    except FileNotFoundError:
        raise
    except ValueError:
        raise
    except Exception as e:
        raise Exception(f"Decryption failed: {str(e)}")


# NEW: helper to get watch directory from config.ini (if present)
def _get_watch_directory():
    cfg = configparser.ConfigParser()
    if not os.path.exists("config.ini"):
        return None
    try:
        cfg.read("config.ini")
        if cfg.has_option("Settings", "watch_directory"):
            return cfg.get("Settings", "watch_directory")
    except Exception:
        return None
    return None


# NEW: helper to check if candidate path is inside base path
def _is_inside(base_dir, candidate_path):
    if not base_dir:
        return False
    base = os.path.abspath(base_dir)
    cand = os.path.abspath(candidate_path)
    try:
        return os.path.commonpath([base, cand]) == base
    except Exception:
        return False


# NEW: quoted-path cleaner for pasted paths (handles Windows “Copy as path”)
import os
import re
import unicodedata

def _clean_path(p: str) -> str:
    """
    Cleans a Windows path pasted from 'Copy as Path' or manual input.
    Removes surrounding quotes (", ', “, ”, etc.), trims whitespace,
    expands environment variables (~, %USERPROFILE%), and normalizes the path.
    """
    if not p:
        return p

    # Normalize Unicode (handles smart quotes and special spaces)
    p = unicodedata.normalize("NFKC", p)

    # Strip leading/trailing spaces and hidden chars (zero-width, BOM)
    p = p.strip().replace("\u200b", "").replace("\ufeff", "")

    # Remove wrapping quotes of any type
    p = re.sub(r'^[\'\"“”‘’]+|[\'\"“”‘’]+$', '', p)

    # Expand user (~) and environment (%VAR%)
    p = os.path.expanduser(os.path.expandvars(p))

    # Normalize to absolute Windows path
    return os.path.normpath(p)


# NEW: decrypt directly to a chosen output path (avoids writing inside watched dir)
def decrypt_file_to(encrypted_file_path, password, output_path, delete_encrypted=True):
    """
    Decrypts encrypted_file_path and writes plaintext directly to output_path.
    """
    try:
        if not os.path.exists(encrypted_file_path):
            raise FileNotFoundError(f"Encrypted file not found: {encrypted_file_path}")

        with open(encrypted_file_path, 'rb') as f:
            encrypted_data = f.read()

        if len(encrypted_data) < 48:
            raise ValueError("Invalid encrypted file format")

        salt = encrypted_data[:16]
        nonce = encrypted_data[16:32]
        tag = encrypted_data[32:48]
        ciphertext = encrypted_data[48:]

        key = PBKDF2(password, salt, dkLen=32, count=100000, hmac_hash_module=SHA256)
        cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)

        try:
            plaintext = cipher.decrypt_and_verify(ciphertext, tag)
        except ValueError:
            raise ValueError("Decryption failed: Wrong password or corrupted file")

        if plaintext is None or len(plaintext) == 0:
            raise ValueError("Decryption produced no data")

        # Ensure destination directory exists and writable
        target_dir = os.path.dirname(os.path.abspath(output_path)) or "."
        os.makedirs(target_dir, exist_ok=True)
        if not os.access(target_dir, os.W_OK):
            raise Exception(f"No write permission to directory: {target_dir}")

        with open(output_path, 'wb') as f:
            f.write(plaintext)
            f.flush()
            os.fsync(f.fileno())

        if not os.path.exists(output_path):
            raise Exception("Output write failed")

        if delete_encrypted:
            try:
                os.remove(encrypted_file_path)
            except Exception as del_err:
                print(f"Warning: Could not delete encrypted file: {del_err}")

        print(f"File decrypted successfully: {output_path}")
        return output_path

    except Exception as e:
        raise Exception(f"Decryption failed: {str(e)}")


# NEW: small UI helpers for browsing
def _pick_restore_file(default_name):
    root = tk.Tk()
    root.withdraw()
    path = filedialog.asksaveasfilename(
        title="Save restored file as",
        initialfile=default_name,
        defaultextension="",
        filetypes=[("All Files", "*.*")]
    )
    root.destroy()
    return path or ""


def main():
    """
    Simple command-line interface for testing the encryption/decryption functions.
    """
    print("=== File Encryption/Decryption Tool ===")
    # Loop until user chooses Exit
    while True:
        print("1. Encrypt a file")
        print("2. Decrypt a file")
        print("3. Exit")
        choice = input("Enter your choice (1, 2, or 3): ").strip()

        if choice == '1':
            file_path = _clean_path(input("Enter path to file to encrypt:"))
            password = getpass.getpass("Enter password: ")
            try:
                encrypted_path = encrypt_file(file_path, password)
                print(f"Success! Encrypted file created: {encrypted_path}")
            except Exception as e:
                print(f"Error: {e}")

        elif choice == '2':
            encrypted_path = _clean_path(input("Enter path to encrypted file (.enc):"))
            password = getpass.getpass("Enter password: ")

            # NEW: offer to browse a safe output location outside the watch dir
            try:
                # Derive default restored name
                if encrypted_path.lower().endswith(".enc"):
                    default_name = os.path.basename(encrypted_path)[:-4]
                else:
                    default_name = os.path.basename(encrypted_path) + ".restored"

                use_browse = input("Browse output location? (y/N): ").strip().lower() == "y"
                if use_browse:
                    out_path = _pick_restore_file(default_name)
                    if not out_path:
                        print("No output selected. Decryption cancelled.")
                        continue

                    watch_dir = _get_watch_directory()
                    if _is_inside(watch_dir, out_path):
                        print("Selected path is inside the watched folder; pick a different location.")
                        continue

                    # Decrypt directly to chosen path
                    decrypted_path = decrypt_file_to(encrypted_path, password, out_path)
                    print(f"Success! Decrypted file created: {decrypted_path}")
                else:
                    # Fall back to original behavior (writes next to .enc file)
                    decrypted_path = decrypt_file(encrypted_path, password)
                    print(f"Success! Decrypted file created: {decrypted_path}")

            except Exception as e:
                print(f"Error: {e}")

        elif choice == '3':
            print("Exiting. Goodbye.")
            break

        else:
            print("Invalid choice!")


if __name__ == "__main__":
    main()