dashboard.py

"""NetLens dashboard entrypoint."""

import time
import threading

import pandas as pd
import plotly.express as px
import streamlit as st
from scapy.all import sniff

from config import (
    APP_DESCRIPTION,
    APP_TITLE,
    NETWORK_SCAN_INTERVAL,
    STREAMLIT_PAGE_CONFIG,
)
from src.network_scanner import NetworkScanner
from src.packet_processor import PacketProcessor
from src.utils import setup_logging

logger = setup_logging()


def start_packet_capture() -> PacketProcessor:
    """Start packet capture in a background thread."""
    processor = PacketProcessor()

    def capture_packets() -> None:
        try:
            logger.info("[CAPTURE] Starting packet sniffer")
            sniff(prn=processor.process_packet, store=False)
        except PermissionError:
            logger.error("[CAPTURE] Permission denied - run with sudo")
            processor.error_msg = "Permission denied. Run with sudo."
        except OSError as exc:
            if "Permission" in str(exc):
                logger.error("[CAPTURE] Permission denied - run with sudo")
                processor.error_msg = "Permission denied. Run with sudo."
            else:
                logger.error(f"[CAPTURE] Error: {exc}")
                processor.error_msg = str(exc)
        except Exception as exc:
            logger.error(f"[CAPTURE] Error: {exc}")
            processor.error_msg = str(exc)

    threading.Thread(target=capture_packets, daemon=True).start()
    time.sleep(0.2)
    return processor


def start_network_scanner() -> NetworkScanner:
    """Start automatic network scanning in a background thread."""
    scanner = NetworkScanner()

    def scan_periodically() -> None:
        logger.info("[SCANNER] Performing initial network scan")
        result = scanner.scan()
        logger.info(f"[SCANNER] Initial scan complete - found {len(result)} devices")

        while True:
            try:
                time.sleep(NETWORK_SCAN_INTERVAL)
                logger.info("[SCANNER] Running periodic scan")
                result = scanner.scan()
                logger.info(f"[SCANNER] Periodic scan complete - found {len(result)} devices")
            except Exception as exc:
                logger.error(f"[SCANNER] Error in network scanner: {exc}")

    threading.Thread(target=scan_periodically, daemon=True).start()

    # Let the first scan begin before first render.
    time.sleep(2)
    return scanner


def _show_overall_view(df: pd.DataFrame, devices_df: pd.DataFrame, processor: PacketProcessor) -> None:
    """Render the overall network dashboard view."""
    st.subheader("🌐 Overall Network Statistics")

    col1, col2 = st.columns(2)

    with col1:
        st.write("**Top Devices by Bandwidth**")
        top_devices = devices_df.nlargest(10, "Bandwidth (MB)")[ ["Device", "Bandwidth (MB)"] ]
        fig_bw = px.bar(
            top_devices,
            x="Bandwidth (MB)",
            y="Device",
            orientation="h",
            title="Bandwidth Usage (MB)",
        )
        st.plotly_chart(fig_bw, width="stretch")

    with col2:
        st.write("**Top Devices by Packets**")
        top_packets = devices_df.nlargest(10, "Packets")[ ["Device", "Packets"] ]
        fig_packets = px.bar(
            top_packets,
            x="Packets",
            y="Device",
            orientation="h",
            title="Packets Sent",
        )
        st.plotly_chart(fig_packets, width="stretch")

    st.subheader("📊 Network Traffic Analysis")
    col1, col2, col3 = st.columns(3)

    with col1:
        if not df.empty and "protocol" in df.columns:
            protocol_counts = df["protocol"].value_counts()
            fig_protocol = px.pie(
                values=protocol_counts.values,
                names=protocol_counts.index,
                title="Protocol Distribution",
            )
            st.plotly_chart(fig_protocol, width="stretch")
        else:
            st.info("No protocol data yet.")

    with col2:
        services = processor.get_services_data()
        if len(services) > 0:
            top_services = services.head(10)
            fig_services = px.bar(
                x=top_services.values,
                y=top_services.index,
                orientation="h",
                title="Top Services",
            )
            st.plotly_chart(fig_services, width="stretch")
        else:
            st.info("No service data yet.")

    with col3:
        timestamps, mb_per_second = processor.get_bandwidth_data()
        if timestamps:
            fig_bandwidth = px.line(
                x=timestamps,
                y=mb_per_second,
                title="Bandwidth Over Time (MB/s)",
            )
            st.plotly_chart(fig_bandwidth, width="stretch")
        else:
            st.info("No bandwidth data yet.")

    st.subheader("📋 All Connected Devices")
    st.dataframe(devices_df, width="stretch", hide_index=True)


