ssh-netbox-wrapper/testenv/seed/seed.py

"""Fill both NetBox instances of the netssh test environment.

Idempotent: objects are looked up before they are created, so running the
seed again leaves NetBox unchanged. Writes /state/env (tokens) and
/state/netssh.yaml (base config for the tests).
"""

import json
import os
import sys
import time
import urllib.error
import urllib.request

FSN = "http://10.231.0.20:8080"
LAB = "http://10.231.0.21:8080"
# Superuser v2 token from compose.yaml (SUPERUSER_API_KEY / SUPERUSER_API_TOKEN).
V2_TOKEN = "nbt_netsshtest01.0123456789abcdef0123456789abcdef01234567"
STATE = "/state"


class NetBox:
    def __init__(self, url):
        self.url = url

    def call(self, method, path, body=None):
        data = json.dumps(body).encode() if body is not None else None
        req = urllib.request.Request(self.url + "/api/" + path, data=data, method=method)
        req.add_header("Authorization", "Bearer " + V2_TOKEN)
        req.add_header("Content-Type", "application/json")
        req.add_header("Accept", "application/json")
        try:
            with urllib.request.urlopen(req, timeout=30) as resp:
                raw = resp.read()
                return json.loads(raw) if raw else None
        except urllib.error.HTTPError as err:
            sys.exit(f"{method} {path}: HTTP {err.code}: {err.read().decode(errors='replace')}")

    def find(self, path, **filters):
        query = "&".join(f"{k}={urllib.request.quote(str(v))}" for k, v in filters.items())
        results = self.call("GET", f"{path}?{query}&limit=2")["results"]
        return results[0] if results else None

    def ensure(self, path, lookup, body):
        """Return the object matching lookup, creating it from body if missing."""
        found = self.find(path, **lookup)
        if found:
            return found
        return self.call("POST", path, body)

    def wait(self):
        for _ in range(120):
            try:
                self.call("GET", "status/")
                return
            except SystemExit:
                pass
            except (urllib.error.URLError, ConnectionError, TimeoutError):
                pass
            time.sleep(5)
        sys.exit(f"NetBox {self.url} did not become ready")


def base(nb, site, cluster_type, cluster):
    for tag in ("pve", "web", "bastion", "db", "ipv6"):
        nb.ensure("extras/tags/", {"slug": tag}, {"name": tag, "slug": tag})
    site = nb.ensure("dcim/sites/", {"slug": site}, {"name": site, "slug": site, "status": "active"})
    mfr = nb.ensure("dcim/manufacturers/", {"slug": "generic"}, {"name": "Generic", "slug": "generic"})
    role = nb.ensure("dcim/device-roles/", {"slug": "server"}, {"name": "Server", "slug": "server", "color": "2196f3"})
    dtype = nb.ensure("dcim/device-types/", {"slug": "generic-server"},
                      {"manufacturer": mfr["id"], "model": "Generic Server", "slug": "generic-server"})
    ctype = nb.ensure("virtualization/cluster-types/", {"slug": cluster_type},
                      {"name": cluster_type, "slug": cluster_type})
    cluster = nb.ensure("virtualization/clusters/", {"name": cluster},
                        {"name": cluster, "type": ctype["id"], "status": "active"})
    return {"site": site["id"], "role": role["id"], "type": dtype["id"], "cluster": cluster["id"]}


def device(nb, ids, name, ips, tags=(), cluster=False, comments=""):
    """ips: list of (interface, address, primary) with primary in (None, 4, 6)."""
    body = {"name": name, "role": ids["role"], "device_type": ids["type"], "site": ids["site"],
            "status": "active", "tags": [{"slug": t} for t in tags], "comments": comments}
    if cluster:
        body["cluster"] = ids["cluster"]
    dev = nb.ensure("dcim/devices/", {"name": name}, body)
    primary = {}
    for ifname, address, family in ips:
        iface = nb.ensure("dcim/interfaces/", {"device_id": dev["id"], "name": ifname},
                          {"device": dev["id"], "name": ifname, "type": "1000base-t"})
        ip = nb.ensure("ipam/ip-addresses/", {"address": address},
                       {"address": address, "assigned_object_type": "dcim.interface", "assigned_object_id": iface["id"]})
        if family:
            primary[f"primary_ip{family}"] = ip["id"]
    if primary:
        nb.call("PATCH", f"dcim/devices/{dev['id']}/", primary)


