"""A lean WS-Management client for Intel AMT.

AMT speaks WS-Management (DMTF DSP0226) with the CIM binding: SOAP 1.2 envelopes POSTed to
``/wsman``, on port 16992 over plain HTTP or 16993 over TLS, authenticated with HTTP Digest.
This client covers the four operations a driver needs: ``Get`` one instance, ``Enumerate``
plus ``Pull`` all instances of a class, and invoke a method. Instances come back as plain
dicts keyed by local element names; a repeated element becomes a list, an element with
children a nested dict, and every value a string, as on the wire.

The envelopes follow Intel's own go-wsman-messages, which every AMT since version 6 accepts.
"""

from __future__ import annotations

import itertools
from collections.abc import Callable, Iterable, Mapping
from typing import TYPE_CHECKING, Any

# escape and quoteattr build the request envelopes; nothing here parses with SAX.
from xml.sax.saxutils import escape, quoteattr  # nosec B406

import requests
from defusedxml import DefusedXmlException
from defusedxml.ElementTree import ParseError, fromstring
from requests.auth import HTTPDigestAuth

if TYPE_CHECKING:
    # For the annotations only; answers are parsed by defusedxml.
    import xml.etree.ElementTree as ET  # nosec B405

SOAP_NS = "http://www.w3.org/2003/05/soap-envelope"
ADDRESSING_NS = "http://schemas.xmlsoap.org/ws/2004/08/addressing"
WSMAN_NS = "http://schemas.dmtf.org/wbem/wsman/1/wsman.xsd"
ENUMERATION_NS = "http://schemas.xmlsoap.org/ws/2004/09/enumeration"
TRANSFER_NS = "http://schemas.xmlsoap.org/ws/2004/09/transfer"
ANONYMOUS = ADDRESSING_NS + "/role/anonymous"

PORT = 16992
TLS_PORT = 16993

#: Class name prefix -> the schema its resource URI lives under.
_SCHEMAS = {
    "AMT_": "http://intel.com/wbem/wscim/1/amt-schema/1/",
    "IPS_": "http://intel.com/wbem/wscim/1/ips-schema/1/",
    "CIM_": "http://schemas.dmtf.org/wbem/wscim/1/cim-schema/2/",
}
#: Pages an enumeration may take. AMT's longest lists are a few dozen instances, and it
#: returns all of them in one page when asked for 999; this only stops a runaway loop.
_MAX_PAGES = 100
_HEADERS = {"Content-Type": "application/soap+xml;charset=UTF-8"}

Instance = dict[str, Any]


class WsmanError(Exception):
    """AMT did not answer, or answered with something other than the expected response."""


class WsmanAuthError(WsmanError):
    """AMT refused the credentials (HTTP 401 after the Digest exchange)."""


class WsmanFault(WsmanError):
    """AMT answered with a SOAP fault."""

    def __init__(self, code: str, subcode: str, reason: str, detail: str) -> None:
        self.code = code
        self.subcode = subcode
        self.reason = reason
        self.detail = detail
        super().__init__(f"{subcode or code}: {reason}" + (f" ({detail})" if detail else ""))


def resource_uri(resource: str) -> str:
    """The resource URI of a class, from its prefix (``AMT_``, ``IPS_``, ``CIM_``)."""
    for prefix, base in _SCHEMAS.items():
        if resource.startswith(prefix):
            return base + resource
    raise ValueError(f"{resource!r} is not an AMT, IPS or CIM class")


def _local(tag: str) -> str:
    return tag.rsplit("}", 1)[-1]


def _value(element: ET.Element) -> str | Instance:
    """An element as a string (no children) or a dict of its children."""
    children = list(element)
    if not children:
        return (element.text or "").strip()
    result: Instance = {}
    for child in children:
        key, value = _local(child.tag), _value(child)
        if key not in result:
            result[key] = value
        elif isinstance(result[key], list):
            result[key].append(value)
        else:
            result[key] = [result[key], value]
    return result


