App.py

import streamlit as st
import os
import tempfile
import traceback
from Secure_transmit import transmit_text, transmit_file, hide_in_pdf, hide_in_file
from Extractor import recover_from_image, recover_from_pdf, recover_from_file

st.set_page_config(page_title="StegMatrix", layout="wide")

st.title("🛡️ StegMatrix")
st.markdown("**Cybersecurity & Cloud Computing Specialization**")

with st.sidebar:
    st.header("🔑 AES Engine")
    PASS = st.text_input("Master Password", type="password", value="ismagi2026")

def save_upload(uploaded_file):
    if uploaded_file:
        suffix = os.path.splitext(uploaded_file.name)[1]
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            tmp.write(uploaded_file.getvalue())
            return tmp.name
    return None

tab1, tab2 = st.tabs(["📤 ENCODER (Run Scenarios A-I)", "🔓 DECODER (Extract Secrets)"])

# ==========================================
# TAB 1: THE ENCODER
with tab1:
    st.header("Select Your Exact Test Scenario")
    
    scenario = st.selectbox("Choose the scenario to execute:", [
        "A. HIDE A PHOTO INSIDE A PHOTO",
        "B. HIDE TYPED TEXT INSIDE A PHOTO",
        "C. HIDE A TEXT FILE INSIDE A PDF",
        "D. HIDE TYPED TEXT INSIDE A PDF",
        "E. HIDE A PDF INSIDE A TEXT FILE",
        "F. HIDE TYPED TEXT INSIDE A TEXT FILE",
        "G. HIDE A PHOTO INSIDE A TEXT FILE",
        "H. HIDE A PDF INSIDE A PHOTO",
        "I. HIDE A PDF INSIDE A PDF"
    ])

    st.markdown("---")
    
    col1, col2 = st.columns(2)
    carrier_file = None
    secret_file = None
    secret_text = ""

    # J'ai retiré les restrictions type=['...'] pour que tu puisses voir tous tes fichiers !
    with col1:
        if scenario in ["A. HIDE A PHOTO INSIDE A PHOTO", "B. HIDE TYPED TEXT INSIDE A PHOTO", "H. HIDE A PDF INSIDE A PHOTO"]:
            carrier_file = st.file_uploader("1. Upload Carrier PHOTO (PNG/JPG)")
        elif scenario in ["C. HIDE A TEXT FILE INSIDE A PDF", "D. HIDE TYPED TEXT INSIDE A PDF", "I. HIDE A PDF INSIDE A PDF"]:
            carrier_file = st.file_uploader("1. Upload Carrier PDF")
        elif scenario in ["E. HIDE A PDF INSIDE A TEXT FILE", "F. HIDE TYPED TEXT INSIDE A TEXT FILE", "G. HIDE A PHOTO INSIDE A TEXT FILE"]:
            carrier_file = st.file_uploader("1. Upload Carrier TEXT FILE")

    with col2:
        if "TYPED TEXT" in scenario:
            secret_text = st.text_area("2. Type your Secret Message")
        elif "PHOTO INSIDE" in scenario:
            secret_file = st.file_uploader("2. Upload Secret PHOTO")
        elif "PDF INSIDE" in scenario:
            secret_file = st.file_uploader("2. Upload Secret PDF")
        elif "TEXT FILE INSIDE" in scenario:
            secret_file = st.file_uploader("2. Upload Secret TEXT FILE")

    if st.button(f"🚀 Execute {scenario.split('.')[0]}"):
        if carrier_file and (secret_file or secret_text):
            c_path = save_upload(carrier_file)
            s_path = save_upload(secret_file) if secret_file else None
            base_folder = os.getcwd() 

            with st.spinner("Processing..."):
                try:
                    if scenario == "A. HIDE A PHOTO INSIDE A PHOTO":
                        output_name = os.path.join(base_folder, "image_in_image.png")
                        transmit_file(s_path, c_path, output_name, PASS)
                    elif scenario == "B. HIDE TYPED TEXT INSIDE A PHOTO":
                        output_name = os.path.join(base_folder, "text_in_image.png")
                        transmit_text(secret_text, c_path, output_name, PASS)
                    elif scenario == "C. HIDE A TEXT FILE INSIDE A PDF":
                        output_name = os.path.join(base_folder, "message_in_pdf.pdf")
                        hide_in_pdf(s_path, c_path, output_name, PASS, is_text=False)
                    elif scenario == "D. HIDE TYPED TEXT INSIDE A PDF":
                        output_name = os.path.join(base_folder, "text_in_pdf.pdf")
                        hide_in_pdf(secret_text, c_path, output_name, PASS, is_text=True)
                    elif scenario == "E. HIDE A PDF INSIDE A TEXT FILE":
                        output_name = os.path.join(base_folder, "pdf_hidden_in_text.txt")
                        hide_in_file(s_path, c_path, output_name, PASS, is_text=False)
                    elif scenario == "F. HIDE TYPED TEXT INSIDE A TEXT FILE":
                        output_name = os.path.join(base_folder, "hidden_in_text.txt")
                        hide_in_file(secret_text, c_path, output_name, PASS, is_text=True)
                    elif scenario == "G. HIDE A PHOTO INSIDE A TEXT FILE":
                        output_name = os.path.join(base_folder, "photo_hidden_in_text.txt")
                        hide_in_file(s_path, c_path, output_name, PASS, is_text=False)
                    elif scenario == "H. HIDE A PDF INSIDE A PHOTO":
                        output_name = os.path.join(base_folder, "pdf_in_photo.png")
                        transmit_file(s_path, c_path, output_name, PASS)
                    elif scenario == "I. HIDE A PDF INSIDE A PDF":
                        output_name = os.path.join(base_folder, "pdf_in_pdf.pdf")
                        hide_in_pdf(s_path, c_path, output_name, PASS, is_text=False)

                    if os.path.exists(output_name):
                        st.success(f"✅ Success! Generated: {os.path.basename(output_name)}")
                        with open(output_name, "rb") as f:
                            st.download_button("📥 Download Secure File", f, file_name=os.path.basename(output_name))
                    else:
                        st.error("⚠️ Error: File not created. (For 'PDF in Photo', ensure the Carrier Photo is large enough to hold the PDF!)")
                
                except Exception as e:
                    st.error(f"Error: {e}")
        else:
            st.error("Please provide both the Carrier and the Secret.")

