#!/usr/bin/env python3
"""
NSX Infrastructure Security Groups
Creates or updates NSX security groups for shared infrastructure servers.

Edit the CONFIGURATION section below before running, then:
    python infra/security_groups.py --nsx-manager <IP> --username admin --password <pass>

Add --dry-run to preview the API payloads without making any changes.
"""

import argparse
import json
import sys
import urllib3
import requests
from typing import Any

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

# ── CONFIGURATION ──────────────────────────────────────────────────────────────

# Prefix applied to every group's display name and ID.
# Change this to match your environment's naming standard.
GROUP_PREFIX = "SG-Infra-"

# NSX tag scope used for all infrastructure server tags.
# Tag format: scope=TAG_SCOPE, value=<server_type>  e.g. scope="Infra", value="DNS Server"
TAG_SCOPE = "Infra"

# NSX Policy domain.  "default" is correct for most environments.
DOMAIN = "default"

# Group definitions.
#
# Each entry key becomes part of the group ID (URL-safe, no spaces).
# Fields:
#   display_name  : suffix appended to GROUP_PREFIX  →  "SG-Infra-DNS-Servers"
#   tag_value     : tag value (under TAG_SCOPE) that dynamically adds matching VMs
#   ip_addresses  : static IPs or CIDR ranges to include  (list of strings)
#   vm_ext_ids    : vCenter VM External IDs / MoRefs to include  (list of strings)
#                   Find these in vCenter (e.g. "vm-1042") or via the NSX VM inventory API.
#
# Members from all three sources are OR'd together in the group expression.

GROUPS: dict[str, dict[str, Any]] = {
    "DNS": {
        "display_name": "DNS-Servers",
        "tag_value": "DNS Server",
        "ip_addresses": [
            # "192.168.10.53",
            # "192.168.10.54",
        ],
        "vm_ext_ids": [
            # "vm-1001",
        ],
    },
    "NTP": {
        "display_name": "NTP-Servers",
        "tag_value": "NTP Server",
        "ip_addresses": [
            # "192.168.10.123",
            # "192.168.10.124",
        ],
        "vm_ext_ids": [
            # "vm-1002",
        ],
    },
    "DC": {
        "display_name": "Domain-Controllers",
        "tag_value": "Domain Controller",
        "ip_addresses": [
            # "192.168.10.10",
            # "192.168.10.11",
        ],
        "vm_ext_ids": [
            # "vm-1003",
            # "vm-1004",
        ],
    },
}

# ── END CONFIGURATION ──────────────────────────────────────────────────────────


def group_id(key: str) -> str:
    """Derive a stable, URL-safe NSX group ID from the config key and prefix."""
    prefix = GROUP_PREFIX.rstrip("-").replace(" ", "-")
    return f"{prefix}-{key}"


def build_expressions(cfg: dict[str, Any]) -> list[dict[str, Any]]:
    """
    Build the NSX Policy expression list for a group.

    Each source type (IPs, tag-matched VMs, explicit VMs) is OR'd together.
    Tag matching uses a NestedExpression to enforce both scope AND value so
    that tags from other scopes with the same value are not inadvertently included.

    Expression order matters in the NSX API — the ConjunctionOperator must sit
    between two non-operator expressions; the first expression must never be
    a ConjunctionOperator.
    """
    expressions: list[dict[str, Any]] = []

    def _or():
        """Append an OR separator between expression blocks."""
        if expressions:
            expressions.append({
                "resource_type": "ConjunctionOperator",
                "conjunction_operator": "OR",
            })

    # ── Static IP addresses / CIDRs ──────────────────────────────────────────
    ip_list = [ip for ip in cfg.get("ip_addresses", []) if ip]
    if ip_list:
        _or()
        expressions.append({
            "resource_type": "IPAddressExpression",
            "ip_addresses": ip_list,
        })

    # ── Dynamic tag membership (scope=TAG_SCOPE AND value=tag_value) ─────────
    # NestedExpression combines two Condition objects with AND:
    #   Condition 1 — matches the tag SCOPE field   (scope_operator triggers this)
    #   Condition 2 — matches the tag VALUE field
    tag_value = cfg.get("tag_value", "").strip()
    if tag_value:
        _or()
        expressions.append({
            "resource_type": "NestedExpression",
            "expressions": [
                {
                    "resource_type": "Condition",
                    "member_type": "VirtualMachine",
                    "key": "Tag",
                    "operator": "EQUALS",
                    "value": TAG_SCOPE,
                    # scope_operator tells NSX to match this Condition against
                    # the tag's scope field instead of the tag's value field.
                    "scope_operator": "EQUALS",
                },
                {
                    "resource_type": "ConjunctionOperator",
                    "conjunction_operator": "AND",
                },
                {
                    "resource_type": "Condition",
                    "member_type": "VirtualMachine",
                    "key": "Tag",
                    "operator": "EQUALS",
                    "value": tag_value,
                },
            ],
        })

    # ── Explicit VMs by vCenter External ID / MoRef ──────────────────────────
    ext_ids = [v for v in cfg.get("vm_ext_ids", []) if v]
    if ext_ids:
        _or()
        expressions.append({
            "resource_type": "ExternalIDExpression",
            "member_type": "VirtualMachine",
            "external_ids": ext_ids,
        })

    return expressions