def _show_device_view(df: pd.DataFrame, processor: PacketProcessor, scanner: NetworkScanner, selected_ip: str) -> None:
    """Render the selected device dashboard view."""
    device_name = scanner.get_device_name(selected_ip)
    device_info = processor.device_stats.get(selected_ip, {})

    st.subheader(f"📱 Device: {device_name} ({selected_ip})")

    col1, col2, col3, col4 = st.columns(4)

    with col1:
        st.metric("Packets", device_info.get("packets", 0))

    with col2:
        bw_mb = device_info.get("bandwidth", 0) / (1024 * 1024)
        st.metric("Bandwidth", f"{bw_mb:.2f} MB")

    with col3:
        services = len(device_info.get("services", {}))
        st.metric("Services Used", services)

    with col4:
        domains = len(device_info.get("dns_queries", {}))
        st.metric("Unique Domains", domains)

    st.markdown("---")

    tab1, tab2, tab3, tab4 = st.tabs(
        [
            "🌐 Websites Visited",
            "🔌 Services",
            "📊 Traffic Timeline",
            "📋 Recent Packets",
        ]
    )

    with tab1:
        device_dns = device_info.get("dns_queries", {})
        if device_dns:
            dns_data = [
                {"domain": domain, "queries": count}
                for domain, count in sorted(device_dns.items(), key=lambda x: x[1], reverse=True)
            ]
            dns_df = pd.DataFrame(dns_data)

            fig_dns = px.bar(
                dns_df.head(20),
                x="queries",
                y="domain",
                orientation="h",
                title=f"Top 20 Websites Visited by {device_name}",
            )
            st.plotly_chart(fig_dns, width="stretch")

            st.write(f"**All {len(dns_data)} domains visited:**")
            st.dataframe(dns_df, width="stretch", hide_index=True)
        else:
            st.info("No DNS queries captured for this device yet.")

    with tab2:
        device_services = device_info.get("services", {})
        if device_services:
            services_data = [
                {"service": service, "packets": count}
                for service, count in sorted(device_services.items(), key=lambda x: x[1], reverse=True)
            ]
            services_df = pd.DataFrame(services_data)

            fig_svc = px.pie(
                services_df,
                values="packets",
                names="service",
                title=f"Services Used by {device_name}",
            )
            st.plotly_chart(fig_svc, width="stretch")
            st.dataframe(services_df, width="stretch", hide_index=True)
        else:
            st.info("No services captured for this device yet.")

    with tab3:
        if not df.empty and "source" in df.columns:
            device_packets = df[df["source"] == selected_ip].copy()
            if len(device_packets) > 0:
                device_packets["timestamp"] = pd.to_datetime(device_packets["timestamp"])
                device_grouped = device_packets.groupby(device_packets["timestamp"].dt.floor("s")).size()
                fig_timeline = px.line(
                    x=device_grouped.index,
                    y=device_grouped.values,
                    title=f"Packets per Second from {device_name}",
                    labels={"x": "Time", "y": "Packets/sec"},
                )
                st.plotly_chart(fig_timeline, width="stretch")
            else:
                st.info("No traffic captured for this device yet.")
        else:
            st.info("No traffic captured for this device yet.")

    with tab4:
        if not df.empty and "source" in df.columns:
            device_packets = df[df["source"] == selected_ip]
            if len(device_packets) > 0:
                display_cols = ["timestamp", "destination", "service", "protocol", "size", "dns_query"]
                available_cols = [col for col in display_cols if col in device_packets.columns]
                st.dataframe(
                    device_packets.tail(20)[available_cols],
                    width="stretch",
                    hide_index=True,
                )
            else:
                st.info("No packets captured for this device yet.")
        else:
            st.info("No packets captured for this device yet.")


