#!/usr/bin/env python3
import time
import subprocess
import os
import random
from datetime import datetime
import threading
import sys
# Global counter
vnf_switch_count = 0
RUN_DURATION = 2 * 60    # 2 minute
attack_success_count = 0  # ← counts buffer overflow seg fault attacks


begin_time=time.time()

# === GLOBAL CONFIGURATION ===
SSH_KEY = os.path.expanduser("~/.ssh/fresh-vnf-key")
SSH_OPTS = [
    "-o", "BatchMode=yes",
    "-o", "StrictHostKeyChecking=accept-new",
    "-o", "ConnectTimeout=5",
    "-n", "-f",
]
INTERVAL = 30
CRASH_CHECK_INTERVAL = 0.5
LOG_DIR = "/home/ubuntu/scripts"

AGGREGATORS = [
    {
        "name": "cpp-vnf",
        "ip": "10.20.20.108",
        "cmd": "stdbuf -o0 -e0 /home/ubuntu/lidar_pointcloud_concatenator",
        "search_pattern": "lidar_pointcloud_concatenator",
    },
    {
        "name": "python-vnf",
        "ip": "10.20.20.119",
        "cmd": "python3 -u /home/ubuntu/lidar_diversification.py",
        "search_pattern": "lidar_diversification.py",
    },
    {
        "name": "java-vnf",
        "ip": "10.20.20.8",
        "cmd": (
            "bash -lc 'cd /home/ubuntu/pointcloud_java && "
            "stdbuf -o0 -e0 java -cp .:org.eclipse.paho.client.mqttv3-1.2.5.jar "
            "lidar_PointCloudConcatenator'"
        ),
        "search_pattern": "lidar_PointCloudConcatenator",
    },
]

def ssh_run(ip, cmd, capture=False):
    ssh_cmd = ["ssh", *SSH_OPTS, "-i", SSH_KEY, f"ubuntu@{ip}", cmd]
    return subprocess.run(ssh_cmd, capture_output=capture, text=True)

def check_process_alive(ip, pattern):
    result = ssh_run(ip, f"pgrep -f '{pattern}'", capture=True)
    return result.returncode == 0 and result.stdout.strip()

def check_for_crash_indicators(ip, log_path):
    segfault_log = f"tail -n 10 {log_path} 2>/dev/null | grep -i 'segmentation fault' || true"
    dmesg_check = "dmesg | tail -n 50 | grep -i 'segfault' | tail -n 1 || true"
    if ssh_run(ip, segfault_log, capture=True).stdout.strip():
        return True
    if ssh_run(ip, dmesg_check, capture=True).stdout.strip():
        return True
    return False