def to_instance(element: ET.Element | None) -> Instance:
    """An instance element as a dict; an empty or missing one as an empty dict."""
    if element is None:
        return {}
    value = _value(element)
    return value if isinstance(value, dict) else {}


def _fault(element: ET.Element) -> WsmanFault:
    def text(path: str) -> str:
        found = element.find(path)
        return "".join(found.itertext()).strip() if found is not None else ""

    def unprefixed(qname: str) -> str:
        return qname.split(":", 1)[-1]

    s = f"{{{SOAP_NS}}}"
    return WsmanFault(
        code=unprefixed(text(f"{s}Code/{s}Value")),
        subcode=unprefixed(text(f"{s}Code/{s}Subcode/{s}Value")),
        reason=text(f"{s}Reason/{s}Text"),
        detail=text(f"{s}Detail"),
    )


class WsmanTransport:
    """WS-Management over HTTP(S) with Digest auth, to one AMT.

    ``tls`` defaults to on for port 16993 and off otherwise. ``verify`` is passed to
    ``requests``: ``True`` checks the certificate against its CA bundle, a path against
    that CA file, ``False`` not at all. ``retries`` counts extra attempts after a
    connection error; a response, even an error, is never retried. ``recorder``, when set,
    is called with the class, the operation and the raw body of every answer (to record
    fixtures).
    """

    def __init__(
        self,
        host: str,
        username: str,
        password: str,
        port: int | None = None,
        tls: bool | None = None,
        verify: bool | str = True,
        timeout: float = 60,
        retries: int = 1,
    ) -> None:
        self.port = int(port or PORT)
        self.tls = self.port == TLS_PORT if tls is None else bool(tls)
        self.url = f"{'https' if self.tls else 'http'}://{host}:{self.port}/wsman"
        self.timeout = timeout
        self.retries = max(0, int(retries))
        self.session = requests.Session()
        self.session.auth = HTTPDigestAuth(username, password)
        self.session.verify = verify
        self._message_ids = itertools.count(1)
        self.recorder: Callable[[str, str, bytes], None] | None = None

    def close(self) -> None:
        self.session.close()

    # --- operations ----------------------------------------------------------------------

    def get(self, resource: str, selectors: Mapping[str, str] | None = None) -> Instance:
        """The one instance of ``resource`` (or the one ``selectors`` pick)."""
        body = self._post(resource, f"{TRANSFER_NS}/Get", "", selectors)
        return to_instance(next(iter(body), None))

    def enumerate(self, resource: str) -> list[Instance]:
        """Every instance of ``resource``: ``Enumerate``, then ``Pull`` to the end."""
        e = f"{{{ENUMERATION_NS}}}"
        body = self._post(
            resource, f"{ENUMERATION_NS}/Enumerate", f'<Enumerate xmlns="{ENUMERATION_NS}"/>'
        )
        context = body.findtext(f"{e}EnumerateResponse/{e}EnumerationContext")
        instances: list[Instance] = []
        for _ in range(_MAX_PAGES):
            if not context:
                return instances
            body = self._post(
                resource,
                f"{ENUMERATION_NS}/Pull",
                f'<Pull xmlns="{ENUMERATION_NS}"><EnumerationContext>{escape(context)}'
                f"</EnumerationContext><MaxElements>999</MaxElements>"
                f"<MaxCharacters>99999</MaxCharacters></Pull>",
            )
            page = body.find(f"{e}PullResponse")
            if page is None:
                raise WsmanError(f"{resource}: Pull answered without a PullResponse")
            instances.extend(to_instance(item) for item in page.findall(f"{e}Items/*"))
            if page.find(f"{e}EndOfSequence") is not None:
                return instances
            context = page.findtext(f"{e}EnumerationContext")
        raise WsmanError(f"{resource}: no end of the enumeration after {_MAX_PAGES} pages")

    def invoke(
        self,
        resource: str,
        method: str,
        params: Iterable[tuple[str, Any]] = (),
        raw: str = "",
        selectors: Mapping[str, str] | None = None,
        check: bool = True,
    ) -> Instance:
        """Call ``method`` of ``resource`` and return its output parameters.

        ``params`` are sent in order as ``<h:Name>value</h:Name>``; ``raw`` is appended
        unescaped after them, for parameters that are references (``h:`` is the class's
        namespace). With ``check``, a ``ReturnValue`` other than 0 raises.
        """
        uri = resource_uri(resource)
        inner = "".join(f"<h:{name}>{escape(str(value))}</h:{name}>" for name, value in params)
        body = self._post(
            resource,
            f"{uri}/{method}",
            f'<h:{method}_INPUT xmlns:h="{uri}">{inner}{raw}</h:{method}_INPUT>',
            selectors,
        )
        output = to_instance(body.find(f"{{{uri}}}{method}_OUTPUT"))
        code = output.get("ReturnValue")
        if check and code not in (None, "0"):
            raise WsmanError(f"{resource}.{method} failed: ReturnValue {code}")
        return output

    # --- plumbing ------------------------------------------------------------------------

    def _envelope(
        self, uri: str, action: str, body: str, selectors: Mapping[str, str] | None
    ) -> str:
        selector_set = ""
        if selectors:
            selector_set = (
                "<w:SelectorSet>"
                + "".join(
                    f"<w:Selector Name={quoteattr(name)}>{escape(value)}</w:Selector>"
                    for name, value in selectors.items()
                )
                + "</w:SelectorSet>"
            )
        return (
            f'<?xml version="1.0" encoding="utf-8"?>'
            f'<Envelope xmlns="{SOAP_NS}" xmlns:a="{ADDRESSING_NS}" xmlns:w="{WSMAN_NS}">'
            f"<Header><a:Action>{action}</a:Action><a:To>/wsman</a:To>"
            f"<w:ResourceURI>{uri}</w:ResourceURI>"
            f"<a:MessageID>{next(self._message_ids)}</a:MessageID>"
            f"<a:ReplyTo><a:Address>{ANONYMOUS}</a:Address></a:ReplyTo>"
            f"<w:OperationTimeout>PT60S</w:OperationTimeout>{selector_set}</Header>"
            f"<Body>{body}</Body></Envelope>"
        )

    def _post(
        self, resource: str, action: str, body: str, selectors: Mapping[str, str] | None = None
    ) -> ET.Element:
        envelope = self._envelope(resource_uri(resource), action, body, selectors)
        response = self._send(envelope.encode())
        if self.recorder is not None:
            self.recorder(resource, action.rsplit("/", 1)[-1], response.content)
        return self._body(response)

    def _send(self, data: bytes) -> requests.Response:
        error: requests.RequestException | None = None
        for _ in range(self.retries + 1):
            try:
                return self.session.post(
                    self.url, data=data, headers=_HEADERS, timeout=self.timeout
                )
            except requests.RequestException as exc:
                error = exc
        raise WsmanError(f"{self.url}: {error}") from error

    def _body(self, response: requests.Response) -> ET.Element:
        """The SOAP body of a response; a fault, a refusal or anything else raises."""
        if response.status_code == 401:
            raise WsmanAuthError(f"{self.url}: credentials refused (HTTP 401)")
        # A DTD has no business in a WS-Management response; refusing one, in whatever
        # encoding, keeps entity expansion out of the parser.
        body: ET.Element | None
        try:
            body = fromstring(response.content, forbid_dtd=True).find(f"{{{SOAP_NS}}}Body")
        except (ParseError, DefusedXmlException):
            body = None
        if body is not None:
            fault = body.find(f"{{{SOAP_NS}}}Fault")
            if fault is not None:
                raise _fault(fault)
        if body is None or response.status_code >= 400:
            raise WsmanError(f"{self.url}: HTTP {response.status_code}, no WS-Management answer")
        return body