def build_group_payload(key: str, cfg: dict[str, Any]) -> dict[str, Any]:
    """Assemble the full NSX Policy group PATCH body."""
    expressions = build_expressions(cfg)
    if not expressions:
        raise ValueError(
            f"Group '{key}' has no members defined "
            "(add ip_addresses, tag_value, or vm_ext_ids in CONFIGURATION)."
        )

    return {
        "id": group_id(key),
        "display_name": f"{GROUP_PREFIX}{cfg['display_name']}",
        "description": (
            f"Infrastructure {cfg['display_name'].replace('-', ' ')} — "
            f"managed by nsx_blueprints/infra/security_groups.py"
        ),
        "expression": expressions,
    }


def patch_group(
    session: requests.Session,
    base_url: str,
    key: str,
    payload: dict[str, Any],
    dry_run: bool,
) -> None:
    url = (
        f"{base_url}/policy/api/v1/infra/domains/{DOMAIN}/groups/{group_id(key)}"
    )

    print(f"\n{'[DRY RUN] ' if dry_run else ''}Group: {payload['display_name']}")
    print(f"  URL : PATCH {url}")
    print(f"  Body: {json.dumps(payload, indent=4)}")

    if dry_run:
        return

    resp = session.patch(url, json=payload)
    if resp.status_code in (200, 201):
        print(f"  ✓ {resp.status_code} — group created/updated successfully.")
    else:
        print(f"  ✗ {resp.status_code} — {resp.text}", file=sys.stderr)
        resp.raise_for_status()


def build_session(username: str, password: str, verify_ssl: bool) -> requests.Session:
    s = requests.Session()
    s.auth = (username, password)
    s.verify = verify_ssl
    s.headers.update({"Content-Type": "application/json"})
    return s


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Create/update NSX infrastructure security groups."
    )
    p.add_argument("--nsx-manager", required=True, metavar="HOST",
                   help="NSX Manager IP or FQDN (e.g. 192.168.1.10)")
    p.add_argument("--username", required=True,
                   help="NSX Manager username")
    p.add_argument("--password", required=True,
                   help="NSX Manager password")
    p.add_argument("--dry-run", action="store_true",
                   help="Print payloads without calling the API")
    p.add_argument("--no-verify-ssl", action="store_true",
                   help="Disable SSL certificate verification (useful in lab environments)")
    return p.parse_args()


def main() -> None:
    args = parse_args()
    base_url = f"https://{args.nsx_manager.rstrip('/')}"
    session = build_session(args.username, args.password, not args.no_verify_ssl)

    errors: list[str] = []
    for key, cfg in GROUPS.items():
        try:
            payload = build_group_payload(key, cfg)
            patch_group(session, base_url, key, payload, args.dry_run)
        except ValueError as exc:
            print(f"  ⚠ Skipping '{key}': {exc}")
        except requests.HTTPError as exc:
            errors.append(f"{key}: {exc}")

    if errors:
        print("\nErrors encountered:", file=sys.stderr)
        for err in errors:
            print(f"  - {err}", file=sys.stderr)
        sys.exit(1)

    print("\nDone." if not args.dry_run else "\nDry run complete — no changes were made.")


if __name__ == "__main__":
    main()