app.py

from crypto import *


from flask import (
    Flask,
    render_template,
    request,
    send_file,
    flash,
    redirect,
    url_for,
    after_this_request
)

import os

from crypto.text import (
    encrypt_message,
    decrypt_message
)

from crypto.files import (
    encrypt_file,
    decrypt_file
)

app = Flask(__name__)

app.secret_key = os.environ.get(  #random key generated
    "SECRET_KEY",
    os.urandom(24)
)

UPLOAD_FOLDER = "uploads"

os.makedirs(UPLOAD_FOLDER, exist_ok=True)

app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER


# home route
@app.route("/")
def home():
    return render_template("index.html")

# text encryption route
@app.route("/encrypt", methods=["POST"])
def encrypt():
    if request.is_json:
        data = request.get_json() or {}
        message = data.get("message")
        password = data.get("password")
    else:
        message = request.form.get("message")
        password = request.form.get("password")

    if not message or not password:
        if request.is_json or request.headers.get("Accept") == "application/json":
            return {"error": "Message and Password are required."}, 400
        flash("Message and Password are required.")
        return redirect(url_for("home"))

    try:
        encrypted = encrypt_message(
            message,
            password
        )
        if request.is_json or request.headers.get("Accept") == "application/json":
            return {"encrypted_text": encrypted}
        return render_template(
            "index.html",
            encrypted_text=encrypted
        )

    except Exception as e:
        if request.is_json or request.headers.get("Accept") == "application/json":
            return {"error": str(e)}, 500
        flash(str(e))
        return redirect(url_for("home"))
    
# text decryption
@app.route("/decrypt", methods=["POST"])
def decrypt():
    if request.is_json:
        data = request.get_json() or {}
        token = data.get("encrypted_text")
        password = data.get("password")
    else:
        token = request.form.get("encrypted_text")
        password = request.form.get("password")

    if not token or not password:
        if request.is_json or request.headers.get("Accept") == "application/json":
            return {"error": "Encrypted text and password required."}, 400
        flash("Encrypted text and password required.")
        return redirect(url_for("home"))

    try:
        plaintext = decrypt_message(
            token,
            password
        )
        if request.is_json or request.headers.get("Accept") == "application/json":
            return {"decrypted_text": plaintext}
        return render_template(
            "index.html",
            decrypted_text=plaintext
        )

    except Exception:
        if request.is_json or request.headers.get("Accept") == "application/json":
            return {"error": "Wrong password or corrupted token."}, 400
        flash(
            "Wrong password or corrupted token."
        )
        return redirect(url_for("home"))
    
# file encryption route
@app.route("/encrypt-file", methods=["POST"])
def encrypt_uploaded_file():
    file = request.files.get("file")
    password = request.form.get("password")

    if not file or file.filename == "":
        flash("Select a file.")
        return redirect(url_for("home"))
    if not password:
        flash("Password is required.")
        return redirect(url_for("home"))

    input_path = os.path.join(
        app.config["UPLOAD_FOLDER"],
        file.filename
    )
    output_path = input_path + ".enc"

    try:
        file.save(input_path)
        encrypt_file(
            input_path,
            output_path,
            password
        )
        
        # Read the file into memory to release disk handles
        import io
        with open(output_path, "rb") as f:
            file_data = io.BytesIO(f.read())

        # Clean up disk files immediately
        if os.path.exists(input_path):
            try: os.remove(input_path)
            except: pass
        if os.path.exists(output_path):
            try: os.remove(output_path)
            except: pass

        return send_file(
            file_data,
            download_name=file.filename + ".enc",
            as_attachment=True
        )
    except Exception as e:
        flash(f"File encryption failed: {str(e)}")
        if os.path.exists(input_path):
            try: os.remove(input_path)
            except: pass
        if os.path.exists(output_path):
            try: os.remove(output_path)
            except: pass
        return redirect(url_for("home"))

#file decryption route
@app.route("/decrypt-file", methods=["POST"])
def decrypt_uploaded_file():
    file = request.files.get("file")
    password = request.form.get("password")

    if not file or file.filename == "":
        flash("Select a file.")
        return redirect(url_for("home"))
    if not password:
        flash("Password is required.")
        return redirect(url_for("home"))

    input_path = os.path.join(
        app.config["UPLOAD_FOLDER"],
        file.filename
    )

    if file.filename.endswith(".enc"):
        output_filename = file.filename[:-4]
    else:
        output_filename = file.filename + ".decrypted"

    output_path = os.path.join(
        app.config["UPLOAD_FOLDER"],
        output_filename
    )

    try:
        file.save(input_path)
        decrypt_file(
            input_path,
            output_path,
            password
        )
        
        # Read the file into memory to release disk handles
        import io
        with open(output_path, "rb") as f:
            file_data = io.BytesIO(f.read())

        # Clean up disk files immediately
        if os.path.exists(input_path):
            try: os.remove(input_path)
            except: pass
        if os.path.exists(output_path):
            try: os.remove(output_path)
            except: pass

        return send_file(
            file_data,
            download_name=output_filename,
            as_attachment=True
        )
    except Exception:
        flash("File decryption failed: Invalid password or corrupted file.")
        if os.path.exists(input_path):
            try: os.remove(input_path)
            except: pass
        if os.path.exists(output_path):
            try: os.remove(output_path)
            except: pass
        return redirect(url_for("home"))

# main function 

if __name__ == "__main__":
    app.run(debug=True)