startup.py

import os, sys, subprocess, time, hashlib, threading, signal
from source.config import config

ROOT_DIR = os.path.dirname(os.path.abspath(__file__))

# Global process tracking
server_process = None
scanner_process = None
cloudflared_process = None
shutdown_flag = threading.Event()  # Flag to signal shutdown

def get_rust_src_hash():
    hasher = hashlib.sha256()
    
    def hash_file(path):
        try:
            with open(path, 'rb') as f:
                while chunk := f.read(65536):
                    hasher.update(chunk)
        except: pass

    hash_file(os.path.join("rust_core", "Cargo.toml"))
    
    src_dir = os.path.join("rust_core", "src")
    for root, dirs, files in os.walk(src_dir):
        dirs.sort()
        for file in sorted(files):
            if file.endswith(".rs"):
                path = os.path.join(root, file)
                hash_file(path)
                
    return hasher.hexdigest()

def build_rust():
    print("Checking Rust Core integrity...")
    
    current_hash = get_rust_src_hash()
    stored_hash = ""
    
    if os.path.exists("rust_core/.build_hash"):
        try:
            with open("rust_core/.build_hash", 'r') as f:
                stored_hash = f.read().strip()
        except: pass
        
    is_importable = False
    try:
        import medal_diff_core
        is_importable = True
    except ImportError:
        pass
        
    if is_importable and current_hash == stored_hash:
        print("Rust Core is up-to-date. Skipping build.")
        return

    print("Building Native Rust Module...") 
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "show", "maturin"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        subprocess.check_call([sys.executable, "-m", "maturin", "build", "--release"], cwd="rust_core", stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

        wheels_dir = os.path.join("rust_core", "target", "wheels")
        if not os.path.exists(wheels_dir):
            print("Build failed: No wheels directory found.")
            return 
            
        wheels = [f for f in os.listdir(wheels_dir) if f.endswith(".whl")]
        if not wheels:
            print("Build failed: No wheel file found.")
            return
            
        latest_wheel = max([os.path.join(wheels_dir, f) for f in wheels], key=os.path.getctime)
        subprocess.check_call([sys.executable, "-m", "pip", "install", "--force-reinstall", latest_wheel], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        
        with open("rust_core/.build_hash", 'w') as f:
            f.write(current_hash)
            
        print("Rust Module Built & Installed.\n")
        
    except Exception as e:
        print(f"Rust Optimization Load Failed: {e}\nRunning in Python-only mode.")

def terminate_process(proc):
    if proc is None:
        return
    
    try:
        proc.terminate()
        try:
            proc.wait(timeout=3)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait()
    except Exception as e:
        pass

def start_cloudflare_tunnel():
    global cloudflared_process
    
    if not config.cloudflare.get('enabled', False):
        return

    token = os.getenv('CLOUDFLARE_TUNNEL_TOKEN')
    if not token:
        print("Warning: Cloudflare Tunnel enabled but CLOUDFLARE_TUNNEL_TOKEN not found in .env")
        return

    print("Starting Cloudflare Tunnel...")
    try:
        # Check if cloudflared exists
        try:
            subprocess.run(["cloudflared", "--version"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
        except (FileNotFoundError, subprocess.CalledProcessError):
            print("Error: 'cloudflared' not found in PATH. Tunnel will not start.")
            return

        cloudflared_process = subprocess.Popen(
            ["cloudflared", "tunnel", "run", "--token", token],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )
        print("Cloudflare Tunnel started.")
    except Exception as e:
        print(f"Failed to start Cloudflare Tunnel: {e}")

def run_server_loop():
    global server_process
    restart_delay = config.startup.get('server_restart_delay', 1.0)
    error_delay = config.startup.get('error_restart_delay', 5.0)
    
    while not shutdown_flag.is_set():
        try:
            print("Starting Production Web Server subprocess...")
            server_process = subprocess.Popen([sys.executable, "source/server.py"], cwd=ROOT_DIR)
            server_process.wait()
            if shutdown_flag.is_set():
                break
            print(f"Server exit. Restarting in {restart_delay}s...")
            time.sleep(restart_delay)
        except Exception as e:
            if shutdown_flag.is_set():
                break
            print(f"Server Loop Error: {e}")
            time.sleep(error_delay)

def shutdown_all():
    global server_process, scanner_process, cloudflared_process
    print("Shutting down...")
    shutdown_flag.set()
    
    if scanner_process:
        terminate_process(scanner_process)
    if server_process:
        terminate_process(server_process)
    if cloudflared_process:
        print("Stopping Cloudflare Tunnel...")
        terminate_process(cloudflared_process)
    
    print("All processes terminated.")
    sys.exit(0)

if __name__ == "__main__":
    from dotenv import load_dotenv
    load_dotenv()

    # Handle Ctrl+C
    signal.signal(signal.SIGINT, lambda s, f: shutdown_all())
    signal.signal(signal.SIGTERM, lambda s, f: shutdown_all())
    
    # Start server thread if enabled
    if config.server.get('enabled', True):
        server_thread = threading.Thread(target=run_server_loop, daemon=True)
        server_thread.start()
    
    # Start Cloudflare Tunnel if enabled
    start_cloudflare_tunnel()

    # Build rust module
    build_rust()

    # Verify directory permissions
    print("Verifying storage permissions...")
    try:
        saves_dir = config.storage.get('local_saves', 'Saves')
        test_file = os.path.join(ROOT_DIR, saves_dir, ".perm_test")
        os.makedirs(os.path.dirname(test_file), exist_ok=True)
        with open(test_file, 'w') as f: f.write("ok")
        os.remove(test_file)
        print("Storage permissions OK.")

    except Exception as e:
        print(f"CRITICAL: Storage permission check failed: {e}")
        print("Check your Docker volume mounts and host folder permissions.")
        sys.exit(1)
    
    scanner_delay = config.startup.get('scanner_restart_delay', 1.0)
    error_delay = config.startup.get('error_restart_delay', 5.0)
    
    while not shutdown_flag.is_set():
        try:
            print("Launching scanner (source/main.py)...")
            scanner_process = subprocess.Popen([sys.executable, "source/main.py"], cwd=ROOT_DIR)
            scanner_process.wait()
            if shutdown_flag.is_set():
                break
            print(f"Scanner exit. Restarting in {scanner_delay}s...")
            time.sleep(scanner_delay)
        except Exception as e:
            if shutdown_flag.is_set():
                break
            print(f"Scanner Error: {e}")
            time.sleep(error_delay)