nsx_rule_export/nsx_migrate.py

#!/usr/bin/env python3
"""
NSX Migration Tool
Exports and imports DFW rules, Security Groups, and Tags from/to NSX Manager (3.x/4.x/9.x)
Supports default infra, VCF Projects, and VPC contexts on both source and destination.
"""

import argparse
import json
import sys
import urllib.request
import urllib.error
import urllib.parse
import base64
import ssl
import os


# ---------------------------------------------------------------------------
# HTTP helpers
# ---------------------------------------------------------------------------

def build_headers(username, password):
    credentials = base64.b64encode(f"{username}:{password}".encode()).decode()
    return {
        "Authorization": f"Basic {credentials}",
        "Content-Type": "application/json",
        "Accept": "application/json",
    }


def make_request(url, headers, method="GET", data=None, verify_ssl=True):
    ctx = ssl.create_default_context() if verify_ssl else ssl._create_unverified_context()
    body = json.dumps(data).encode() if data else None
    req = urllib.request.Request(url, data=body, headers=headers, method=method)
    try:
        with urllib.request.urlopen(req, context=ctx) as resp:
            raw = resp.read()
            return json.loads(raw) if raw else {}
    except urllib.error.HTTPError as e:
        body = e.read().decode(errors="replace")
        raise RuntimeError(f"HTTP {e.code} {e.reason} — {url}\n{body}")


def paginate(base_url, headers, verify_ssl):
    """Collect all results from a paginated NSX endpoint."""
    results = []
    cursor = None
    while True:
        url = base_url + (f"&cursor={cursor}" if cursor else "")
        resp = make_request(url, headers, verify_ssl=verify_ssl)
        results.extend(resp.get("results", []))
        cursor = resp.get("cursor")
        if not cursor:
            break
    return results


# ---------------------------------------------------------------------------
# Context / path builder
# ---------------------------------------------------------------------------

def build_api_paths(base_url, org="default", project=None, vpc=None):
    """
    Build a dict of endpoint base paths for the given NSX context.

    Three contexts are supported:
      1. Default (no project/VPC):
            Policy base  →  /policy/api/v1/infra/domains/default
      2. Project (VCF multi-tenancy, no VPC):
            Policy base  →  /policy/api/v1/orgs/{org}/projects/{project}/infra/domains/default
      3. VPC (isolated network within a project):
            Policy base  →  /policy/api/v1/orgs/{org}/projects/{project}/vpcs/{vpc}
            (VPCs have groups/security-policies directly; no domain level)

    The fabric VM API (/api/v1/fabric/...) is always global — it is not
    scoped to a project or VPC.
    """
    fabric_base = f"{base_url}/api/v1/fabric"

    if vpc:
        if not project:
            raise ValueError("--vpc requires --project to be set")
        policy_base = (
            f"{base_url}/policy/api/v1/orgs/{org}/projects/{project}/vpcs/{vpc}"
        )
        context_label = f"org={org}, project={project}, vpc={vpc}"
    elif project:
        policy_base = (
            f"{base_url}/policy/api/v1/orgs/{org}/projects/{project}/infra/domains/default"
        )
        context_label = f"org={org}, project={project}"
    else:
        policy_base = f"{base_url}/policy/api/v1/infra/domains/default"
        context_label = "default"

    return {
        "context_label":    context_label,
        "fabric_vms":       f"{fabric_base}/virtual-machines",
        "fabric_vm_tags":   f"{fabric_base}/virtual-machines?action=update_tags",
        "groups":           f"{policy_base}/groups",
        "security_policies": f"{policy_base}/security-policies",
    }


def paths_from_args(base_url, org, project, vpc):
    """Validate and build API paths, exiting on misconfiguration."""
    try:
        return build_api_paths(base_url, org=org, project=project, vpc=vpc)
    except ValueError as e:
        print(f"ERROR: {e}", file=sys.stderr)
        sys.exit(1)


# ---------------------------------------------------------------------------
# EXPORT
# ---------------------------------------------------------------------------

