#!/usr/bin/env bash

# Exit immediately if any command fails
set -e

# Make sure the script is running with superuser privileges
if [ "$EUID" -ne 0 ]; then
  echo "Error: This deployment script must be run as root (sudo)." >&2
  exit 1
fi

# ====================================================
# ⚙️ CONFIGURATION: Set your Client ID & Central URL
# ====================================================
CLIENT_ID=""

# For Production deployment, use: "http://rmm-backend.seekright.com"
# For Local WSL testing on your laptop, use: "http://172.18.16.1:8000"
CENTRAL_URL="http://172.18.16.1:8000"

# Check if Client ID is set; if not, check command line arguments, otherwise default to hostname
if [ -z "$CLIENT_ID" ]; then
  CLIENT_ID="${1:-$(hostname)}"
fi

echo "🚀 Preparing SeekRight RMM Agent installation..."
echo "📍 Target Client ID: $CLIENT_ID"
echo "🌐 Central Server URL: $CENTRAL_URL"

# 1. Install System Dependencies
echo "🔄 Updating system package indexes and installing python dependencies..."
apt-get update -y
apt-get install -y python3 python3-pip python3-psutil

# 2. Create Destination Directories
echo "📁 Creating installation workspace..."
mkdir -p /opt/seekright-agent
chmod 755 /opt/seekright-agent

# 3. Dynamically Generate Python Agent File
echo "📝 Generating client agent file..."
cat << 'EOF' > /opt/seekright-agent/client_agent_prototype.py
import urllib.request
import json
import time
import subprocess
import psutil
import os
import sys
import string

CLIENT_ID = "TEMPLATE_CLIENT_ID"
CENTRAL_BASE = "TEMPLATE_CENTRAL_URL"

CENTRAL_API_URL = f"{CENTRAL_BASE}/api/get-command?client_id={CLIENT_ID}"
CENTRAL_TELEMETRY_URL = f"{CENTRAL_BASE}/api/telemetry?client_id={CLIENT_ID}"

platform_name = "nt" if os.name == "nt" else "posix"
CENTRAL_COMMANDS_URL = f"{CENTRAL_BASE}/api/get-commands?platform={platform_name}"

CACHED_COMMANDS = {}

def load_allowed_commands():
    global CACHED_COMMANDS
    try:
        req = urllib.request.Request(CENTRAL_COMMANDS_URL)
        req.add_header('ngrok-skip-browser-warning', 'true')
        with urllib.request.urlopen(req, timeout=3) as response:
            commands = json.loads(response.read().decode())
            if commands:
                CACHED_COMMANDS = commands
            return CACHED_COMMANDS
    except Exception as e:
        print(f"    -> [!] Failed to load commands from central API: {e}")
        if CACHED_COMMANDS:
            print("    -> Returning last successfully cached whitelist.")
            return CACHED_COMMANDS
        return {"whoami": ["whoami"]}

ROCKET_CHAT_WEBHOOK = "https://text.seekright.com/hooks/6a0eb41ea88367bc6e101f0c/gvXfD8zSNRti43kEabqdE9tMCvNi9zAAoGfbao7cxcKbiW6R"

def send_rocket_chat_notification(title, message, color="#764FA5"):
    try:
        if not message:
            message = "No output provided."
        payload = {
            "alias": f"Agent: {CLIENT_ID}",
            "text": "Execution Report",
            "attachments": [{
                "title": title,
                "text": message,
                "color": color
            }]
        }
        data = json.dumps(payload).encode('utf-8')
        req = urllib.request.Request(ROCKET_CHAT_WEBHOOK, data=data, headers={'Content-Type': 'application/json'}, method='POST')
        urllib.request.urlopen(req, timeout=3)
    except Exception as e:
        print(f"    -> [!] Failed to send webhook: {e}")

