napalm-amt/tests/conftest.py

"""Shared helpers: an Intel AMT WS-Management service served from the recorded responses."""

import re
import xml.etree.ElementTree as ET
from pathlib import Path

import pytest
import responses

FIXTURES = Path(__file__).parent / "fixtures"
DMT = FIXTURES / "dmt-wsman-messages"

HOST = "amt.example.net"
URL = f"http://{HOST}:16992/wsman"

_SOAP = "http://www.w3.org/2003/05/soap-envelope"
_ADDRESSING = "http://schemas.xmlsoap.org/ws/2004/08/addressing"
_WSMAN = "http://schemas.dmtf.org/wbem/wsman/1/wsman.xsd"
_ENUMERATION = "http://schemas.xmlsoap.org/ws/2004/09/enumeration"


def load(name: str, fixtures: Path = DMT) -> str:
    return (fixtures / name).read_text()


def _end_of_sequence(resource_uri: str) -> str:
    """The last, empty page of an enumeration."""
    return (
        f'<?xml version="1.0" encoding="UTF-8"?>'
        f'<a:Envelope xmlns:a="{_SOAP}" xmlns:b="{_ADDRESSING}" xmlns:c="{_WSMAN}" '
        f'xmlns:g="{_ENUMERATION}"><a:Header>'
        f"<b:Action>{_ENUMERATION}/PullResponse</b:Action>"
        f"<c:ResourceURI>{resource_uri}</c:ResourceURI></a:Header>"
        f"<a:Body><g:PullResponse><g:Items/><g:EndOfSequence/></g:PullResponse></a:Body>"
        f"</a:Envelope>"
    )


def _enumerate_response(resource_uri: str, context: str) -> str:
    return (
        f'<?xml version="1.0" encoding="UTF-8"?>'
        f'<a:Envelope xmlns:a="{_SOAP}" xmlns:b="{_ADDRESSING}" xmlns:c="{_WSMAN}" '
        f'xmlns:g="{_ENUMERATION}"><a:Header>'
        f"<b:Action>{_ENUMERATION}/EnumerateResponse</b:Action>"
        f"<c:ResourceURI>{resource_uri}</c:ResourceURI></a:Header>"
        f"<a:Body><g:EnumerateResponse><g:EnumerationContext>{context}</g:EnumerationContext>"
        f"</g:EnumerateResponse></a:Body></a:Envelope>"
    )


class Request:
    """One WS-Management request as the mock received it."""

    def __init__(self, body: bytes | str):
        self.xml = body.decode() if isinstance(body, bytes) else body
        root = ET.fromstring(self.xml)
        self.action = root.findtext(f"{{{_SOAP}}}Header/{{{_ADDRESSING}}}Action") or ""
        self.resource_uri = root.findtext(f"{{{_SOAP}}}Header/{{{_WSMAN}}}ResourceURI") or ""
        self.resource = self.resource_uri.rsplit("/", 1)[-1]
        self.operation = self.action.rsplit("/", 1)[-1]
        self.context = root.findtext(f".//{{{_ENUMERATION}}}EnumerationContext")
        self.body = root.find(f"{{{_SOAP}}}Body")
        header = root.find(f"{{{_SOAP}}}Header")
        self.selectors = {s.get("Name"): s.text for s in header.iter(f"{{{_WSMAN}}}Selector")}


class AmtServer:
    """Answer WS-Management requests from ``<Resource>.<operation>.xml`` files.

    ``Get`` -> ``<Resource>.get.xml``, ``Pull`` -> ``<Resource>.pull.xml`` (a page without
    ``EndOfSequence`` is followed by an empty last page), a method -> ``<Resource>.<Method>.xml``.
    A resource without a file is answered with the fault for an unknown resource URI, as
    AMT does. ``override(resource, operation, status, body)`` replaces one answer.
    """

    def __init__(self, rsps, fixtures: Path = DMT, url: str = URL):
        self.fixtures = fixtures
        self.requests: list[Request] = []
        self.overrides: dict[tuple[str, str], tuple[int, str]] = {}
        self._finished: set[str] = set()
        rsps.add_callback(responses.POST, re.compile(re.escape(url)), callback=self._handle)

    def override(self, resource: str, operation: str, status: int, body: str) -> None:
        self.overrides[(resource, operation)] = (status, body)

    def calls(self, operation: str | None = None) -> list[Request]:
        return [r for r in self.requests if operation is None or r.operation == operation]

    def _file(self, name: str) -> tuple[int, str]:
        path = self.fixtures / name
        if not path.exists():
            return 400, load("fault_destination_unreachable.xml")
        return 200, path.read_text()

    def _handle(self, request):
        req = Request(request.body)
        self.requests.append(req)
        status, text = self._answer(req)
        return status, {"Content-Type": "application/soap+xml;charset=UTF-8"}, text

    def _answer(self, req: Request) -> tuple[int, str]:
        if (req.resource, req.operation) in self.overrides:
            return self.overrides[(req.resource, req.operation)]
        if req.operation == "Get":
            return self._file(f"{req.resource}.get.xml")
        if req.operation == "Enumerate":
            if not (self.fixtures / f"{req.resource}.pull.xml").exists():
                return self._file(f"{req.resource}.pull.xml")
            return 200, _enumerate_response(req.resource_uri, f"ctx-{req.resource}")
        if req.operation == "Pull":
            if req.resource in self._finished:
                return 200, _end_of_sequence(req.resource_uri)
            status, text = self._file(f"{req.resource}.pull.xml")
            if "EndOfSequence" not in text:
                self._finished.add(req.resource)
            return status, text
        return self._file(f"{req.resource}.{req.operation}.xml")


@pytest.fixture
def amt():
    """A mocked AMT on http://amt.example.net:16992 answering from the DMT recordings."""
    with responses.RequestsMock(assert_all_requests_are_fired=False) as rsps:
        yield AmtServer(rsps)


def instances(name: str, fixtures: Path = DMT) -> list[dict]:
    """The instances in a recorded Get or Pull response, as the transport returns them."""
    from napalm_amt.wsman import to_instance

    body = ET.fromstring(load(name, fixtures)).find(f"{{{_SOAP}}}Body")
    page = body.find(f"{{{_ENUMERATION}}}PullResponse")
    if page is not None:
        return [to_instance(item) for item in page.findall(f"{{{_ENUMERATION}}}Items/*")]
    return [to_instance(body[0])]


def patched(name: str, old: str, new: str, count: int = -1) -> str:
    """A recorded response with ``old`` replaced; fails when ``old`` is not in it."""
    text = load(name)
    assert old in text, f"{old!r} not in {name}"
    return text.replace(old, new, count)