def export_tags(paths, headers, verify_ssl):
    """
    NSX tags live on VMs — not as standalone objects.
    Collect all unique tags and their per-VM assignments.
    The fabric VM API is always global (not project/VPC scoped).
    """
    print("  Collecting tags from virtual machines...")
    vms = paginate(f"{paths['fabric_vms']}?page_size=1000", headers, verify_ssl)

    tag_assignments = []
    seen_tags = set()
    all_tags = []

    for vm in vms:
        vm_tags = vm.get("tags", [])
        if vm_tags:
            tag_assignments.append({
                "resource_type": "VirtualMachine",
                "external_id": vm.get("external_id"),
                "display_name": vm.get("display_name"),
                "tags": vm_tags,
            })
            for t in vm_tags:
                key = (t.get("scope", ""), t.get("tag", ""))
                if key not in seen_tags:
                    seen_tags.add(key)
                    all_tags.append({"scope": t.get("scope", ""), "tag": t.get("tag", "")})

    print(f"  Found {len(all_tags)} unique tags across {len(vms)} VMs.")
    return {"tags": all_tags, "assignments": tag_assignments}


def export_security_groups(paths, headers, verify_ssl):
    """Export NSX Groups (security groups) from the resolved Policy API path."""
    print("  Collecting security groups...")
    groups = paginate(f"{paths['groups']}?page_size=1000", headers, verify_ssl)
    print(f"  Found {len(groups)} security groups.")
    return groups


def export_dfw_rules(paths, headers, verify_ssl):
    """Export all DFW security policies and their rules."""
    print("  Collecting DFW security policies...")
    policies = paginate(f"{paths['security_policies']}?page_size=1000", headers, verify_ssl)

    exported = []
    for policy in policies:
        policy_id = policy.get("id")
        print(f"    Fetching rules for policy: {policy.get('display_name', policy_id)}")
        rules = paginate(
            f"{paths['security_policies']}/{policy_id}/rules?page_size=1000",
            headers, verify_ssl,
        )
        exported.append({**policy, "_rules": rules})

    total_rules = sum(len(p["_rules"]) for p in exported)
    print(f"  Found {len(exported)} policies with {total_rules} total rules.")
    return exported


def run_export(args):
    base_url = f"https://{args.host}"
    headers = build_headers(args.username, args.password)
    verify_ssl = not args.no_verify_ssl

    paths = paths_from_args(base_url, args.src_org, args.src_project, args.src_vpc)

    print(f"\n[Export] Connecting to {args.host}")
    print(f"  Context : {paths['context_label']}")

    # Tags
    print("\n[1/3] Exporting tags...")
    tags_data = export_tags(paths, headers, verify_ssl)
    with open(args.tags_file, "w") as f:
        json.dump(tags_data, f, indent=2)
    print(f"  Saved -> {args.tags_file}")

    # Security Groups
    print("\n[2/3] Exporting security groups...")
    groups_data = export_security_groups(paths, headers, verify_ssl)
    with open(args.sec_groups_file, "w") as f:
        json.dump(groups_data, f, indent=2)
    print(f"  Saved -> {args.sec_groups_file}")

    # DFW Rules
    print("\n[3/3] Exporting DFW rules...")
    dfw_data = export_dfw_rules(paths, headers, verify_ssl)
    with open(args.dfw_rules_file, "w") as f:
        json.dump(dfw_data, f, indent=2)
    print(f"  Saved -> {args.dfw_rules_file}")

    print("\n[Export complete]")


# ---------------------------------------------------------------------------
# IMPORT — helpers
# ---------------------------------------------------------------------------

IMPORT_STRIP_KEYS = {
    # Standard read-only metadata
    "_create_time", "_last_modified_time", "_create_user", "_last_modified_user",
    "_system_owned", "_protection", "_revision", "_self",
    # Server-assigned identity fields (also present in nested objects like rule conditions)
    "unique_id", "realization_id", "owner_id",
    # VCF/multi-tenancy path — environment-specific, must not carry across to a different instance
    "remote_path",
}


def clean_for_import(obj):
    """Recursively remove read-only/server-managed fields before PUT/PATCH."""
    if isinstance(obj, dict):
        return {k: clean_for_import(v) for k, v in obj.items() if k not in IMPORT_STRIP_KEYS}
    if isinstance(obj, list):
        return [clean_for_import(i) for i in obj]
    return obj