def collect_and_send_telemetry():
    try:
        gpus = []
        try:
            result = subprocess.run(
                ["nvidia-smi", "--query-gpu=name,utilization.gpu,temperature.gpu,memory.used,memory.total,power.draw,power.limit,fan.speed", "--format=csv,noheader"], 
                capture_output=True, text=True
            )
            if result.returncode == 0:
                for line in result.stdout.strip().split('\n'):
                    if line:
                        parts = [p.strip() for p in line.split(',')]
                        if len(parts) >= 8:
                            gpus.append({
                                "name": parts[0],
                                "utilization": parts[1],
                                "temp": parts[2] + "C",
                                "memory_used": parts[3],
                                "memory_total": parts[4],
                                "power_draw": parts[5],
                                "power_limit": parts[6],
                                "fan_speed": parts[7]
                            })
        except FileNotFoundError:
            pass

        disks = []
        for part in psutil.disk_partitions(all=False):
            if 'loop' in part.device or 'loop' in part.fstype or part.fstype == '':
                continue
            if part.mountpoint.startswith('/snap') or part.fstype == 'squashfs':
                continue
            if part.mountpoint.startswith('/boot') or part.mountpoint.startswith('/mnt/wsl'):
                continue
            try:
                usage = psutil.disk_usage(part.mountpoint)
                disks.append({
                    "mount": part.mountpoint,
                    "total_gb": round(usage.total / (1024**3), 2),
                    "free_gb": round(usage.free / (1024**3), 2),
                    "percent": usage.percent
                })
            except PermissionError:
                continue

        payload = {
            "cpu_percent": psutil.cpu_percent(interval=0.5),
            "memory_percent": psutil.virtual_memory().percent,
            "memory_total_gb": round(psutil.virtual_memory().total / (1024**3), 2),
            "memory_free_gb": round(psutil.virtual_memory().available / (1024**3), 2),
            "disks": disks,
            "gpus": gpus
        }
        
        data = json.dumps(payload).encode('utf-8')
        req = urllib.request.Request(CENTRAL_TELEMETRY_URL, data=data, headers={'Content-Type': 'application/json'}, method='POST')
        req.add_header('ngrok-skip-browser-warning', 'true')
        urllib.request.urlopen(req, timeout=3)
        print("[*] Telemetry pushed successfully.")
    except Exception as e:
        print(f"[!] Failed to push telemetry: {e}")

def send_command_result_to_server(command: str, returncode: int, stdout: str, stderr: str):
    try:
        base_url = CENTRAL_API_URL.split("/api/")[0]
        url = f"{base_url}/api/command-result?client_id={CLIENT_ID}"
        payload = {
            "command": command,
            "returncode": returncode,
            "stdout": stdout,
            "stderr": stderr
        }
        data = json.dumps(payload).encode('utf-8')
        req = urllib.request.Request(
            url, 
            data=data, 
            headers={'Content-Type': 'application/json'}, 
            method='POST'
        )
        req.add_header('ngrok-skip-browser-warning', 'true')
        with urllib.request.urlopen(req, timeout=3) as response:
            pass
        print("    -> Sent command execution output to Central Server.")
    except Exception as e:
        print(f"    -> [!] Failed to send command result to server: {e}")

def poll_server():
    try:
        print(f"[*] Checking for commands...")
        req = urllib.request.Request(CENTRAL_API_URL)
        req.add_header('ngrok-skip-browser-warning', 'true')
        with urllib.request.urlopen(req, timeout=3) as response:
            data = json.loads(response.read().decode())
            command_key = data.get("command", "none")
            
            if command_key == "none":
                print("    -> No pending commands.")
            else:
                allowed_commands = load_allowed_commands()
                if command_key in allowed_commands:
                    print(f"    -> Executing approved command: '{command_key}'")
                    safe_cmd_list = allowed_commands[command_key]
                    
                    is_shutdown_or_restart = any(x in command_key.lower() for x in ["reboot", "restart", "kill_python"])
                    if is_shutdown_or_restart:
                        print("    -> Shutdown/restart command detected. Sending pre-execution success notification...")
                        send_rocket_chat_notification(
                            f"🔄 Executing: {command_key}", 
                            f"System/agent is performing a scheduled action: {' '.join(safe_cmd_list)}\nConnection may drop temporarily.", 
                            "#FFA500"
                        )
                        time.sleep(3)
                        
                    try:
                        result = subprocess.run(safe_cmd_list, capture_output=True, text=True)
                        output = result.stdout.strip()
                        error_output = result.stderr.strip() if result.stderr else ""
                        
                        # Post execution results to Central Server
                        send_command_result_to_server(command_key, result.returncode, output, error_output)
                        
                        if not is_shutdown_or_restart:
                            if result.returncode == 0:
                                print(f"    -> Output: {output}")
                                send_rocket_chat_notification(f"✅ Success: {command_key}", output, "#00FF00")
                            else:
                                print(f"    -> Error: {error_output}")
                                send_rocket_chat_notification(f"❌ Failed: {command_key}", error_output, "#FF0000")
                    except Exception as ex:
                        print(f"    -> [!] Execution Error: {ex}")
                        send_command_result_to_server(command_key, -1, "", str(ex))
                        if not is_shutdown_or_restart:
                            send_rocket_chat_notification(f"❌ Execution Error: {command_key}", str(ex), "#FF0000")
                else:
                    warning_msg = f"Server requested unknown/malicious command: '{command_key}'. Ignored."
                    print(f"    -> [!] SECURITY WARNING: {warning_msg}")
                    send_rocket_chat_notification("⚠️ SECURITY WARNING", warning_msg, "#FFA500")
    except Exception as e:
        print(f"    -> [!] Failed to connect to central server: {e}")

