from PyPDF2 import PdfReader, PdfWriter

def encrypt_pdf(input_pdf, output_pdf, password):
    """Encrypts a PDF file with a password."""
    reader = PdfReader(input_pdf)
    writer = PdfWriter()

    for page in reader.pages:
        writer.add_page(page)

    writer.encrypt(password)

    with open(output_pdf, "wb") as f:
        writer.write(f)

    print(f"Encrypted PDF saved at {output_pdf}")

def decrypt_pdf(input_pdf, output_pdf, password):
    """Decrypts a password-protected PDF file."""
    reader = PdfReader(input_pdf)
    
    if reader.is_encrypted:
        reader.decrypt(password)

    writer = PdfWriter()
    
    for page in reader.pages:
        writer.add_page(page)

    with open(output_pdf, "wb") as f:
        writer.write(f)

    print(f"Decrypted PDF saved at {output_pdf}")
