Ragnar.py

#ragnar.py
# This script defines the main execution flow for the Ragnar application. It initializes and starts
# various components such as network scanning, display, and web server functionalities. The Ragnar 
# application serves as a comprehensive IoT security tool designed for network analysis and penetration testing.
# It integrates various modules to provide a unified platform for cybersecurity professionals and enthusiasts.

# Essential imports for the application
import asyncio
import os
import signal
import threading
import time
import atexit
try:
    import fcntl  # Unix only
except ImportError:
    fcntl = None
import logging
from collections import defaultdict

# ragnar.py


import threading
import signal
import logging
import time
import sys
import subprocess
from init_shared import shared_data
from display import Display, handle_exit_display
from comment import Commentaireia
from orchestrator import Orchestrator
from logger import Logger
from wifi_manager import WiFiManager
from env_manager import load_env

logger = Logger(name="Ragnar.py", level=logging.DEBUG)

class Ragnar:
    """Main class for Ragnar. Manages the primary operations of the application."""
    def __init__(self, shared_data):
        self.shared_data = shared_data
        self.commentaire_ia = Commentaireia()
        self.orchestrator_thread = None
        self.orchestrator = None
        self.wifi_manager = WiFiManager(shared_data)

        # Set reference to this instance in shared_data for other modules
        self.shared_data.ragnar_instance = self
        self.shared_data.headless_mode = False

        # Reference to display instance (will be set when display is started)
        self.display = None

        # Web portal lifecycle management (stop during wardriving without WiFi)
        self._web_portal_lock = threading.Lock()
        self._web_thread_ref = None  # Reference to current web server thread

        # PiSugar button listener (for Ragnar/Pwnagotchi swap via hardware button)
        self.pisugar_listener = None
        try:
            from pisugar_button import PiSugarButtonListener
            self.pisugar_listener = PiSugarButtonListener(shared_data)
        except ImportError:
            pass

    def run(self):
        """Main loop for Ragnar. Waits for Wi-Fi connection and starts Orchestrator."""
        logger.info("=" * 70)
        logger.info("RAGNAR MAIN THREAD STARTING")
        logger.info("=" * 70)
        
        # Start PiSugar button listener (if available)
        if self.pisugar_listener:
            self.pisugar_listener.start()

        # If wardriving-on-boot is enabled, kick wardriving off BEFORE the WiFi
        # manager starts. wifi_manager.start() blocks ~5s minimum (longer with
        # no AP available), and wardriving doesn't need WiFi connectivity —
        # it scans interfaces directly. Running it in a daemon thread also
        # lets us parallelize the GPS spin-up with WiFi startup, so the
        # e-paper sees `running == True` within ~1s of boot instead of 20+.
        wardriving_boot = self.shared_data.config.get('wardriving_on_boot', False)
        if wardriving_boot:
            logger.info("Wardriving-on-boot enabled. Starting wardriving early, orchestrator will be suppressed.")
            # Wardriving and On-Screen Network Diagnostic mode both hijack the
            # e-Paper display and the HAT keys, and the diagnostic layer wins the
            # render race (see display.py) — leaving Ragnar stuck on the netdiag
            # field-test screen even though wardriving is scanning underneath.
            # They are mutually exclusive, so if diagnostic mode was left on,
            # turn it off before wardriving takes over the panel.
            if self.shared_data.config.get('network_diagnostic_mode', False):
                logger.info("Disabling On-Screen Network Diagnostic mode: wardriving-on-boot takes the display.")
                self.shared_data.config['network_diagnostic_mode'] = False
                try:
                    self.shared_data.save_config()
                except Exception as e:
                    logger.warning(f"Could not persist network_diagnostic_mode off at boot: {e}")
            self.shared_data.manual_mode = True
            threading.Thread(
                target=self._start_wardriving_on_boot,
                daemon=True,
                name="wardriving-on-boot",
            ).start()

        # Initialize Wi-Fi management system
        logger.info("Starting Wi-Fi management system...")
        self.wifi_manager.start()
        logger.info("Wi-Fi management system started")

        # Main loop to keep Ragnar running
        logger.info("Entering main Ragnar loop...")

        loop_count = 0
        while not self.shared_data.should_exit:
            loop_count += 1
            if loop_count % 6 == 1:  # Log every 60 seconds (6 iterations * 10 sec)
                logger.info(f"Ragnar main loop iteration {loop_count}, manual_mode={self.shared_data.manual_mode}")
            
            if not self.shared_data.manual_mode:
                self.check_and_start_orchestrator()
            # Sleep in 1-second chunks so should_exit is checked quickly
            for _ in range(10):
                if self.shared_data.should_exit:
                    break
                time.sleep(1)
        
        logger.info("Ragnar main loop exited")



    def check_and_start_orchestrator(self):
        """Check Wi-Fi and start the orchestrator if connected."""
        wifi_connected = self.wifi_manager.check_wifi_connection()
        logger.debug(f"WiFi connection check: {wifi_connected}")
        
        if wifi_connected:
            self.shared_data.wifi_connected = True
            if self.orchestrator_thread is None or not self.orchestrator_thread.is_alive():
                logger.info("WiFi detected - attempting to start Orchestrator...")
                self.start_orchestrator()
        else:
            self.shared_data.wifi_connected = False
            if not self.wifi_manager.startup_complete:
                logger.info("Waiting for Wi-Fi management system to complete startup...")
            else:
                logger.debug("Waiting for Wi-Fi connection to start Orchestrator...")

    def start_orchestrator(self):
        """Start the orchestrator thread."""
        # Always clear manual mode so the main loop will auto-start
        # the orchestrator when Wi-Fi becomes available
        self.shared_data.manual_mode = False
        self.shared_data.orchestrator_should_exit = False

        # Use Wi-Fi manager's connection check
        if self.wifi_manager.check_wifi_connection():
            self.shared_data.wifi_connected = True
            if self.orchestrator_thread is None or not self.orchestrator_thread.is_alive():
                logger.info("Starting Orchestrator thread...")
                self.orchestrator = Orchestrator()
                self.orchestrator_thread = threading.Thread(target=self.orchestrator.run)
                self.orchestrator_thread.start()
                logger.info("Orchestrator thread started, automatic mode activated.")
            else:
                logger.info("Orchestrator thread is already running.")
        else:
            logger.warning("Cannot start Orchestrator yet: Wi-Fi is not connected. Will auto-start when connected.")

    def stop_orchestrator(self):
        """Stop the orchestrator thread."""
        self.shared_data.manual_mode = True
        logger.info("Stop button pressed. Manual mode activated & Stopping Orchestrator...")
        if self.orchestrator_thread is not None and self.orchestrator_thread.is_alive():
            logger.info("Stopping Orchestrator thread...")
            self.shared_data.orchestrator_should_exit = True
            self.orchestrator_thread.join(timeout=3)
            logger.info("Orchestrator thread stopped.")
            self.shared_data.ragnarorch_status = "IDLE"
            self.shared_data.ragnarstatustext2 = ""
            self.shared_data.manual_mode = True

    def _start_wardriving_on_boot(self):
        """Auto-start wardriving engine at boot (called when wardriving_on_boot is True).

        Runs in a daemon thread; never block the main loop. GPSManager has its
        own internal reconnect loop so we only need a single start() — if the
        USB GPS isn't enumerated yet, scanning still proceeds and GPS attaches
        as soon as the device shows up.
        """
        try:
            from wardriving import WardrivingEngine
            self._wd_engine = WardrivingEngine(self.shared_data)
            device_name = self.shared_data.config.get('wardriving_device_name', 'Ragnar')
            gps_port = self.shared_data.config.get('wardriving_gps_port', '') or None
            result = self._wd_engine.start(device_name=device_name, gps_port=gps_port)
            logger.info(f"Wardriving auto-started on boot: {result}")
            # Store reference so webapp can find it
            import webapp_modern
            webapp_modern._wardriving_engine = self._wd_engine
        except Exception as e:
            logger.error(f"Failed to auto-start wardriving on boot: {e}")

    def start_web_portal(self):
        """Start the web portal if not already running (called when WiFi connects during wardriving)."""
        if not self.shared_data.config.get('websrv', True):
            return False
        with self._web_portal_lock:
            if self._web_thread_ref and self._web_thread_ref.is_alive():
                logger.debug("Web portal already running, skipping start")
                return False
            logger.info("Starting web portal (WiFi connected during wardriving)...")
            self.shared_data.webapp_should_exit = False
            from webapp_modern import run_server
            t = threading.Thread(target=run_server, daemon=True, name="web-server")
            t.start()
            self._web_thread_ref = t
            self.shared_data.web_portal_active = True
            logger.info("Web portal started")
            return True

    def stop_web_portal(self):
        """Stop the web portal to save memory (called when WiFi drops during wardriving)."""
        with self._web_portal_lock:
            if not self._web_thread_ref or not self._web_thread_ref.is_alive():
                logger.debug("Web portal not running, skipping stop")
                self.shared_data.web_portal_active = False
                return False
            logger.info("Stopping web portal to save memory (wardriving without WiFi)...")
            self.shared_data.webapp_should_exit = True
            try:
                import webapp_modern
                webapp_modern.socketio.stop()
            except Exception as e:
                logger.warning(f"socketio.stop() error: {e}")
            self._web_thread_ref.join(timeout=5)
            self.shared_data.web_portal_active = False
            logger.info("Web portal stopped")
            return True

    def stop(self):
        """Stop Ragnar and cleanup all resources."""
        logger.info("Stopping Ragnar...")

        # Stop PiSugar listener
        if self.pisugar_listener:
            self.pisugar_listener.stop()

        # Stop orchestrator
        self.stop_orchestrator()

        # Stop Wi-Fi manager
        if hasattr(self, 'wifi_manager'):
            self.wifi_manager.stop()
        
        # Set exit flags
        self.shared_data.should_exit = True
        self.shared_data.orchestrator_should_exit = True
        self.shared_data.display_should_exit = True
        self.shared_data.webapp_should_exit = True
        
        logger.info("Ragnar stopped successfully")

    def is_wifi_connected(self):
        """Legacy method - use wifi_manager for new code."""
        if hasattr(self, 'wifi_manager'):
            return self.wifi_manager.check_wifi_connection()
        else:
            # Fallback to original method
            result = subprocess.Popen(['nmcli', '-t', '-f', 'active', 'dev', 'wifi'], stdout=subprocess.PIPE, text=True).communicate()[0]
            return 'yes' in result

    
    @staticmethod
    def start_display():
        """Start the display thread"""
        display = Display(shared_data)
        display_thread = threading.Thread(target=display.run)
        display_thread.start()
        
        # Store display instance in shared_data for access by other modules
        shared_data.display_instance = display
        
        return display_thread

