#!/usr/bin/env python3
"""Wallarm Docker Node Manager"""

import json, os, sys, time, subprocess, shutil

VERSION = "2.0.0"
BASE = "/opt/fw"
APP = f"{BASE}/app"
CONF = f"{APP}/fw.conf"
STATE = f"{APP}/state.json"
DOCKER_DIR = f"{BASE}/docker"
OFFLINE_DIR = f"{BASE}/offline"
WALLARM_IMAGE = os.environ.get("WALLARM_IMAGE", "wallarm/node:6.11.0-rc1")

# Ensure Docker is in PATH
os.environ["PATH"] = f"{DOCKER_DIR}/bin:{os.environ.get('PATH', '')}"

def run(cmd, timeout=120):
    return subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout)

def load_config():
    if not os.path.exists(CONF): return {}
    with open(CONF) as f: return json.load(f)

def load_state():
    if not os.path.exists(STATE): return {"nodes": []}
    with open(STATE) as f: return json.load(f)

def save_state(s):
    os.makedirs(APP, exist_ok=True)
    with open(STATE, "w") as f: json.dump(s, f, indent=2)

# ─── Docker ────────────────────────────────────────────────────────────

def install_docker():
    if shutil.which("docker"):
        return
    print("Installing Docker...")
    run("apt-get update -qq && apt-get install -y -qq docker.io 2>/dev/null || yum install -y -q docker 2>/dev/null", timeout=120)
    run("systemctl enable docker && systemctl start docker")

def deploy_container(name, token, cloud, port, upstream_ip, upstream_port, labels, mode):
    api_host = "us1.api.wallarm.com" if cloud == "US" else "api.wallarm.com"

    # Ensure Docker is running
    install_docker()

    # Pull image if needed
    if run(f"docker images -q {WALLARM_IMAGE}").stdout.strip() == "":
        print(f"[{name}] Pulling Wallarm image...")
        run(f"docker pull {WALLARM_IMAGE}", timeout=300)

    # Remove old container
    container = f"wallarm-{name}"
    run(f"docker rm -f {container} 2>/dev/null", timeout=10)

    host_dir = f"{BASE}/{name}"
    os.makedirs(host_dir, exist_ok=True)

    # Generate start.sh (declarative config + startup in one file)
    start_path = f"{host_dir}/start.sh"
    with open(start_path, "w") as f:
        f.write(f"""#!/bin/bash
# Wallarm node: {name}
docker rm -f {container} 2>/dev/null || true
docker run -d \\
  --name {container} \\
  --restart always \\
  -p {port}:80 \\
  -e WALLARM_API_TOKEN='{token}' \\
  -e WALLARM_API_HOST={api_host} \\
  -e WALLARM_MODE={mode} \\
  -v {host_dir}/nginx.conf:/etc/nginx/http.d/default.conf:ro \\
  {WALLARM_IMAGE}
sleep 3
docker ps --filter name={container} --format '{{{{.Status}}}}'
""")
    os.chmod(start_path, 0o755)

    # Write nginx.conf with production defaults
    nginx_conf = f"{host_dir}/nginx.conf"
    if not os.path.exists(nginx_conf):
        with open(nginx_conf, "w") as f:
            f.write(f"""server {{
    listen 80;
    server_name _;
    client_max_body_size 1024m;
    proxy_read_timeout 300s;
    proxy_send_timeout 300s;

    location /wallarm-status {{
        wallarm_status on;
        allow 127.0.0.0/8;
        deny all;
    }}

    location / {{
        proxy_pass http://{upstream_ip}:{upstream_port};
        proxy_set_header Host $host;
        proxy_set_header X-Real-IP $remote_addr;
        proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
        proxy_set_header X-Forwarded-Proto $scheme;
    }}
}}
""")

    print(f"[{name}] Starting...")
    r = run(f"bash {start_path}", timeout=30)
    if r.returncode != 0:
        print(f"❌ {name}: {r.stderr}")
        return False

    # Save state
    s = load_state()
    s.setdefault("nodes", []).append({
        "name": name, "type": "docker", "port": int(port),
        "upstream_ip": upstream_ip, "upstream_port": int(upstream_port),
        "status": "running", "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ")
    })
    save_state(s)
    print(f"✅ {name} running on port {port} (docker ps --filter name={container})")
    return True

def remove_container(name):
    container = f"wallarm-{name}"
    run(f"docker rm -f {container} 2>/dev/null", timeout=10)
    shutil.rmtree(f"/opt/wallarm-{name}", ignore_errors=True)
    s = load_state()
    s["nodes"] = [n for n in s.get("nodes", []) if n["name"] != name]
    save_state(s)
    print(f"✅ {name} removed")