def object_exists(url, headers, verify_ssl):
    """Return True if the object exists (HTTP 200), False on 404, raise on other errors."""
    try:
        make_request(url, headers, verify_ssl=verify_ssl)
        return True
    except RuntimeError as e:
        if "HTTP 404" in str(e):
            return False
        raise


def prompt_overwrite(name, resource_type, yes_to_all):
    """
    Ask the user whether to overwrite an existing object.
    Returns (should_write: bool, updated_yes_to_all: bool).
    Accepts: y=yes, n=no, a=yes to all remaining.
    """
    if yes_to_all:
        return True, True
    while True:
        try:
            answer = input(
                f"    {resource_type} '{name}' already exists. Overwrite? [y=yes / n=no / a=yes to all]: "
            ).strip().lower()
        except EOFError:
            print("    (non-interactive, defaulting to skip)")
            return False, False
        if answer == "y":
            return True, False
        if answer == "n":
            return False, False
        if answer == "a":
            return True, True


def resolve_vm(paths, headers, verify_ssl, ext_id, display_name):
    """
    Resolve a VM on the destination NSX by external_id first, then fall back
    to display_name. The fabric VM API is always global (not project/VPC scoped).
    Returns (external_id, existing_tags) or (None, []).
    """
    # Primary: match by external_id (same vCenter MoRef)
    try:
        vms = paginate(
            f"{paths['fabric_vms']}?external_id={ext_id}&page_size=10",
            headers, verify_ssl,
        )
        if vms:
            return vms[0]["external_id"], vms[0].get("tags", [])
    except RuntimeError:
        pass

    # Fallback: match by display_name (cross-vCenter)
    if display_name:
        try:
            vms = paginate(
                f"{paths['fabric_vms']}?display_name={urllib.parse.quote(display_name)}&page_size=10",
                headers, verify_ssl,
            )
            if vms:
                if len(vms) > 1:
                    print(f"    WARNING {display_name}: {len(vms)} VMs matched by name — using first result")
                return vms[0]["external_id"], vms[0].get("tags", [])
        except RuntimeError:
            pass

    return None, []


def merge_tag_lists(existing, incoming):
    """Union of two tag lists, deduplicated by (scope, tag) pair."""
    seen = {(t.get("scope", ""), t.get("tag", "")) for t in existing}
    merged = list(existing)
    for t in incoming:
        key = (t.get("scope", ""), t.get("tag", ""))
        if key not in seen:
            merged.append(t)
            seen.add(key)
    return merged


# ---------------------------------------------------------------------------
# IMPORT — resources
# ---------------------------------------------------------------------------

def import_tags(paths, headers, verify_ssl, tags_file, conflict, merge_tags, step_label="[1/3]"):
    """
    Import VM tag assignments.

    conflict:   "overwrite" | "skip" | "prompt"
    merge_tags: if True, union imported tags with existing instead of replacing.
                Ignored when conflict=skip (existing tags are left untouched entirely).
    Note: the fabric VM API used for tag updates is always global, not scoped to
          the destination project/VPC.
    """
    print(f"\n{step_label} Importing tags from {tags_file} ...")
    print(f"  Conflict mode : {conflict}" + (" + merge" if merge_tags else ""))
    with open(tags_file) as f:
        data = json.load(f)

    assignments = data.get("assignments", [])
    ok, skipped, warn, fail = 0, 0, 0, 0
    yes_to_all = False

    for entry in assignments:
        ext_id = entry.get("external_id")
        incoming_tags = entry.get("tags", [])
        name = entry.get("display_name", ext_id)

        resolved_id, existing_tags = resolve_vm(paths, headers, verify_ssl, ext_id, name)

        if not resolved_id:
            print(f"    WARNING {name}: VM not found by external_id or display_name — skipping")
            warn += 1
            continue

        has_existing_tags = bool(existing_tags)

        if has_existing_tags:
            if conflict == "skip":
                print(f"    SKIP  {name}: VM already has {len(existing_tags)} tag(s)")
                skipped += 1
                continue
            elif conflict == "prompt":
                do_write, yes_to_all = prompt_overwrite(name, "VM tags", yes_to_all)
                if not do_write:
                    skipped += 1
                    continue

        if has_existing_tags and merge_tags:
            final_tags = merge_tag_lists(existing_tags, incoming_tags)
            tag_note = (
                f"{len(final_tags)} tags "
                f"({len(incoming_tags)} imported + {len(existing_tags)} existing, merged)"
            )
        else:
            final_tags = incoming_tags
            tag_note = f"{len(final_tags)} tags"

        payload = {"external_id": resolved_id, "tags": final_tags}
        try:
            make_request(paths["fabric_vm_tags"], headers, method="POST", data=payload, verify_ssl=verify_ssl)
            print(f"    OK    {name} ({tag_note})")
            ok += 1
        except RuntimeError as e:
            print(f"    FAIL  {name}: {e}")
            fail += 1

    print(f"  Tags: {ok} ok, {skipped} skipped, {warn} warnings, {fail} failed.")