def handle_exit(sig, frame, display_thread, ragnar_thread, web_thread):
    """Handles the termination of the main, display, and web threads."""
    logger.info("Received exit signal, initiating clean shutdown...")

    # Stop Ragnar instance first
    if hasattr(shared_data, 'ragnar_instance') and shared_data.ragnar_instance:
        shared_data.ragnar_instance.stop()

    # Set all exit flags
    shared_data.should_exit = True
    shared_data.orchestrator_should_exit = True
    shared_data.display_should_exit = True
    shared_data.webapp_should_exit = True

    # Encrypt database on shutdown if auth is configured
    try:
        from webapp_modern import auth_mgr
        auth_mgr.shutdown_encrypt()
    except Exception as e:
        logger.error(f"Shutdown encryption failed: {e}")

    # Stop individual threads (fast timeouts - systemd will SIGKILL after 5s anyway)
    handle_exit_display(sig, frame, display_thread, exit_process=False)

    if display_thread and display_thread.is_alive():
        display_thread.join(timeout=1)
    if ragnar_thread and ragnar_thread.is_alive():
        ragnar_thread.join(timeout=1)
    if web_thread and web_thread.is_alive():
        web_thread.join(timeout=1)

    logger.info("Main loop finished. Clean exit.")
    sys.exit(0)



