sientia-dataops-opc-ingestor/e2e/opc_server.py

"""In-process OPC UA servers for the E2E suite (asyncua `Server`).

Each `OpcTestServer` listens on its own localhost port and exposes string NodeIds
under namespace 2 (`ns=2;s=<name>`), so a slot can address them the same way the
Orchestrator addresses a real server. Writes received from the gateway are recorded
(decoded `WriteValue`s) through the server's PreWrite callback.
"""

import socket
import time
from pathlib import Path

from asyncua import Server, ua
from asyncua.common.callback import CallbackType
from asyncua.crypto.cert_gen import setup_self_signed_certificate
from cryptography.x509.oid import ExtendedKeyUsageOID

NAMESPACE_URI = 'urn:sientia:e2e'

_VARIANT_TYPES = {
    'float': ua.VariantType.Float,
    'double': ua.VariantType.Double,
    'int': ua.VariantType.Int32,
    'bool': ua.VariantType.Boolean,
    'string': ua.VariantType.String,
}


def free_port() -> int:
    with socket.socket() as s:
        s.bind(('127.0.0.1', 0))
        return s.getsockname()[1]


def address(name: str) -> str:
    """NodeId string of a variable created by `OpcTestServer.add_variable`."""
    return f'ns=2;s={name}'


async def make_certificate(
    folder: Path, name: str, app_uri: str, server: bool
) -> tuple[Path, Path]:
    """Self-signed application certificate (DER) + private key (PEM) with `app_uri` as
    the SAN URI, as asyncua expects from both ends of a secure channel."""
    folder.mkdir(parents=True, exist_ok=True)
    cert, key = folder / f'{name}.der', folder / f'{name}.pem'
    usage = ExtendedKeyUsageOID.SERVER_AUTH if server else ExtendedKeyUsageOID.CLIENT_AUTH
    await setup_self_signed_certificate(
        key, cert, app_uri, 'localhost', [usage], {'commonName': name}
    )
    return cert, key


class OpcTestServer:
    """One asyncua server. `secured=True` exposes only Basic256Sha256 SignAndEncrypt."""

    def __init__(self, name: str, cert_dir: Path | None = None, secured: bool = False):
        self.name = name
        self.port = free_port()
        self.secured = secured
        self.cert_dir = cert_dir
        self.server: Server | None = None
        self.variables: dict[str, tuple[object, str, bool]] = {}
        self.writes: list[ua.WriteValue] = []
        self.write_times: list[float] = []  # time.monotonic() of each received write

    @property
    def url(self) -> str:
        return f'opc.tcp://127.0.0.1:{self.port}'

    def add_variable(
        self, name: str, value: object, data_type: str = 'float', writable: bool = True
    ) -> str:
        """Declare a variable (created on `start`, and again on every restart)."""
        self.variables[name] = (value, data_type, writable)
        return address(name)

    async def start(self) -> None:
        server = Server()
        await server.init()
        server.set_endpoint(self.url)
        server.set_server_name(self.name)
        if self.secured:
            assert self.cert_dir is not None, 'a secured server needs a cert_dir'
            cert, key = await make_certificate(
                self.cert_dir, f'server-{self.name}', f'urn:sientia:e2e:{self.name}', True
            )
            await server.load_certificate(str(cert))
            await server.load_private_key(str(key))
            server.set_security_policy([ua.SecurityPolicyType.Basic256Sha256_SignAndEncrypt])
        idx = await server.register_namespace(NAMESPACE_URI)
        assert idx == 2, f'unexpected namespace index {idx}'
        for name, (value, data_type, writable) in self.variables.items():
            node = await server.nodes.objects.add_variable(
                address(name), name, value, _VARIANT_TYPES[data_type]
            )
            if writable:
                await node.set_writable()
        server.subscribe_server_callback(CallbackType.PreWrite, self._record_write)
        await server.start()
        self.server = server

    async def stop(self) -> None:
        if self.server is not None:
            await self.server.stop()
            self.server = None

    async def restart(self) -> None:
        """Stop and start again on the same port (values reset to their initial ones)."""
        await self.stop()
        await self.start()

    async def set(self, name: str, value: object) -> None:
        """Change a value server-side (raises a data change for subscribed clients)."""
        assert self.server is not None
        _, data_type, _ = self.variables[name]
        node = self.server.get_node(address(name))
        await node.write_value(ua.Variant(value, _VARIANT_TYPES[data_type]))

    async def read(self, name: str) -> ua.DataValue:
        assert self.server is not None
        return await self.server.get_node(address(name)).read_data_value()

    def session_count(self) -> int:
        """Client sessions currently open on this server."""
        assert self.server is not None
        return len(self.server.iserver._external_sessions)

    async def _record_write(self, event, _dispatcher) -> None:
        if not event.is_external:  # the server updates its own nodes (ServerStatus, ...)
            return
        for value in event.request_params.NodesToWrite:
            self.writes.append(value)
            self.write_times.append(time.monotonic())