def import_security_groups(paths, headers, verify_ssl, sec_groups_file, conflict, step_label="[2/3]"):
    """
    Import security groups into the resolved destination context.

    conflict: "overwrite" | "skip" | "prompt"
    """
    print(f"\n{step_label} Importing security groups from {sec_groups_file} ...")
    print(f"  Conflict mode : {conflict}")
    print(f"  Target        : {paths['groups']}")
    with open(sec_groups_file) as f:
        groups = json.load(f)

    ok, skipped, fail = 0, 0, 0
    yes_to_all = False

    for group in groups:
        group_id = group.get("id")
        name = group.get("display_name", group_id)
        group_url = f"{paths['groups']}/{group_id}"

        if conflict != "overwrite":
            exists = object_exists(group_url, headers, verify_ssl)
            if exists:
                if conflict == "skip":
                    print(f"    SKIP  {name}: group already exists")
                    skipped += 1
                    continue
                elif conflict == "prompt":
                    do_write, yes_to_all = prompt_overwrite(name, "Security Group", yes_to_all)
                    if not do_write:
                        skipped += 1
                        continue

        payload = clean_for_import(group)
        try:
            make_request(group_url, headers, method="PATCH", data=payload, verify_ssl=verify_ssl)
            print(f"    OK    {name}")
            ok += 1
        except RuntimeError as e:
            print(f"    FAIL  {name}: {e}")
            fail += 1

    print(f"  Security groups: {ok} ok, {skipped} skipped, {fail} failed.")


def import_dfw_rules(paths, headers, verify_ssl, dfw_rules_file, conflict, step_label="[3/3]"):
    """
    Import DFW security policies and their rules into the resolved destination context.

    conflict: "overwrite" | "skip" | "prompt"
    Rules inside a brand-new policy are always written without individual checks.
    Rules inside an existing policy are individually checked against the conflict mode.
    """
    print(f"\n{step_label} Importing DFW rules from {dfw_rules_file} ...")
    print(f"  Conflict mode : {conflict}")
    print(f"  Target        : {paths['security_policies']}")
    with open(dfw_rules_file) as f:
        policies = json.load(f)

    ok_policies, skip_policies, fail_policies = 0, 0, 0
    ok_rules, skip_rules, fail_rules = 0, 0, 0
    yes_to_all = False

    for policy in policies:
        policy_id = policy.get("id")
        policy_name = policy.get("display_name", policy_id)
        rules = policy.pop("_rules", [])
        policy_url = f"{paths['security_policies']}/{policy_id}"

        policy_existed = False
        if conflict != "overwrite":
            policy_existed = object_exists(policy_url, headers, verify_ssl)
            if policy_existed:
                if conflict == "skip":
                    print(f"    SKIP  policy: {policy_name} (and its {len(rules)} rule(s))")
                    skip_policies += 1
                    skip_rules += len(rules)
                    continue
                elif conflict == "prompt":
                    do_write, yes_to_all = prompt_overwrite(policy_name, "DFW Policy", yes_to_all)
                    if not do_write:
                        skip_policies += 1
                        skip_rules += len(rules)
                        continue

        payload = clean_for_import(policy)
        try:
            make_request(policy_url, headers, method="PATCH", data=payload, verify_ssl=verify_ssl)
            print(f"    OK    policy: {policy_name}")
            ok_policies += 1
        except RuntimeError as e:
            print(f"    FAIL  policy: {policy_name}: {e}")
            fail_policies += 1
            continue

        for rule in rules:
            rule_id = rule.get("id")
            rule_name = rule.get("display_name", rule_id)
            rule_url = f"{policy_url}/rules/{rule_id}"

            if conflict != "overwrite" and policy_existed:
                rule_exists = object_exists(rule_url, headers, verify_ssl)
                if rule_exists:
                    if conflict == "skip":
                        print(f"      SKIP  rule: {rule_name}")
                        skip_rules += 1
                        continue
                    elif conflict == "prompt":
                        do_write, yes_to_all = prompt_overwrite(rule_name, "DFW Rule", yes_to_all)
                        if not do_write:
                            skip_rules += 1
                            continue

            rule_payload = clean_for_import(rule)
            try:
                make_request(rule_url, headers, method="PATCH", data=rule_payload, verify_ssl=verify_ssl)
                ok_rules += 1
            except RuntimeError as e:
                print(f"      FAIL  rule: {rule_name}: {e}")
                fail_rules += 1

    print(f"  Policies: {ok_policies} ok, {skip_policies} skipped, {fail_policies} failed.")
    print(f"  Rules:    {ok_rules} ok, {skip_rules} skipped, {fail_rules} failed.")


