#!/usr/bin/env python3
"""Stage UniFi Policy Engine firewall policies for the current segmentation.

Default mode is dry-run. Pass --apply to create/update desired custom policies.
Reads credentials from the runtime automation .env on PD by default.

This helper targets the newer Policy Engine / zone-based firewall model via:
  /proxy/network/v2/api/site/<site>/firewall/zone
  /proxy/network/v2/api/site/<site>/firewall-policies
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import os
from copy import deepcopy
from pathlib import Path
from typing import Any

DEFAULT_RUNTIME_HELPER = Path('/mnt/docker-ssd/docker/compose/automation/bin/unifi_stage_low_risk_objects.py')
DEFAULT_ENV_PATH = Path('/mnt/docker-ssd/docker/compose/automation/.env')


def load_helper(path: Path):
    spec = importlib.util.spec_from_file_location('unifi_stage', path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f'Could not load helper module from {path}')
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def desired_policies() -> list[dict[str, Any]]:
    restricted_subnets = ['10.5.10.0/24', '10.5.20.0/24', '192.168.1.0/24']
    return [
        {
            'name': 'Allow Internal to Untrusted',
            'enabled': True,
            'action': 'ALLOW',
            'protocol': 'all',
            'ip_version': 'BOTH',
            'logging': False,
            'create_allow_respond': True,
            'connection_state_type': 'ALL',
            'connection_states': [],
            'match_ip_sec': False,
            'match_opposite_protocol': False,
            'icmp_typename': 'ANY',
            'icmp_v6_typename': 'ANY',
            'source_zone_name': 'Internal',
            'destination_zone_name': 'Untrusted',
            'schedule': {'mode': 'ALWAYS'},
            'source': {
                'matching_target': 'ANY',
                'port_matching_type': 'ANY',
                'match_opposite_ports': False,
            },
            'destination': {
                'matching_target': 'ANY',
                'port_matching_type': 'ANY',
                'match_opposite_ports': False,
            },
        },
        {
            'name': 'Block Untrusted to Gateway Admin Surfaces',
            'enabled': True,
            'action': 'BLOCK',
            'protocol': 'tcp',
            'ip_version': 'IPV4',
            'logging': False,
            'create_allow_respond': False,
            'connection_state_type': 'ALL',
            'connection_states': [],
            'match_ip_sec': False,
            'match_opposite_protocol': False,
            'icmp_typename': 'ANY',
            'icmp_v6_typename': 'ANY',
            'description': 'Block IoT, Camera, and Old IoT clients from UniFi gateway admin TCP ports while preserving explicit DHCP and mDNS exceptions plus external DNS via 10.5.30.53.',
            'source_zone_name': 'Untrusted',
            'destination_zone_name': 'Gateway',
            'schedule': {'mode': 'ALWAYS'},
            'source': {
                'matching_target': 'IP',
                'matching_target_type': 'SPECIFIC',
                'ips': restricted_subnets,
                'match_mac': False,
                'match_opposite_ips': False,
                'match_opposite_ports': False,
                'port_matching_type': 'ANY',
            },
            'destination': {
                'matching_target': 'IP',
                'matching_target_type': 'SPECIFIC',
                'ips': ['10.5.0.1'],
                'match_opposite_ips': False,
                'match_opposite_ports': False,
                'port_matching_type': 'SPECIFIC',
                'port': '22,80,443,8443,9443',
            },
        },
        {
            'name': 'Allow Untrusted to DNS Resolver',
            'enabled': True,
            'action': 'ALLOW',
            'protocol': 'tcp_udp',
            'ip_version': 'IPV4',
            'logging': False,
            'create_allow_respond': False,
            'connection_state_type': 'ALL',
            'connection_states': [],
            'match_ip_sec': False,
            'match_opposite_protocol': False,
            'icmp_typename': 'ANY',
            'icmp_v6_typename': 'ANY',
            'description': 'Allow IoT, Camera, and Old IoT clients to reach the Pi-hole DNS resolver on PD after DHCP DNS cutover.',
            'source_zone_name': 'Untrusted',
            'destination_zone_name': 'Internal',
            'schedule': {'mode': 'ALWAYS'},
            'source': {
                'matching_target': 'IP',
                'matching_target_type': 'SPECIFIC',
                'ips': restricted_subnets,
                'match_mac': False,
                'match_opposite_ips': False,
                'match_opposite_ports': False,
                'port_matching_type': 'ANY',
            },
            'destination': {
                'matching_target': 'IP',
                'matching_target_type': 'SPECIFIC',
                'ips': ['10.5.30.53'],
                'match_opposite_ips': False,
                'match_opposite_ports': False,
                'port_matching_type': 'SPECIFIC',
                'port': '53',
            },
        },
        {
            'name': 'Block Untrusted to Gateway Default Services',
            'enabled': True,
            'action': 'BLOCK',
            'protocol': 'all',
            'ip_version': 'BOTH',
            'logging': False,
            'create_allow_respond': False,
            'connection_state_type': 'ALL',
            'connection_states': [],
            'match_ip_sec': False,
            'match_opposite_protocol': False,
            'icmp_typename': 'ANY',
            'icmp_v6_typename': 'ANY',
            'description': 'Block general access from IoT, Camera, and Old IoT to gateway services after DNS cutover, preserving only explicit DHCP and mDNS exceptions.',
            'source_zone_name': 'Untrusted',
            'destination_zone_name': 'Gateway',
            'schedule': {'mode': 'ALWAYS'},
            'source': {
                'matching_target': 'IP',
                'matching_target_type': 'SPECIFIC',
                'ips': restricted_subnets,
                'match_mac': False,
                'match_opposite_ips': False,
                'match_opposite_ports': False,
                'port_matching_type': 'ANY',
            },
            'destination': {
                'matching_target': 'ANY',
                'match_opposite_ports': False,
                'port_matching_type': 'ANY',
            },
        },
        {
            'name': 'Allow Untrusted to Gateway DHCP',
            'enabled': True,
            'action': 'ALLOW',
            'protocol': 'udp',
            'ip_version': 'BOTH',
            'logging': False,
            'create_allow_respond': False,
            'connection_state_type': 'ALL',
            'connection_states': [],
            'match_ip_sec': False,
            'match_opposite_protocol': False,
            'icmp_typename': 'ANY',
            'icmp_v6_typename': 'ANY',
            'description': 'Preserve DHCP from IoT, Camera, and Old IoT to the gateway while the broader gateway shield is in place.',
            'source_zone_name': 'Untrusted',
            'destination_zone_name': 'Gateway',
            'schedule': {'mode': 'ALWAYS'},
            'source': {
                'matching_target': 'IP',
                'matching_target_type': 'SPECIFIC',
                'ips': restricted_subnets,
                'match_mac': False,
                'match_opposite_ips': False,
                'match_opposite_ports': False,
                'port_matching_type': 'SPECIFIC',
                'port': '68',
            },
            'destination': {
                'matching_target': 'ANY',
                'match_opposite_ports': False,
                'port_matching_type': 'SPECIFIC',
                'port': '67',
            },
        },
        {
            'name': 'Allow Untrusted to Gateway mDNS',
            'enabled': True,
            'action': 'ALLOW',
            'protocol': 'udp',
            'ip_version': 'IPV4',
            'logging': False,
            'create_allow_respond': False,
            'connection_state_type': 'ALL',
            'connection_states': [],
            'match_ip_sec': False,
            'match_opposite_protocol': False,
            'icmp_typename': 'ANY',
            'icmp_v6_typename': 'ANY',
            'description': 'Preserve mDNS multicast from IoT, Camera, and Old IoT while the broader gateway shield is in place.',
            'source_zone_name': 'Untrusted',
            'destination_zone_name': 'Gateway',
            'schedule': {'mode': 'ALWAYS'},
            'source': {
                'matching_target': 'IP',
                'matching_target_type': 'SPECIFIC',
                'ips': restricted_subnets,
                'match_mac': False,
                'match_opposite_ips': False,
                'match_opposite_ports': False,
                'port_matching_type': 'SPECIFIC',
                'port': '5353',
            },
            'destination': {
                'matching_target': 'IP',
                'matching_target_type': 'SPECIFIC',
                'ips': ['224.0.0.251'],
                'match_opposite_ips': False,
                'match_opposite_ports': False,
                'port_matching_type': 'SPECIFIC',
                'port': '5353',
            },
        },
    ]


POLICY_MUTABLE_KEYS = {
    'name',
    'enabled',
    'action',
    'protocol',
    'ip_version',
    'logging',
    'create_allow_respond',
    'connection_state_type',
    'connection_states',
    'match_ip_sec',
    'match_opposite_protocol',
    'icmp_typename',
    'icmp_v6_typename',
    'source',
    'destination',
    'schedule',
    'description',
}

TEMP_POLICY_NAMES = {'DORIS-TEMP'}


def canonical_policy(policy: dict[str, Any]) -> dict[str, Any]:
    out = {key: deepcopy(policy.get(key)) for key in POLICY_MUTABLE_KEYS if key in policy}
    out.setdefault('description', '')
    out.setdefault('connection_state_type', 'ALL')
    out.setdefault('connection_states', [])
    out.setdefault('logging', False)
    out.setdefault('match_ip_sec', False)
    out.setdefault('match_opposite_protocol', False)
    out.setdefault('icmp_typename', 'ANY')
    out.setdefault('icmp_v6_typename', 'ANY')
    out.setdefault('schedule', {'mode': 'ALWAYS'})
    return out


def resolve_zone_ids(zone_data: list[dict[str, Any]]) -> dict[str, str]:
    by_name: dict[str, str] = {}
    for zone in zone_data:
        name = zone.get('name')
        zid = zone.get('_id')
        if name and zid:
            by_name[name] = zid
    return by_name


def materialize_policy(spec: dict[str, Any], zone_ids: dict[str, str]) -> dict[str, Any]:
    source_zone = spec['source_zone_name']
    dest_zone = spec['destination_zone_name']
    if source_zone not in zone_ids:
        raise KeyError(f'Missing source zone: {source_zone}')
    if dest_zone not in zone_ids:
        raise KeyError(f'Missing destination zone: {dest_zone}')

    policy = {k: deepcopy(v) for k, v in spec.items() if not k.endswith('_zone_name')}
    policy['description'] = policy.get('description', '')
    policy['source'] = deepcopy(policy['source'])
    policy['destination'] = deepcopy(policy['destination'])
    policy['source']['zone_id'] = zone_ids[source_zone]
    policy['destination']['zone_id'] = zone_ids[dest_zone]
    return policy


def main() -> int:
    parser = argparse.ArgumentParser(description='Stage UniFi Policy Engine firewall policies.')
    parser.add_argument('--helper', default=str(DEFAULT_RUNTIME_HELPER), help='Path to runtime UniFi helper on PD')
    parser.add_argument('--env-file', default=str(DEFAULT_ENV_PATH), help='Path to runtime .env on PD')
    parser.add_argument('--site', default='default')
    parser.add_argument('--apply', action='store_true')
    args = parser.parse_args()

    helper = load_helper(Path(args.helper))
    helper.load_env_file(Path(args.env_file))
    client = helper.UniFiClient(
        base_url=os.environ.get('UNIFI_BASE_URL', 'https://10.5.0.1'),
        username=os.environ['UNIFI_USERNAME'],
        password=os.environ['UNIFI_PASSWORD'],
        verify_tls=helper.env_bool('UNIFI_VERIFY_TLS', False),
        timeout=20,
    )
    client.login()

    zones_payload = client.get(f'/proxy/network/v2/api/site/{args.site}/firewall/zone')
    zone_data_raw = zones_payload['data'] if isinstance(zones_payload, dict) and 'data' in zones_payload else zones_payload
    zone_data = zone_data_raw if isinstance(zone_data_raw, list) else []
    zone_ids = resolve_zone_ids(zone_data)

    policies_path = f'/proxy/network/v2/api/site/{args.site}/firewall-policies'
    policies_payload = client.get(policies_path)
    policy_data_raw = policies_payload['data'] if isinstance(policies_payload, dict) and 'data' in policies_payload else policies_payload
    policy_data = policy_data_raw if isinstance(policy_data_raw, list) else []
    custom_by_name = {p.get('name'): p for p in policy_data if not p.get('predefined') and p.get('name')}

    result: dict[str, Any] = {
        'mode': 'apply' if args.apply else 'dry-run',
        'zones': sorted(zone_ids.keys()),
        'changes': [],
    }

    for desired in desired_policies():
        materialized = materialize_policy(desired, zone_ids)
        existing = custom_by_name.get(materialized['name'])
        desired_canonical = canonical_policy(materialized)
        if existing is None:
            result['changes'].append({'name': materialized['name'], 'action': 'create', 'payload': desired_canonical})
            if args.apply:
                status, payload = client._json_request('POST', policies_path, desired_canonical)
                if status >= 300:
                    result['changes'][-1]['error'] = {'status': status, 'payload': payload}
            continue

        existing_canonical = canonical_policy(existing)
        if existing_canonical != desired_canonical:
            update_payload = deepcopy(existing)
            for key, value in desired_canonical.items():
                update_payload[key] = deepcopy(value)
            result['changes'].append({
                'name': materialized['name'],
                'action': 'update',
                'before': existing_canonical,
                'after': desired_canonical,
            })
            if args.apply:
                status, payload = client._json_request('PUT', f"{policies_path}/{existing['_id']}", update_payload)
                if status >= 300:
                    result['changes'][-1]['error'] = {'status': status, 'payload': payload}
        else:
            result['changes'].append({'name': materialized['name'], 'action': 'noop'})

    for name in sorted(TEMP_POLICY_NAMES):
        existing = custom_by_name.get(name)
        if not existing:
            continue
        result['changes'].append({'name': name, 'action': 'delete-temp'})
        if args.apply:
            status, payload = client._json_request('DELETE', f"{policies_path}/{existing['_id']}")
            if status >= 300:
                result['changes'][-1]['error'] = {'status': status, 'payload': payload}

    print(json.dumps(result, indent=2, sort_keys=True))
    return 0


if __name__ == '__main__':
    raise SystemExit(main())