def get_dir_size(path):
    total_size = 0
    try:
        for dirpath, dirnames, filenames in os.walk(path):
            for f in filenames:
                fp = os.path.join(dirpath, f)
                try:
                    total_size += os.path.getsize(fp)
                except Exception:
                    pass
    except Exception:
        pass
    return total_size

def format_size(size_in_bytes):
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if size_in_bytes < 1024.0:
            return f"{size_in_bytes:.2f} {unit}"
        size_in_bytes /= 1024.0
    return f"{size_in_bytes:.2f} PB"

def check_takeleap_folders():
    target_names = {"SHIFT", "Error_Videos", "UPLOAD_FOLDER"}
    parent_found = False
    files_found = False
    
    if os.name == "nt":
        drives = [f"{letter}:\\" for letter in string.ascii_uppercase if os.path.exists(f"{letter}:\\")]
    else:
        drives = ["/"]
        
    for drive in drives:
        for root, dirs, files in os.walk(drive, topdown=True):
            if os.name == "nt":
                dirs[:] = [d for d in dirs if d.lower() not in {"windows", "program files", "program files (x86)", "appdata", "programdata", "$recycle.bin", "system volume information", "node_modules", ".git"}]
            else:
                if root == "/":
                    dirs[:] = [d for d in dirs if d.lower() not in {"proc", "sys", "dev", "var", "lib", "run", "boot", "snap", "node_modules", ".git"}]
                elif root == "/mnt":
                    dirs[:] = [d for d in dirs if not (len(d) == 1 and d.isalpha()) and d.lower() not in {"wslg", "wsl"}]
                
            if os.path.basename(root).upper() == "TAKELEAP":
                parent_found = True
                for sub in dirs:
                    if sub in target_names:
                        subpath = os.path.join(root, sub)
                        try:
                            for item in os.listdir(subpath):
                                if os.path.isfile(os.path.join(subpath, item)):
                                    files_found = True
                                    break
                        except Exception:
                            pass
                    if files_found:
                        break
            if files_found:
                break
        if files_found:
            break
            
    if not parent_found:
        print("Folder not found")
    else:
        print("true" if files_found else "false")

def list_takeleap_parent():
    found_any = False
    
    if os.name == "nt":
        drives = [f"{letter}:\\" for letter in string.ascii_uppercase if os.path.exists(f"{letter}:\\")]
    else:
        drives = ["/"]
        
    for drive in drives:
        for root, dirs, files in os.walk(drive, topdown=True):
            if os.name == "nt":
                dirs[:] = [d for d in dirs if d.lower() not in {"windows", "program files", "program files (x86)", "appdata", "programdata", "$recycle.bin", "system volume information", "node_modules", ".git"}]
            else:
                if root == "/":
                    dirs[:] = [d for d in dirs if d.lower() not in {"proc", "sys", "dev", "var", "lib", "run", "boot", "snap", "node_modules", ".git"}]
                elif root == "/mnt":
                    dirs[:] = [d for d in dirs if not (len(d) == 1 and d.isalpha()) and d.lower() not in {"wslg", "wsl"}]
                
            if os.path.basename(root).upper() == "TAKELEAP":
                found_any = True
                total_size = get_dir_size(root)
                print(f"\n[+] Found TAKELEAP at: {root} (Total Size: {format_size(total_size)})")
                if dirs:
                    print("    Subfolders present:")
                    for d in dirs:
                        subpath = os.path.join(root, d)
                        sub_size = get_dir_size(subpath)
                        print(f"      - {d} (Size: {format_size(sub_size)})")
                else:
                    print("    (No subfolders found inside)")
                    
    if not found_any:
        print("No TAKELEAP folders found on this system.")