def start_vnf(vnf, kill_old=True):
    global vnf_switch_count
    global attack_success_count

    ip = vnf["ip"]
    cmd = vnf["cmd"]
    log_path = f"{LOG_DIR}/{vnf['name']}.log"
    pat = vnf.get("search_pattern", cmd)

    if kill_old:
        ssh_run(ip, f"pkill -f '{pat}' || true")

    print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}] Launching {vnf['name']} at {ip}")
    remote_cmd = f"mkdir -p {LOG_DIR} && nohup {cmd} > {log_path} 2>&1 < /dev/null & sleep 0.5; sync"
    ''' if vnf["name"] == "cpp-vnf":
       ssh_run(ip, remote_cmd)
       result = ssh_run(ip, f"tail -n 5 {log_path}", capture=True)
      # ssh_output = (result.stdout + result.stderr).strip()
        if "segmentation fault" in result.stdout.lower():
          attack_success_count += 1
          print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}] ATTACK SUCCESSFUL!")
          print(f"Attack Success Count: {attack_success_count}")
    else:'''
    ssh_run(ip, remote_cmd)  # ← original, no capture for java/python VNF


    #Capture ssh output to catch bash-level segfault message
    #print(f"[DEBUG] ssh output: '{ssh_output}'")  # ← remove after testi
    #ssh_run(ip, remote_cmd)
    print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}] Started.")

    time.sleep(0.5)
    pid = ssh_run(ip, f"pgrep -f '{pat}' | head -n1", capture=True).stdout.strip() or "n/a"
    exe_path = "(n/a)"
    if pid != "n/a":
        r = ssh_run(ip, f"readlink -f /proc/{pid}/exe 2>/dev/null || true", capture=True)
        exe_path = (r.stdout or "").strip() or "(n/a)"
        vnf_switch_count += 1
        print(f"VNF Switch Count: {vnf_switch_count}")

    time.sleep(1)
    ssh_run(ip, "sync")
    tail = ssh_run(ip, f"tail -n 10 {log_path} || true", capture=True).stdout.strip() or "(log empty)"

    print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}] pid: {pid}")
    print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}] exe: {exe_path}")
    print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}] last log lines:\n{tail}")

    alive = pid != "n/a" and ssh_run(ip, f"ps -p {pid}", capture=True).returncode == 0
    #crashed = not alive or check_for_crash_indicators(ip, log_path)
    #segfault_in_ssh = "segmentation fault" in tail.lower()
    crashed = not alive or  check_for_crash_indicators(ip, log_path)
    '''
    if segfault_in_ssh:
        attack_success_count += 1
        print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}] ATTACK SUCCESSFUL!")
        print(f"Attack Success Count: {attack_success_count}")    
'''

    if crashed:
        print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}]   {vnf['name']} crashed or failed to start properly!")
        return True, pid

    return False, pid

def monitor_vnf_health(vnf, stop_event):
    ip = vnf["ip"]
    log_path = f"{LOG_DIR}/{vnf['name']}.log"
    pattern = vnf.get("search_pattern", vnf["cmd"])

    while not stop_event.is_set():
        if not check_process_alive(ip, pattern):
            print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}]  {vnf['name']} process died!")
            return True
        if check_for_crash_indicators(ip, log_path):
            print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}]  {vnf['name']} segmentation fault detected!")
            return True
        stop_event.wait(CRASH_CHECK_INTERVAL)

    return False

print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}] MQTT Controller started — rotating every {INTERVAL}s")
print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}] Crash detection enabled — checking every {CRASH_CHECK_INTERVAL}s")

last_vnf_name = None

while True:
    attempts = 0
    while True:
        random.shuffle(AGGREGATORS)
        if AGGREGATORS[0]["name"] != last_vnf_name or attempts >= 10:
            break
        attempts += 1

    for vnf in AGGREGATORS:
        elapsed = time.time() - begin_time
        if elapsed >= RUN_DURATION:
          print(f"\n[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Time limit reached.")
          print(f"Total VNF Switches: {vnf_switch_count}")
          print(f"Attack Success Count: {attack_success_count}")

          sys.exit(0)

        crashed_on_start, pid = start_vnf(vnf, kill_old=True)
        if crashed_on_start:
            print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}] Switching due to startup failure...")
            continue

        stop_monitoring = threading.Event()
        monitor_thread = threading.Thread(
            #target=lambda: monitor_vnf_health(vnf, stop_monitoring),
             # Correct — captures current vnf
            target=lambda v=vnf: monitor_vnf_health(v, stop_monitoring),
            daemon=True
        )
        monitor_thread.start()

        start_time = time.time()
        
         


        crashed_during_run = False
        while time.time() - start_time < INTERVAL:
            if not monitor_thread.is_alive():
                crashed_during_run = True
                break
            time.sleep(0.1)

        stop_monitoring.set()
        last_vnf_name = vnf["name"]

        if crashed_during_run:
            print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}]  Switching due to crash detection!")
            attack_success_count += 1
            print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}] ATTACK SUCCESSFUL!")
            print(f"Attack Success Count: {attack_success_count}")
        else:
            print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}] Sleeping for {INTERVAL}s...\n")