# ==========================================
# TAB 2: THE DECODER (COMPLETELY FIXED)
# ==========================================
with tab2:
    st.header("Forensic Extraction")
    
    col3, col4 = st.columns(2)
    with col3:
        carrier_type = st.radio("1. What is the Carrier format?", ["Photo", "PDF", "Text File"])
    with col4:
        # LE SECRET : C'est ça qui va déterminer l'extension de sortie !
        secret_format = st.radio("2. What exactly are we extracting?", ["Typed Text", "Photo (PNG)", "Document (PDF)", "Text File (TXT)"])

    is_text = (secret_format == "Typed Text")
    stego_file = st.file_uploader("3. Upload the Suspect File (No format restrictions)")

    if st.button("🔍 Extract Secret") and stego_file:
        st_path = save_upload(stego_file)
        base_folder = os.getcwd() 
        
        # On donne la BONNE extension en fonction de ce qu'on cherche
        if secret_format == "Typed Text":
            out_name = os.path.join(base_folder, "recovered_secret_text.txt")
        elif secret_format == "Photo (PNG)":
            out_name = os.path.join(base_folder, "recovered_secret_image.png")
        elif secret_format == "Document (PDF)":
            out_name = os.path.join(base_folder, "recovered_secret_doc.pdf")
        elif secret_format == "Text File (TXT)":
            out_name = os.path.join(base_folder, "recovered_secret_file.txt")
        
        with st.spinner("Extracting..."):
            try:
                if carrier_type == "Photo":
                    recover_from_image(st_path, out_name, PASS, is_text=is_text)
                elif carrier_type == "PDF":
                    recover_from_pdf(st_path, out_name, PASS, is_text=is_text)
                elif carrier_type == "Text File":
                    recover_from_file(st_path, out_name, PASS, is_text=is_text)

                st.success("🔓 Cryptographic lock bypassed. Data recovered!")
                
                # LA CORRECTION POUR AFFICHER LE TEXTE À L'ÉCRAN
                if is_text:
                    if os.path.exists(out_name):
                        with open(out_name, "r", encoding="utf-8", errors="ignore") as f:
                            recovered_msg = f.read()
                        st.subheader("📝 Extracted Message:")
                        st.code(recovered_msg)
                    else:
                        st.warning("Could not read the file. Check terminal.")
                
                # TÉLÉCHARGEMENT POUR LES FICHIERS
                else:
                    if os.path.exists(out_name):
                        with open(out_name, "rb") as f:
                            st.download_button(f"💾 Download Extracted {secret_format}", f, file_name=os.path.basename(out_name))
                    else:
                        st.error("Extraction failed. File not found.")
            
            except Exception as e:
                # Si l'erreur est liée au mot de passe (InvalidToken)
                if "InvalidToken" in repr(e):
                    st.error("❌ Accès Refusé : Mot de passe incorrect ou signature cryptographique invalide.")
                    st.info("💡 Mécanisme de sécurité AES déclenché : La tentative de déchiffrement a été bloquée.")
                # Pour toutes les autres erreurs inconnues
                else:
                    st.error(f"⚠️ Échec de l'extraction : {repr(e)}")
                    st.warning("🔍 Détails techniques (Traceback) :")
                    st.code(traceback.format_exc())