workflow_steps.py

"""Workflow steps that are shared across multiple workflows."""

import json
from typing import Any

from orchestrator import inputstep, step
from orchestrator.config.assignee import Assignee
from orchestrator.types import State, UUIDstr
from orchestrator.utils.json import json_dumps
from pydantic import ConfigDict
from pydantic_forms.core import FormPage
from pydantic_forms.types import FormGenerator
from pydantic_forms.validators import Label

from gso.products.product_blocks.router import RouterRole
from gso.products.product_types.iptrunk import Iptrunk
from gso.products.product_types.router import Router
from gso.services import lso_client, subscriptions
from gso.settings import load_oss_params


def _deploy_base_config(
    subscription: dict[str, Any],
    tt_number: str,
    callback_route: str,
    process_id: UUIDstr,
    *,
    dry_run: bool,
) -> None:
    inventory = subscription["router"]["router_fqdn"]

    extra_vars = {
        "wfo_router_json": subscription,
        "dry_run": dry_run,
        "verb": "deploy",
        "commit_comment": f"GSO_PROCESS_ID: {process_id} - TT_NUMBER: {tt_number} - Deploy base config",
    }

    lso_client.execute_playbook(
        playbook_name="base_config.yaml",
        callback_route=callback_route,
        inventory=inventory,
        extra_vars=extra_vars,
    )


@step("[DRY RUN] Deploy base config")
def deploy_base_config_dry(
    subscription: dict[str, Any],
    tt_number: str,
    callback_route: str,
    process_id: UUIDstr,
) -> State:
    """Perform a dry run of provisioning base config on a router."""
    _deploy_base_config(subscription, tt_number, callback_route, process_id, dry_run=True)

    return {"subscription": subscription}


@step("[FOR REAL] Deploy base config")
def deploy_base_config_real(
    subscription: dict[str, Any],
    tt_number: str,
    callback_route: str,
    process_id: UUIDstr,
) -> State:
    """Deploy base config on a router using the provisioning proxy."""
    _deploy_base_config(subscription, tt_number, callback_route, process_id, dry_run=False)

    return {"subscription": subscription}


@step("[FOR REAL] Set ISIS metric to very high value")
def set_isis_to_max(subscription: Iptrunk, process_id: UUIDstr, callback_route: str, tt_number: str) -> State:
    """Workflow step for setting the :term:`ISIS` metric to an arbitrarily high value to drain a link."""
    old_isis_metric = subscription.iptrunk.iptrunk_isis_metric
    params = load_oss_params()
    subscription.iptrunk.iptrunk_isis_metric = params.GENERAL.isis_high_metric
    extra_vars = {
        "wfo_trunk_json": json.loads(json_dumps(subscription)),
        "dry_run": False,
        "verb": "deploy",
        "config_object": "isis_interface",
        "commit_comment": f"GSO_PROCESS_ID: {process_id} - TT_NUMBER: {tt_number} - Deploy config for "
        f"{subscription.iptrunk.geant_s_sid}",
    }

    lso_client.execute_playbook(
        playbook_name="iptrunks.yaml",
        callback_route=callback_route,
        inventory=f"{subscription.iptrunk.iptrunk_sides[0].iptrunk_side_node.router_fqdn}\n"
        f"{subscription.iptrunk.iptrunk_sides[1].iptrunk_side_node.router_fqdn}\n",
        extra_vars=extra_vars,
    )

    return {
        "subscription": subscription,
        "old_isis_metric": old_isis_metric,
    }


@step("Run show commands after base config install")
def run_checks_after_base_config(subscription: dict[str, Any], callback_route: str) -> None:
    """Workflow step for running show commands after installing base config."""
    lso_client.execute_playbook(
        playbook_name="base_config_checks.yaml",
        callback_route=callback_route,
        inventory=subscription["router"]["router_fqdn"],
        extra_vars={"wfo_router_json": subscription},
    )


@inputstep("Prompt for new SharePoint checklist", assignee=Assignee.SYSTEM)
def prompt_sharepoint_checklist_url(checklist_url: str) -> FormGenerator:
    """Prompt the operator with the checklist in SharePoint for approving a new subscription."""

    class SharepointPrompt(FormPage):
        model_config = ConfigDict(title="Complete new checklist")

        info_label_1: Label = f"A new checklist has been created at: {checklist_url}"
        info_label_2: Label = "Click proceed to finish the workflow."

    yield SharepointPrompt

    return {}


@step("[DRY RUN] Add all PE routers to P router iBGP group")
def add_all_pe_to_p_dry(subscription: dict[str, Any], callback_route: str) -> None:
    """Perform a dry run of adding the list of all PE routers to the new P router."""
    extra_vars = {
        "dry_run": True,
        "subscription": subscription,
        "pe_router_list": subscriptions.get_active_pe_router_dict(),
        "verb": "add_pe_to_p",
    }

    lso_client.execute_playbook(
        playbook_name="update_ibgp_mesh.yaml",
        callback_route=callback_route,
        inventory=subscription["router"]["router_fqdn"],
        extra_vars=extra_vars,
    )


@step("[FOR REAL] Add all PE routers to P router iBGP group")
def add_all_pe_to_p_real(
    subscription: dict[str, Any], callback_route: str, tt_number: str, process_id: UUIDstr
) -> None:
    """Add the list of all PE routers to the new P router."""
    extra_vars = {
        "dry_run": False,
        "subscription": subscription,
        "pe_router_list": subscriptions.get_active_pe_router_dict(),
        "commit_comment": f"GSO_PROCESS_ID: {process_id} - TT_NUMBER: {tt_number} - Update iBGP mesh",
        "verb": "add_pe_to_p",
    }

    lso_client.execute_playbook(
        playbook_name="update_ibgp_mesh.yaml",
        callback_route=callback_route,
        inventory=subscription["router"]["router_fqdn"],
        extra_vars=extra_vars,
    )


@step("Calculate list of all active PE routers")
def calculate_pe_router_list() -> State:
    """Calculate a list of all active PE routers in the network."""
    all_routers = [
        Router.from_subscription(r["subscription_id"]) for r in subscriptions.get_active_router_subscriptions()
    ]
    all_pe_routers = [router for router in all_routers if router.router.router_role == RouterRole.PE]

    return {"pe_router_list": all_pe_routers}


@step("Calculate list of all active P routers")
def calculate_p_router_list() -> State:
    """Calculate a list of all active P routers in the network."""
    all_routers = [
        Router.from_subscription(r["subscription_id"]) for r in subscriptions.get_active_router_subscriptions()
    ]
    all_pe_routers = [router for router in all_routers if router.router.router_role == RouterRole.P]

    return {"pe_router_list": all_pe_routers}


def generate_inventory(router_list: list[Router]) -> dict[str, Any]:
    """Generate an Ansible-compatible inventory for executing playbooks. Contains all active routers."""
    return {
        "all": {
            "hosts": {
                router.router.router_fqdn: {
                    "lo4": str(router.router.router_lo_ipv4_address),
                    "lo6": str(router.router.router_lo_ipv6_address),
                    "vendor": str(router.router.vendor),
                }
                for router in router_list
            }
        },
    }