def run_import(args):
    base_url = f"https://{args.host}"
    headers = build_headers(args.username, args.password)
    verify_ssl = not args.no_verify_ssl
    conflict = args.conflict
    merge_tags = args.merge_tags

    paths = paths_from_args(base_url, args.dst_org, args.dst_project, args.dst_vpc)

    print(f"\n[Import] Connecting to {args.host}")
    print(f"  Context : {paths['context_label']}")

    action = args.action

    # Validate required input files before starting
    if action in ("import", "import-tags"):
        if not os.path.exists(args.tags_file):
            print(f"ERROR: Tags file not found: {args.tags_file}", file=sys.stderr)
            sys.exit(1)
    if action in ("import", "import-groups"):
        if not os.path.exists(args.sec_groups_file):
            print(f"ERROR: Security groups file not found: {args.sec_groups_file}", file=sys.stderr)
            sys.exit(1)
    if action == "import":
        if not os.path.exists(args.dfw_rules_file):
            print(f"ERROR: DFW rules file not found: {args.dfw_rules_file}", file=sys.stderr)
            sys.exit(1)

    if action == "import-tags":
        import_tags(paths, headers, verify_ssl, args.tags_file, conflict, merge_tags, step_label="[1/1]")

    elif action == "import-groups":
        import_security_groups(paths, headers, verify_ssl, args.sec_groups_file, conflict, step_label="[1/1]")

    else:
        # Full import — order matters: tags → groups → rules
        import_tags(paths, headers, verify_ssl, args.tags_file, conflict, merge_tags, step_label="[1/3]")
        import_security_groups(paths, headers, verify_ssl, args.sec_groups_file, conflict, step_label="[2/3]")
        import_dfw_rules(paths, headers, verify_ssl, args.dfw_rules_file, conflict, step_label="[3/3]")

    print("\n[Import complete]")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="NSX Manager export/import tool for DFW rules, security groups, and tags.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Actions:
  export          Export all resources (tags, security groups, DFW rules) to JSON files
  import          Import all resources in order: tags → security groups → DFW rules
  import-tags     Import only tags (applies VM tag assignments)
  import-groups   Import only security groups

Context (project / VPC):
  Use --src-* args for export, --dst-* args for import.
  If no project is specified the default infra space is used.
  --dst-vpc / --src-vpc requires the corresponding --*-project to also be set.

  Context resolution:
    No project/VPC   →  /policy/api/v1/infra/domains/default/...
    Project only     →  /policy/api/v1/orgs/{org}/projects/{project}/infra/domains/default/...
    Project + VPC    →  /policy/api/v1/orgs/{org}/projects/{project}/vpcs/{vpc}/...

  Note: VM tag lookups always use the global fabric API regardless of project/VPC context.

Conflict modes (--conflict):
  overwrite       Always write, replacing any existing object           [default]
  skip            Leave existing objects untouched, import only new ones
  prompt          Ask interactively for each conflict (y / n / a=yes-to-all)

Tag merge (--merge-tags, import/import-tags only):
  Without flag    Replace all tags on the VM with the imported set
  With flag       Union imported tags with existing VM tags (no removals)

