wallarm/python/deploy.py

#!/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"
WALLARM_IMAGE = "wallarm/node:6.13.0"
DOCKER_BIN_URL = "https://git.sechpoint.app/customer-engineering/wallarm/raw/branch/main/docker/binaries/docker-29.2.1.tgz"

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 if exists
    container = f"wallarm-{name}"
    run(f"docker rm -f {container} 2>/dev/null", timeout=10)

    # Run container
    monitoring_port = int(port) + 10
    print(f"[{name}] Starting container on port {port}...")
    cmd = f"""docker run -d \
        --name {container} \
        --restart unless-stopped \
        -p {port}:{port} \
        -p {monitoring_port}:{monitoring_port} \
        -e WALLARM_API_TOKEN={token} \
        -e WALLARM_API_HOST={api_host} \
        -e WALLARM_MODE={mode} \
        -e NGINX_PORT={port}"""
    if upstream_ip:
        cmd += f" -e WALLARM_UPSTREAM=http://{upstream_ip}:{upstream_port}"
    cmd += f" {WALLARM_IMAGE}"

    r = run(cmd, 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 Docker Manager v{VERSION} ═══\n")
    
    # Check for fw.conf
    cfg = load_config()
    if not cfg:
        print("No fw.conf found. Create one? [Y/n]: ", end="")
        if input().strip().lower() in ("", "y", "yes"):
            create_config()
            cfg = load_config()
        else:
            print("Create /opt/fw/app/fw.conf manually or use the menu to add nodes.")

    while True:
        print("\n─── Menu ───")
        print("  [1] Deploy all")
        print("  [2] Add a node")
        print("  [3] Status")
        print("  [4] Remove a node")
        print("  [5] Edit a node")
        print("  [q] Quit")
        c = input("\nChoice: ").strip()
        if c == "1": deploy_all()
        elif c == "2": add_node()
        elif c == "3": container_status()
        elif c == "4": remove_menu()
        elif c == "5": edit_menu()
        elif c.lower() == "q": break

def create_config():
    print("\n─── Create Node Config ───")
    name = input("Node name: ").strip()
    if not name: return
    token = input("Wallarm token: ").strip()
    cloud = input("Cloud [EU]: ").strip() or "EU"
    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"
    mode = input("Mode [monitoring]: ").strip() or "monitoring"
    
    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)
    print(f"✅ {name} added to fw.conf")

def add_node():
    create_config()
    cfg = load_config()
    if cfg:
        # Deploy just the newly added node
        for name in list(cfg.keys()):
            if not any(n["name"] == name for n in load_state().get("nodes", [])):
                deploy_container(name, cfg[name]["token"], cfg[name].get("cloud","EU"),
                                str(cfg[name].get("port","8081")), cfg[name].get("upstream_ip",""),
                                str(cfg[name].get("upstream_port","80")), "",
                                cfg[name].get("mode","monitoring"))

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.")

if __name__ == "__main__":
    main()