def container_status():
    r = run("docker ps --filter 'name=wallarm-' --format '{{.Names}}\t{{.Status}}\t{{.Ports}}'")
    if not r.stdout.strip():
        print("No containers running.")
        return
    print("\n─── Containers ───")
    for line in r.stdout.strip().split("\n"):
        print(f"  {line}")

# ─── Deploy All ───────────────────────────────────────────────────────

def deploy_all():
    cfg = load_config()
    nodes = cfg.get("nodes", cfg)  # support both {"nodes":{...}} and flat format
    if not nodes:
        print("No nodes in fw.conf")
        return
    for name, nc in nodes.items():
        if not isinstance(nc, dict): continue
        deploy_container(name, nc["token"], nc.get("cloud","EU"),
                        str(nc.get("port","8081")), nc.get("upstream_ip","127.0.0.1"),
                        str(nc.get("upstream_port","80")), nc.get("labels",""),
                        nc.get("mode","monitoring"))

# ─── Menu ─────────────────────────────────────────────────────────────

def main():
    if len(sys.argv) > 1 and sys.argv[1] == "--deploy-all":
        return deploy_all()

    print(f"═══ Wallarm Node Manager v{VERSION} ═══\n")
    print("Preflight checks...")
    if not shutil.which("docker"):
        print("  ❌ Docker not found. Run setup.sh first.")
        sys.exit(1)
    print("  ✅ Docker ready\n")

    cfg = load_config()
    deployed = {n["name"] for n in load_state().get("nodes", [])}

    while True:
        undeployed = [k for k in cfg if k not in deployed]
        
        print("\n─── Menu ───")
        print("  [1] Add a new node")
        
        n = 2
        node_map = {}
        ops = {}
        for name in undeployed:
            print(f"  [{n}] Deploy {name} (port {cfg[name].get('port','?')})")
            node_map[str(n)] = name
            n += 1
        
        if deployed:
            print(f"  [{n}] Edit nodes"); ops['edit'] = str(n); n += 1
            print(f"  [{n}] Delete a node"); ops['del'] = str(n); n += 1
        print(f"  [{n}] Status"); ops['status'] = str(n); n += 1
        print(f"  [{n}] Remote tunnel"); ops['tunnel'] = str(n); n += 1
        print("  [q] Quit")
        
        c = input("\nChoice: ").strip()
        
        if c == "1":
            add_node()
            cfg = load_config()
            deployed = {n["name"] for n in load_state().get("nodes", [])}
        elif c in node_map:
            name = node_map[c]
            nc = cfg[name]
            deploy_container(name, nc["token"], nc.get("cloud","EU"),
                            str(nc.get("port","8081")), nc.get("upstream_ip",""),
                            str(nc.get("upstream_port","80")), "",
                            nc.get("mode","monitoring"))
            deployed.add(name)
        elif c == ops.get('edit') and deployed:
            edit_nodes()
        elif c == ops.get('del') and deployed:
            remove_menu()
        elif c == ops.get('status'):
            container_status()
        elif c == ops.get('tunnel'):
            start_tunnel()
        elif c.lower() == "q":
            break


def add_node():
    print("\n─── Add New Node ───")
    name = input("Node name: ").strip()
    if not name: return
    token = input("Wallarm token: ").strip()
    if not token: print("Token required."); return
    print("Cloud region:")
    print("  [1] EU (api.wallarm.com)")
    print("  [2] US (us1.api.wallarm.com)")
    c = input("Choose [1]: ").strip()
    cloud = "EU" if c == "2" else "EU"
    if c == "2": cloud = "US"
    port = input("Port [8081]: ").strip() or "8081"
    upstream_ip = input("Upstream IP [127.0.0.1]: ").strip() or "127.0.0.1"
    upstream_port = input("Upstream port [80]: ").strip() or "80"
    print("\nTraffic mode:")
    print("  [1] monitoring     — detect only")
    print("  [2] safe_blocking  — block definitely malicious")
    print("  [3] block          — block all attacks")
    print("  [4] off            — disable")
    m = input("Choose [1]: ").strip()
    mode = {"1":"monitoring","2":"safe_blocking","3":"block","4":"off"}.get(m, "monitoring")
    
    # Save to fw.conf
    cfg = load_config()
    cfg[name] = {"token": token, "cloud": cloud, "port": port,
                  "upstream_ip": upstream_ip, "upstream_port": upstream_port,
                  "labels": f"group={name}", "mode": mode}
    os.makedirs(APP, exist_ok=True)
    with open(CONF, "w") as f: json.dump(cfg, f, indent=2)
    
    # Deploy immediately
    print(f"\nDeploying {name}...")
    deploy_container(name, token, cloud, port, upstream_ip, upstream_port, f"group={name}", mode)