def _atexit_encrypt():
    """Safety net: encrypt DB on process exit if auth is configured."""
    try:
        from webapp_modern import auth_mgr
        auth_mgr.shutdown_encrypt()
    except Exception:
        pass

atexit.register(_atexit_encrypt)

if __name__ == "__main__":
    # Load environment variables from .env file at the very beginning
    load_env()

    logger.info("Starting threads")

    try:
        logger.info("Loading shared data config...")
        shared_data.load_config()

        # Always reset pwnagotchi_mode on Ragnar startup so Ragnar
        # is the canonical mode after every boot / service restart.
        shared_data.config['pwnagotchi_mode'] = 'ragnar'
        shared_data.config['pwnagotchi_last_status'] = 'Ragnar service is running'
        shared_data.save_config()

        # Clean up leftover pwnagotchi state (mon0, services)
        # Run all cleanup commands in parallel to avoid sequential timeouts
        logger.info("Cleaning up leftover pwnagotchi state...")
        cleanup_cmds = [
            (['ip', 'link', 'set', 'mon0', 'down'], 5),
            (['iw', 'mon0', 'del'], 5),
            (['systemctl', 'stop', 'pwnagotchi'], 10),
            (['systemctl', 'stop', 'bettercap'], 10),
            (['systemctl', 'stop', 'ragnar-swap-button'], 5),
        ]
        cleanup_procs = []
        for cmd, _ in cleanup_cmds:
            try:
                cleanup_procs.append(
                    (cmd, subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL))
                )
            except Exception:
                pass
        # Wait for all with a single combined timeout (10s max instead of 30s sequential)
        for cmd, proc in cleanup_procs:
            try:
                proc.wait(timeout=10)
            except Exception:
                try:
                    proc.kill()
                except Exception:
                    pass

        # wipe_epd stays as ExecStartPre (separate process) to avoid GPIO conflicts
        # with the Display's EPDHelper instance that shares the same pins.

        # Start the web server first so it becomes available even if the
        # EPD display initialisation is slow (e-paper displays can take
        # 10-30 seconds to initialise, especially larger panels like the 4.26").
        if shared_data.config["websrv"]:
            logger.info("Starting the web server...")
            from webapp_modern import run_server
            web_thread = threading.Thread(target=run_server, daemon=True, name="web-server")
            web_thread.start()
        else:
            web_thread = None

        logger.info("Starting display thread...")
        shared_data.display_should_exit = False  # Initialize display should_exit
        display_thread = Ragnar.start_display()

        logger.info("Starting Ragnar thread...")
        ragnar = Ragnar(shared_data)
        shared_data.ragnar_instance = ragnar  # Assigner l'instance de Ragnar à shared_data

        # Link display instance to ragnar instance
        if hasattr(shared_data, 'display_instance'):
            ragnar.display = shared_data.display_instance

        # Register web thread with Ragnar instance for portal lifecycle management
        ragnar._web_thread_ref = web_thread
        shared_data.web_portal_active = web_thread is not None

        ragnar_thread = threading.Thread(target=ragnar.run)
        ragnar_thread.start()

        signal.signal(signal.SIGINT, lambda sig, frame: handle_exit(sig, frame, display_thread, ragnar_thread, web_thread))
        signal.signal(signal.SIGTERM, lambda sig, frame: handle_exit(sig, frame, display_thread, ragnar_thread, web_thread))

    except Exception as e:
        import traceback
        logger.error(f"An exception occurred during thread start: {e}")
        # Log the full traceback — a bare str(e) (e.g. "cannot import name
        # 'run_server' from 'webapp_modern'") hides where the failure actually
        # happened, which turns a two-minute diagnosis into a guessing game.
        logger.error("Full traceback:\n" + traceback.format_exc())
        if 'display_thread' in locals():
            handle_exit_display(signal.SIGINT, None, display_thread)
        exit(1)