def list_takeleap_subfolder(subfolder_name):
    found_any = False
    
    if os.name == "nt":
        drives = [f"{letter}:\\" for letter in string.ascii_uppercase if os.path.exists(f"{letter}:\\")]
    else:
        drives = ["/"]
        
    for drive in drives:
        for root, dirs, files in os.walk(drive, topdown=True):
            if os.name == "nt":
                dirs[:] = [d for d in dirs if d.lower() not in {"windows", "program files", "program files (x86)", "appdata", "programdata", "$recycle.bin", "system volume information", "node_modules", ".git"}]
            else:
                if root == "/":
                    dirs[:] = [d for d in dirs if d.lower() not in {"proc", "sys", "dev", "var", "lib", "run", "boot", "snap", "node_modules", ".git"}]
                elif root == "/mnt":
                    dirs[:] = [d for d in dirs if not (len(d) == 1 and d.isalpha()) and d.lower() not in {"wslg", "wsl"}]
                
            if os.path.basename(root).upper() == "TAKELEAP":
                for sub in dirs:
                    if sub.lower() == subfolder_name.lower():
                        subpath = os.path.join(root, sub)
                        try:
                            found_any = True
                            sub_size = get_dir_size(subpath)
                            items = [item for item in os.listdir(subpath) if os.path.isfile(os.path.join(subpath, item))]
                            print(f"\n[+] Folder: {subpath} (Total Size: {format_size(sub_size)})")
                            if items:
                                for idx, item in enumerate(items, 1):
                                    print(f"  {idx}. {item}")
                            else:
                                print("  (Empty)")
                        except Exception as e:
                            print(f"  [!] Error accessing folder {subpath}: {e}")
                            
    if not found_any:
        print(f"No TAKELEAP/{subfolder_name} folders found on this system.")

if __name__ == "__main__":
    if len(sys.argv) > 1:
        if sys.argv[1] == "--check-takeleap":
            check_takeleap_folders()
            sys.exit(0)
        elif sys.argv[1] == "--list-takeleap":
            list_takeleap_parent()
            sys.exit(0)
        elif sys.argv[1] == "--list-takeleap-sub" and len(sys.argv) > 2:
            list_takeleap_subfolder(sys.argv[2])
            sys.exit(0)
        else:
            print(f"Unknown argument or missing parameters: {sys.argv[1]}")
            sys.exit(1)
    else:
        print(f"Starting Agent for {CLIENT_ID}...")
        try:
            while True:
                collect_and_send_telemetry()
                poll_server()
                time.sleep(10)
        except KeyboardInterrupt:
            print("\n[*] Agent stopped.")
EOF

# Inject the Client ID and Central URL dynamically
sed -i "s/TEMPLATE_CLIENT_ID/$CLIENT_ID/g" /opt/seekright-agent/client_agent_prototype.py
sed -i "s|TEMPLATE_CENTRAL_URL|$CENTRAL_URL|g" /opt/seekright-agent/client_agent_prototype.py
chmod 644 /opt/seekright-agent/client_agent_prototype.py

# 4. Create Background systemd Service
echo "⚙️ Configuring background systemd daemon service..."
cat << 'EOF' > /etc/systemd/system/seekright-agent.service
[Unit]
Description=SeekRight RMM Background Agent
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
# -u enables unbuffered stdout output for direct real-time logging via journalctl
ExecStart=/usr/bin/python3 -u /opt/seekright-agent/client_agent_prototype.py
WorkingDirectory=/opt/seekright-agent
Restart=always
RestartSec=5
User=root
Group=root
StandardOutput=syslog
StandardError=syslog
SyslogIdentifier=seekright-agent

[Install]
WantedBy=multi-user.target
EOF

chmod 644 /etc/systemd/system/seekright-agent.service

# 5. Enable and Launch Service
echo "🔄 Reloading system daemons and starting agent service..."
systemctl daemon-reload
systemctl enable seekright-agent.service
systemctl restart seekright-agent.service

echo ""
echo "✨ ============================================= ✨"
echo "✅ SeekRight RMM Agent successfully deployed!"
echo "📍 Location: /opt/seekright-agent/client_agent_prototype.py"
echo "📋 System Service: seekright-agent.service"
echo "✨ ============================================= ✨"
echo ""
echo "📈 Displaying real-time execution logs (Press Ctrl+C to exit log view):"
echo "--------------------------------------------------------"
journalctl -u seekright-agent.service -n 15 -f