app.py

import streamlit as st
import pandas as pd
import numpy as np
import joblib
from tensorflow import keras
import plotly.express as px
from datetime import datetime

# Load model and tools
model = keras.models.load_model('saved_model/threat_detection_model.keras')
scaler = joblib.load('saved_model/scaler.pkl')
label_encoder = joblib.load('saved_model/label_encoder.pkl')
with open('saved_model/feature_names.txt') as f:
    feature_names = f.read().splitlines()

# Streamlit setup
st.set_page_config(page_title="Network Threat Monitoring", layout="wide")
st.title("📡 Network Threat Monitoring Dashboard")

# Session state
if 'threats' not in st.session_state:
    st.session_state.threats = pd.DataFrame(columns=[
        'timestamp', 'source_ip', 'dest_ip', 'threat_class', 'threat_level', 'protocol'
    ])
if 'stats' not in st.session_state:
    st.session_state.stats = {'total_packets': 0, 'threat_packets': 0}

# Align features
def align_features(test_df, expected_features):
    aligned_df = pd.DataFrame(columns=expected_features)
    missing = []

    for feature in expected_features:
        if feature in test_df.columns:
            aligned_df[feature] = test_df[feature]
        else:
            aligned_df[feature] = 0
            missing.append(feature)

    if missing:
        st.warning(f"{len(missing)} features were missing and filled with 0.")
        with st.expander("🔍 View Missing Features"):
            st.write(missing)

    return aligned_df

# Analyze and predict
def analyze_dataframe(df):
    st.subheader("🔬 Analyzing Data...")

    original_attacks_present = None
    if 'attack_cat' in df.columns:
        df['attack_cat'] = df['attack_cat'].fillna('normal').str.strip().str.lower()
        dist = df['attack_cat'].value_counts().reset_index()
        dist.columns = ['Attack Category', 'Count']
        original_attacks_present = dist

        st.markdown("📊 *Original Attack Category Distribution (Before Resampling):*")
        st.dataframe(dist)
        st.plotly_chart(px.bar(dist, x='Attack Category', y='Count', title="Class Imbalance Overview"), use_container_width=True)

    if st.checkbox("🔎 Preview uploaded file features"):
        st.write(sorted(df.columns.tolist()))

    X = align_features(df, feature_names)
    X_scaled = scaler.transform(X)
    preds = model.predict(X_scaled)

    threat_classes = label_encoder.inverse_transform(np.argmax(preds, axis=1))
    threat_levels = np.max(preds, axis=1)

    # Ensure actual IP addresses from the file are used
    df_result = pd.DataFrame({
        'timestamp': [datetime.now()] * len(df),
        'source_ip': df['srcip'] if 'srcip' in df.columns else 'unknown',  # Use 'srcip' if present
        'dest_ip': df['dstip'] if 'dstip' in df.columns else 'unknown',  # Use 'dstip' if present
        'protocol': df['proto'] if 'proto' in df.columns else 'unknown',
        'threat_class': threat_classes,
        'threat_level': threat_levels
    })

    st.session_state.stats['total_packets'] = len(df_result)
    st.session_state.stats['threat_packets'] = (df_result['threat_class'] != 'normal').sum()
    st.session_state.threats = df_result

    if st.session_state.stats['threat_packets'] == len(df_result):
        st.warning("⚠ All records flagged as threats. This may indicate a mismatch between the data and expected model input.")

    display_dashboard(df_result, original_attacks_present)

# Threat breakdown
def display_dashboard(df_result, original_attacks=None):
    st.subheader("📊 Network Threat Analysis Dashboard")

    col1, col2 = st.columns(2)
    col1.metric("Total Packets Analyzed", f"{len(df_result):,}")
    col2.metric("Detected Threats", f"{(df_result['threat_class'] != 'normal').sum():,}")

    st.plotly_chart(px.line(
        df_result,
        x='timestamp',
        y='threat_level',
        title='Threat Level Over Time'
    ), use_container_width=True)

    with st.expander("📄 Full Threat Log"):
        st.dataframe(df_result.sort_values('threat_level', ascending=False), use_container_width=True)

    # 🧨 Detected categories summary
    st.subheader("🧨 Detected Threat Categories:")
    detected_threats = df_result[df_result['threat_class'] != 'normal']
    if detected_threats.empty:
        st.info("No non-normal threats detected.")
    else:
        summary = detected_threats.groupby('threat_class')
        for category, group in summary:
            st.markdown(f"#### • *{category}* – {len(group):,} packets")
            st.markdown(f"- Example Source IPs: {', '.join(group['source_ip'].unique()[:3])}")
            st.markdown(f"- Example Dest IPs: {', '.join(group['dest_ip'].unique()[:3])}")
            st.markdown("---")

    # Show original attack_cat from file
    if original_attacks is not None:
        st.subheader("📑 Original File Attack Categories Found:")
        known_threats = original_attacks[original_attacks['Attack Category'] != 'normal']
        if not known_threats.empty:
            for _, row in known_threats.iterrows():
                st.markdown(f"- *{row['Attack Category']}*: {row['Count']} rows")
        else:
            st.info("No labeled threats in 'attack_cat' column.")

# Sidebar and upload
st.sidebar.header("📥 Select Input Source")
input_method = st.sidebar.radio("Analyze data by:", ["Upload File", "Use System File"])

if input_method == "Upload File":
    uploaded_file = st.sidebar.file_uploader("Upload a .pcap or .csv file", type=["csv", "pcap"])

    if uploaded_file:
        file_name = uploaded_file.name
        st.info(f"Analyzing {file_name}...")

        if file_name.endswith('.csv'):
            try:
                df = pd.read_csv(uploaded_file)
                analyze_dataframe(df)
            except Exception as e:
                st.error(f"Error during analysis: {str(e)}")
        else:
            st.error("Only CSV format supported in this version. PCAP support coming soon!")

elif input_method == "Use System File":
    system_file_path = "UNSW_NB15_testing-set.csv"
    st.info(f"Analyzing CSV file: {system_file_path}")
    try:
        df = pd.read_csv(system_file_path)
        analyze_dataframe(df)
    except Exception as e:
        st.error(f"Error loading system file: {str(e)}")