def vm(nb, ids, name, address, tags=()):
    machine = nb.ensure("virtualization/virtual-machines/", {"name": name},
                        {"name": name, "cluster": ids["cluster"], "status": "active",
                         "tags": [{"slug": t} for t in tags]})
    iface = nb.ensure("virtualization/interfaces/", {"virtual_machine_id": machine["id"], "name": "eth0"},
                      {"virtual_machine": machine["id"], "name": "eth0"})
    ip = nb.ensure("ipam/ip-addresses/", {"address": address},
                   {"address": address, "assigned_object_type": "virtualization.vminterface",
                    "assigned_object_id": iface["id"]})
    nb.call("PATCH", f"virtualization/virtual-machines/{machine['id']}/", {"primary_ip4": ip["id"]})


def v1_token(nb):
    """A legacy v1 token for the lab instance, reused from /state/env if present."""
    env = read_env()
    if env.get("LAB_TOKEN"):
        return env["LAB_TOKEN"]
    user = nb.find("users/users/", username="admin")
    token = nb.call("POST", "users/tokens/", {"user": user["id"], "version": 1, "description": "netssh v1 test"})
    plaintext = token.get("token") or token.get("plaintext")
    if not plaintext:
        sys.exit(f"v1 token response without plaintext: {token}")
    return plaintext


def read_env():
    env = {}
    try:
        with open(os.path.join(STATE, "env")) as f:
            for line in f:
                key, _, value = line.strip().partition("=")
                env[key] = value
    except FileNotFoundError:
        pass
    return env


def main():
    fsn, lab = NetBox(FSN), NetBox(LAB)
    fsn.wait()
    lab.wait()

    ids = base(fsn, "fsn1", "proxmox", "pve-fsn1")
    for n in (1, 2, 3):
        device(fsn, ids, f"fsn1-pve{n}.hoster.test",
               [("eth0", f"10.231.2.1{n}/24", 4), ("mgmt0", f"10.231.1.1{n}/24", None)],
               tags=("pve",), cluster=True, comments=f"Proxmox node {n}")
    device(fsn, ids, "bastion.hoster.test", [("eth0", "10.231.0.30/24", 4)], tags=("bastion",))
    device(fsn, ids, "v6-host", [("eth0", "fd00:30::60/64", 6)], tags=("ipv6",))
    vm(fsn, ids, "web1", "10.231.2.21/24", tags=("web",))
    vm(fsn, ids, "web2", "10.231.2.22/24", tags=("web",))

    ids = base(lab, "lab", "kubernetes", "lab-k8s")
    device(lab, ids, "lab-db1", [("eth0", "10.231.1.41/24", 4)], tags=("db",))
    vm(lab, ids, "web1", "10.231.1.42/24", tags=("web",))

    lab_token = v1_token(lab)
    os.makedirs(STATE, exist_ok=True)
    with open(os.path.join(STATE, "env"), "w") as f:
        f.write(f"FSN_URL={FSN}\nLAB_URL={LAB}\nFSN_TOKEN={V2_TOKEN}\nLAB_TOKEN={lab_token}\n")
    with open(os.path.join(STATE, "netssh.yaml"), "w") as f:
        f.write(f"""netboxes:
  - name: fsn
    url: "{FSN}"
    token: "{V2_TOKEN}"
  - name: lab
    url: "{LAB}"
    token: "{lab_token}"
""")
    print("seed: done (fsn: 3 nodes, bastion, v6-host, 2 VMs; lab: lab-db1, web1)")


if __name__ == "__main__":
    main()