def remove_menu():
    s = load_state().get("nodes", [])
    if not s: print("No nodes deployed."); return
    print("\n─── Remove Node ───")
    for i, n in enumerate(s):
        print(f"  [{i+1}] {n['name']} (port {n.get('port','')})")
    c = input("\nSelect number: ").strip()
    if not c.isdigit() or int(c) < 1 or int(c) > len(s): return
    remove_container(s[int(c)-1]["name"])

def edit_menu():
    cfg = load_config()
    if not cfg: print("No nodes in fw.conf"); return
    print("\n─── Edit Node ───")
    names = list(cfg.keys())
    for i, name in enumerate(names):
        nc = cfg[name]
        print(f"  [{i+1}] {name} (port {nc.get('port','')})")
    c = input("\nSelect number: ").strip()
    if not c.isdigit() or int(c) < 1 or int(c) > len(names): return
    name = names[int(c)-1]
    nc = cfg[name]
    print(f"\nEditing {name}. Leave blank to keep current value:")
    for key, prompt in [("port","Port"), ("upstream_ip","Upstream IP"), 
                         ("upstream_port","Upstream port"), ("mode","Mode")]:
        v = input(f"{prompt} [{nc.get(key,'')}]: ").strip()
        if v: nc[key] = v
    with open(CONF, "w") as f: json.dump(cfg, f, indent=2)
    print(f"✅ {name} updated in fw.conf. Re-deploy to apply changes.")

    main()
def start_tunnel():
    host = input("Jumphost URL [ssh.sechpoint.app:443]: ").strip() or "ssh.sechpoint.app:443"
    user = input("Username [wallarm-tunnel]: ").strip() or "wallarm-tunnel"
    pw = input("Password or SSH key path: ").strip()
    if not pw: print("No credentials."); return
    print(f"\nShare: ssh -p 9042 {user}@{host}\nCtrl+C to close.\n")
    # Requires paramiko: pip install paramiko
    try:
        import paramiko, socket, threading
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.connect((host.split(":")[0], 443))
        t = paramiko.Transport(sock)
        if pw.startswith("/"): t.connect(username=user, key_filename=pw)
        else: t.connect(username=user, password=pw)
        t.request_port_forward("", 9042, "localhost", 22)
        print("Tunnel active. Ctrl+C to close.")
        threading.Event().wait()
    except ImportError:
        print("Install paramiko: pip3 install paramiko")
    except Exception as e:
        print(f"Tunnel failed: {e}")
def edit_nodes():
    s = load_state().get("nodes", [])
    if not s: print("No nodes."); return
    print("\n─── Edit Node ───")
    for i, n in enumerate(s):
        r = run(f"docker inspect -f '{{{{.State.Status}}}}' wallarm-{n['name']} 2>/dev/null")
        st = "●" if r.stdout.strip() == "running" else "○"
        print(f"  [{i+1}] {st} {n['name']} (port {n.get('port','')})")
    print("  [b] Back")
    c = input("\nSelect node: ").strip()
    if c.lower() == 'b': return
    if not c.isdigit(): return
    idx = int(c) - 1
    if idx < 0 or idx >= len(s): return
    n = s[idx]
    print(f"\nEditing {n['name']}. Enter to keep current:")
    port = input(f"Port [{n.get('port','')}]: ").strip()
    upstream = input(f"Upstream [{n.get('upstream_ip','')}:{n.get('upstream_port','')}]: ").strip()
    if port: n['port'] = int(port)
    if ':' in upstream:
        ip, p = upstream.split(':')
        n['upstream_ip'] = ip; n['upstream_port'] = int(p)
    save_state({"nodes": s})
    print(f"✅ Updated. Restart to apply.")
def remove_menu():
    s = load_state().get("nodes", [])
    if not s: print("No nodes."); return
    print("\n─── Delete Node ───")
    for i, n in enumerate(s):
        print(f"  [{i+1}] {n['name']} (port {n.get('port','')})")
    c = input("\nSelect: ").strip()
    if not c.isdigit(): return
    idx = int(c) - 1
    if 0 <= idx < len(s):
        remove_container(s[idx]["name"])

if __name__ == "__main__":
    main()