def main() -> None:
    """Run the Streamlit dashboard."""
    st.set_page_config(**STREAMLIT_PAGE_CONFIG)

    st.title(APP_TITLE)
    st.markdown(f"**{APP_DESCRIPTION}**")

    if "processor" not in st.session_state:
        st.session_state.processor = start_packet_capture()
        st.session_state.scanner = start_network_scanner()
        st.session_state.start_time = time.time()
        st.session_state.selected_device_ip = None

    processor = st.session_state.processor
    scanner = st.session_state.scanner
    df = processor.get_dataframe()

    if processor.error_msg:
        st.error(f"⚠️ Capture Error: {processor.error_msg}")
        st.info("💡 Tip: Run with sudo for packet capture privilege")
        st.code(
            "sudo /home/vanitas/Desktop/network-dashboard/.venv/bin/streamlit run dashboard.py",
            language="bash",
        )
        return

    col1, col2, col3, col4, col5 = st.columns(5)

    with col1:
        st.metric("Total Packets", len(df))

    with col2:
        duration = time.time() - st.session_state.start_time
        st.metric("Capture Duration", f"{duration:.0f}s")

    with col3:
        devices = scanner.get_devices()
        online = sum(1 for d in devices.values() if d.get("status") == "online")
        total = len(devices)
        st.metric("🟢 Online Devices", f"{online}/{total}")

    with col4:
        dns_count = len(processor.dns_queries)
        st.metric("Unique Domains", dns_count)

    with col5:
        _, mb_per_second = processor.get_bandwidth_data()
        current_bw = mb_per_second[-1] if mb_per_second else 0
        st.metric("Current BW", f"{current_bw:.3f} MB/s")

    st.markdown("---")

    devices_df = processor.get_device_stats_dataframe(scanner)

    discovered = scanner.get_devices()
    if len(discovered) > 0:
        with st.expander("🔍 Network Discovery Status"):
            st.caption(f"Scanner found {len(discovered)} devices on {scanner.network_range}")
            discovery_info = []
            for ip, info in discovered.items():
                discovery_info.append(
                    {
                        "IP": ip,
                        "Hostname": info.get("hostname", "Unknown"),
                        "Status": info.get("status", "unknown"),
                        "Last Seen": info.get("last_seen", "N/A"),
                    }
                )
            if discovery_info:
                st.dataframe(pd.DataFrame(discovery_info), width="stretch", hide_index=True)

    if len(devices_df) == 0:
        st.warning("⏳ Scanning network... No devices found yet. This can take up to 30 seconds.")
        st.info(
            "💡 If devices still do not appear, click Scan Now or run `nmap -sn 192.168.1.0/24` manually."
        )
        return

    device_ips = ["overall"] + list(devices_df["IP"].values)
    device_options = ["📊 Overall Network"] + [
        f"{status} {name} ({ip})"
        for ip, name, status in zip(devices_df["IP"], devices_df["Device"], devices_df["Status"])
    ]

    current_index = 0
    selected_device_ip = st.session_state.get("selected_device_ip")
    if selected_device_ip is not None and selected_device_ip in device_ips:
        current_index = device_ips.index(selected_device_ip)

    col1, col2, col3 = st.columns([3, 1, 1])
    with col1:
        selected = st.selectbox(
            "📱 Select Device to Monitor",
            device_options,
            index=current_index,
            key="device_selector_main",
        )

    with col2:
        if st.button(
            "🔄 Scan Now",
            use_container_width=True,
            key="btn_refresh_devices",
            help="Manually trigger network scan",
        ):
            logger.info("[UI] User clicked Scan Now button")
            scanner.scan()
            st.rerun()

    with col3:
        if st.button(
            "📋 All Devices",
            use_container_width=True,
            key="btn_view_all",
            help="View all devices at once",
        ):
            st.session_state.selected_device_ip = None
            st.rerun()

    st.markdown("---")

    selected_ip = None
    if selected != "📊 Overall Network":
        selected_ip = selected.split("(")[-1].rstrip(")")
        st.session_state.selected_device_ip = selected_ip
    else:
        st.session_state.selected_device_ip = None

    if selected_ip is None:
        _show_overall_view(df, devices_df, processor)
    else:
        _show_device_view(df, processor, scanner, selected_ip)

    time.sleep(2)
    st.rerun()


if __name__ == "__main__":
    main()