Examples:
  # Export from default space
  python nsx_migrate.py export --host 192.168.1.10 -u admin -p VMware1! --no-verify-ssl

  # Export from a VCF project
  python nsx_migrate.py export --host 192.168.1.10 -u admin -p VMware1! --no-verify-ssl \\
      --src-project my-project

  # Export from a VPC inside a project
  python nsx_migrate.py export --host 192.168.1.10 -u admin -p VMware1! --no-verify-ssl \\
      --src-project my-project --src-vpc my-vpc

  # Import to a specific project on the destination
  python nsx_migrate.py import --host 192.168.2.10 -u admin -p VMware1! --no-verify-ssl \\
      --dst-project my-project --conflict skip

  # Import to a VPC on the destination
  python nsx_migrate.py import --host 192.168.2.10 -u admin -p VMware1! --no-verify-ssl \\
      --dst-project my-project --dst-vpc my-vpc

  # Cross-project migration: export from project-a, import to project-b (same host)
  python nsx_migrate.py export --host 192.168.1.10 -u admin -p VMware1! --no-verify-ssl \\
      --src-project project-a
  python nsx_migrate.py import --host 192.168.1.10 -u admin -p VMware1! --no-verify-ssl \\
      --dst-project project-b --conflict skip

  # Import only tags with merge, into a specific project
  python nsx_migrate.py import-tags --host 192.168.2.10 -u admin -p VMware1! --no-verify-ssl \\
      --dst-project my-project --merge-tags

  # Import only security groups, skip existing, into a VPC
  python nsx_migrate.py import-groups --host 192.168.2.10 -u admin -p VMware1! --no-verify-ssl \\
      --dst-project my-project --dst-vpc my-vpc --conflict skip
""",
    )

    parser.add_argument(
        "action",
        choices=["export", "import", "import-tags", "import-groups"],
        help="Action to perform",
    )

    # Connection
    conn = parser.add_argument_group("Connection")
    conn.add_argument("--host", required=True, help="NSX Manager hostname or IP")
    conn.add_argument("--username", "-u", required=True, help="NSX Manager username")
    conn.add_argument("--password", "-p", required=True, help="NSX Manager password")
    conn.add_argument(
        "--no-verify-ssl", action="store_true", default=False,
        help="Disable SSL certificate verification (for self-signed certs)",
    )

    # Source context (used by export)
    src = parser.add_argument_group("Source context (export)")
    src.add_argument("--src-org", default="default", metavar="ORG",
                     help="Source VCF org (default: 'default')")
    src.add_argument("--src-project", default=None, metavar="PROJECT",
                     help="Source VCF project name (omit for default infra space)")
    src.add_argument("--src-vpc", default=None, metavar="VPC",
                     help="Source VPC name within the project (requires --src-project)")

    # Destination context (used by import)
    dst = parser.add_argument_group("Destination context (import)")
    dst.add_argument("--dst-org", default="default", metavar="ORG",
                     help="Destination VCF org (default: 'default')")
    dst.add_argument("--dst-project", default=None, metavar="PROJECT",
                     help="Destination VCF project name (omit for default infra space)")
    dst.add_argument("--dst-vpc", default=None, metavar="VPC",
                     help="Destination VPC name within the project (requires --dst-project)")

    # Import behaviour
    imp = parser.add_argument_group("Import behaviour")
    imp.add_argument(
        "--conflict", choices=["overwrite", "skip", "prompt"], default="overwrite",
        help="How to handle objects that already exist on the destination (default: overwrite)",
    )
    imp.add_argument(
        "--merge-tags", action="store_true", default=False,
        help="Merge imported tags with existing VM tags instead of replacing (import/import-tags only)",
    )

    # Files
    files = parser.add_argument_group("File paths")
    files.add_argument("--dfw-rules-file", default="dfw_rules.json",
                       help="DFW rules file (default: dfw_rules.json)")
    files.add_argument("--sec-groups-file", default="sec_groups.json",
                       help="Security groups file (default: sec_groups.json)")
    files.add_argument("--tags-file", default="tags.json",
                       help="Tags file (default: tags.json)")

    args = parser.parse_args()

    if args.action == "export":
        run_export(args)
    else:
        run_import(args)


if __name__ == "__main__":
    main()