app.py

import streamlit as st
import os
from dotenv import load_dotenv
from google import genai
from google.genai import types

# 1. LOAD ENVIRONMENT VARIABLES
load_dotenv()
api_key = os.getenv("GEMINI_API_KEY")

# 2. INITIALIZE GEMINI 3 CLIENT
client = genai.Client(api_key=api_key)

# 3. PAGE CONFIGURATION
st.set_page_config(
    page_title="Neural-Trace",
    layout="wide",
    page_icon="🛡️",
    initial_sidebar_state="collapsed"
)

# 4. DARK THEME CSS & ALIGNMENT FIXES
st.markdown("""
<style>
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700;800&display=swap');

/* Main App Container */
.stApp {
    background: #0f172a;
    color: #e2e8f0;
    font-family: 'Inter', sans-serif;
}

/* Remove default Streamlit padding */
.main .block-container {
    padding-top: 2rem !important;
    max-width: 1100px;
    margin: auto;
}

/* Hide Streamlit UI elements */
#MainMenu {visibility:hidden;}
footer {visibility:hidden;}
header[data-testid="stHeader"] {background:transparent;}
div[data-testid="stToolbar"] {display:none;}

/* Sticky Navbar */
.navbar {
    position: fixed;
    top: 0;
    left: 0;
    width: 100%;
    padding: 15px 50px;
    background: #020617;
    font-weight: 700;
    font-size: 22px;
    border-bottom: 1px solid #1e293b;
    z-index: 999;
    color: #f1f5f9;
}

/* Hero Section Alignment - Adjusted for Gap */
.badge-container {
    text-align: center;
    margin-top: 40px; /* Reduced to pull content up */
}
.badge {
    display: inline-block;
    padding: 6px 16px;
    border-radius: 999px;
    background: rgba(99,102,241,0.15);
    color: #818cf8;
    font-size: 13px;
    font-weight: 600;
    letter-spacing: 0.5px;
}

.hero {
    text-align: center;
    padding: 10px 10px;
}

.hero h1 {
    font-size: 52px;
    font-weight: 800;
    color: #f1f5f9;
    margin: 10px 0;
    line-height: 1.2;
}

.highlight {
    color: #818cf8;
}

/* Centered subtitle - Fixed Alignment and Spacing */
.subtitle-container {
    display: flex;
    justify-content: center;
    width: 100%;
    margin-bottom: 25px;
}

.subtitle {
    color: #94a3b8;
    font-size: 17px;
    max-width: 680px;
    text-align: center;
    line-height: 1.6;
    margin: 0;
    padding: 0 20px;
}

/* File Uploader styling */
[data-testid="stFileUploadDropzone"] {
    background: #020617 !important;
    border: 2px dashed #1e293b !important;
    border-radius: 16px !important;
    padding: 20px !important;
}

/* Button Styling - Centered */
.stButton > button {
    background: #6366f1 !important;
    color: white !important;
    border-radius: 10px !important;
    font-weight: 600 !important;
    border: none !important;
    width: 100%;
    transition: all 0.3s ease;
}

.stButton > button:hover {
    background: #4f46e5 !important;
    transform: translateY(-1px);
}

.report-title {
    margin-top: 40px;
    color: #818cf8;
    font-weight: 700;
}

/* Footer */
.footer {
    text-align: center;
    padding: 50px 0 20px 0;
    font-size: 14px;
    color: #64748b;
    border-top: 1px solid #1e293b;
    margin-top: 40px;
}
</style>
""", unsafe_allow_html=True)

# 5. UI LAYOUT COMPONENTS
st.markdown('<div class="navbar">🛡️ Neural-Trace</div>', unsafe_allow_html=True)

st.markdown("""
<div class="badge-container"><div class="badge">AI-POWERED BINARY SECURITY ANALYSIS</div></div>
<div class="hero">
    <h1>Security auditing trusted<br>by <span class="highlight">Neural-Trace</span></h1>
</div>
<div class="subtitle-container">
    <p class="subtitle">
        Analyze assembly binaries, detect vulnerabilities, and strengthen systems 
        through intelligent forensic auditing powered by Gemini 3.
    </p>
</div>
""", unsafe_allow_html=True)

# 6. SYSTEM PROMPT & ANALYSIS LOGIC
SYSTEM_PROMPT = """
You are a Senior Malware Analyst. 
Perform forensic audits on Assembly code.
Detect buffer overflows, integer overflows, logic bombs.
Provide threat level and secure Python rewrite.
"""

def analyze_asm(content):
    config = types.GenerateContentConfig(
        system_instruction=SYSTEM_PROMPT,
        thinking_config=types.ThinkingConfig(
            include_thoughts=True,
            thinking_level="HIGH"
        ),
        tools=[types.Tool(code_execution=types.ToolCodeExecution())]
    )
    prompt = f"Perform forensic audit:\n\n{content}"
    return client.models.generate_content(
        model="gemini-3-flash-preview",
        contents=prompt,
        config=config
    )

# 7. MAIN CONTENT AREA
_, col_mid, _ = st.columns([1, 2, 1])
with col_mid:
    uploaded_file = st.file_uploader(
        "Upload ASM/TXT",
        type=["asm","txt"],
        label_visibility="collapsed"
    )

if uploaded_file:
    asm_code = uploaded_file.read().decode("utf-8")
    
    st.markdown('<h3 class="report-title">📄 Uploaded Assembly Code</h3>', unsafe_allow_html=True)
    st.code(asm_code, language="asm", line_numbers=True)

    _, col_btn, _ = st.columns([1, 1, 1])
    with col_btn:
        analyze_btn = st.button("🚀 Start Forensic Audit")

    if analyze_btn:
        with st.spinner("Analyzing assembly code..."):
            try:
                report = analyze_asm(asm_code)
                st.markdown('<h3 class="report-title">🧠 AI Forensic Analysis Report</h3>', unsafe_allow_html=True)
                st.markdown(report.text)

                for part in report.candidates[0].content.parts:
                    if part.thought:
                        with st.expander("👁️ View AI Reasoning Chain", expanded=True):
                            st.info(part.text)
                st.success("Analysis completed.")
            except Exception as e:
                st.error(f"Analysis failed: {e}")
else:
    st.markdown("<br><center><p style='color: #64748b;'>Waiting for file upload to begin analysis...</p></center>", unsafe_allow_html=True)

# 8. FOOTER
st.markdown("""
<div class="footer">
    <strong>Neural-Trace</strong><br>
    Powered by Google Gemini 3 • Developed by Sutharshan Suthakaran
</div>
""", unsafe_allow_html=True)