src/python_testing/TC_PAVSTI_Utils.py

#
#    Copyright (c) 2025 Project CHIP Authors
#    All rights reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.
#

import ipaddress
import logging
import os
import random
import shutil
import tempfile
from enum import StrEnum

import psutil
import requests
from cryptography import x509
from cryptography.hazmat.primitives import serialization
from TC_TLS_Utils import TLSUtils

from matter.interaction_model import Status
from matter.testing.tasks import Subprocess

log = logging.getLogger(__name__)


class SupportedIngestInterface(StrEnum):
    cmaf = "cmaf-ingest"  # Interface 1
    dash = "dash"  # Interface 2, DASH version
    hls = "hls"  # Interface 2, HLS version


class PushAvServerProcess(Subprocess):
    """Class for starting Push AV Server in a subprocess"""

    # Prefix for log messages from push av server
    PREFIX = b"[PUSH_AV_SERVER]"

    # By default this points to the push_av_server in Test Harness
    # TCs utilizing this should expect th_server_app_path otherwise
    DEFAULT_SERVER_PATH = "/root/apps/push_av_server/src/server.py"

    def __init__(
        self,
        server_path: str | None,
        port: int = 1234,
        host: str = "0.0.0.0",
        server_ip: str | None = None,
    ):
        if server_path is None:
            log.error("No path provided for Push AV Server, using the default path for TH: %s", self.DEFAULT_SERVER_PATH)
            server_path = self.DEFAULT_SERVER_PATH
        self._working_directory = os.path.join(tempfile.gettempdir(), "pavstest")
        if os.path.exists(self._working_directory):
            shutil.rmtree(self._working_directory)
        os.makedirs(self._working_directory)
        self.host = host
        self.port = port
        self.base_url = f"https://{self.host}:{self.port}"
        # Build the command list
        command = ["python3", server_path]
        command.extend(
            [
                "--host",
                str(self.host),
                "--working-directory",
                self._working_directory,
                "--strict-mode"
            ]
        )

        if server_ip:
            command.extend(
                [
                    "--server-ip",
                    server_ip
                ]
            )

        # Start the server application
        super().__init__(
            *command,
            output_cb=lambda line, is_stderr: self.PREFIX + line,
        )

    def __del__(self):
        try:
            if os.path.exists(self._working_directory):
                shutil.rmtree(self._working_directory)
        except Exception:
            pass

    def _get_json(self, endpoint: str) -> dict:
        url = f"{self.base_url}{endpoint}"
        response = requests.get(url, verify=False, timeout=5)
        response.raise_for_status()
        return response.json()

    def _post_json(self, endpoint: str, data: dict | None = None) -> dict:
        url = f"{self.base_url}{endpoint}"
        response = requests.post(url, json=data or {}, verify=False, timeout=5)
        response.raise_for_status()
        return response.json()

    def get_root_cert(self) -> bytes:
        """Retrieve the root certificate in DER format."""
        response = self._get_json("/certs/server/root.pem")
        root_cert_pem = response["cert"]["public_cert"].encode("utf-8")
        root_cert = x509.load_pem_x509_certificate(root_cert_pem)
        return root_cert.public_bytes(encoding=serialization.Encoding.DER)

    def get_device_certificate(self, device_name: str = "DUT") -> bytes:
        """Retrieve a device certificate in DER format."""
        response = self._get_json(f"/certs/device/{device_name}.pem")
        device_cert_pem = response["cert"]["public_cert"].encode("utf-8")
        device_cert = x509.load_pem_x509_certificate(device_cert_pem)
        return device_cert.public_bytes(encoding=serialization.Encoding.DER)

    def sign_csr(self, csr_der: bytes, device_name: str = "DUT") -> dict:
        """Submit a CSR for signing and return the server's response."""
        csr = x509.load_der_x509_csr(csr_der)
        csr_pem = csr.public_bytes(serialization.Encoding.PEM).decode("utf-8")
        return self._post_json(f"/certs/{device_name}/sign", {"csr": csr_pem})

    def create_stream(self, interface: SupportedIngestInterface) -> str:
        """Request the server to create a new stream."""
        response = self._post_json(f"/streams?interface={interface}")
        return response["id"]

    def update_expected_track_names(self, stream_id: str, expected_track_names: list[str]) -> None:
        """
        Request the server to set the expected track names for a stream.
        This is required to validate that uploaded track names match the expected list.
        """
        self._post_json(endpoint=f"/streams/{stream_id}/expectedTrackNames", data={"expected_track_names": expected_track_names})


class PAVSTIUtils:
    """Utils for Push AV TC's TLS requirements."""

    # ----------------------------------------------------------------------
    # Precondition setup
    # ----------------------------------------------------------------------

    def get_private_ip(self):
        candidates = {"192": [], "10": []}

        for iface, addrs in psutil.net_if_addrs().items():
            for addr in addrs:
                if addr.family.name == 'AF_INET':
                    ip = addr.address
                    ip_obj = ipaddress.ip_address(ip)
                    if ip_obj.is_private:
                        if ip.startswith("192.168."):
                            candidates["192"].append(ip)
                        elif ip.startswith("10."):
                            candidates["10"].append(ip)

        if candidates["192"]:
            return candidates["192"][0]
        if candidates["10"]:
            return candidates["10"][0]

        raise RuntimeError("No private IP found, specify using --string-arg host_ip <IPv4>")

    async def precondition_provision_tls_endpoint(
        self, server: PushAvServerProcess, host_ip: str | None = None
    ) -> int:
        """Perform provisioning steps to set up TLS endpoint."""
        if host_ip is None:
            # If no host ip specified, try to get private ip
            # this is mainly required when running TCs in Test Harness
            log.error("No host_ip provided in test arguments")
            host_ip = self.get_private_ip()

        # TLS clusters are always on EP0, per spec
        tls_utils = TLSUtils(self, endpoint=0)
        log.info("Using IP: %s as hostname to provision TLS Endpoint", host_ip)
        root_cert_der = server.get_root_cert()
        prc_result = await tls_utils.send_provision_root_command(certificate=root_cert_der)
        tls_utils.assert_valid_caid(prc_result.caid)

        csr_nonce = random.randbytes(32)
        csr_result = await tls_utils.send_csr_command(nonce=csr_nonce)
        tls_utils.assert_valid_ccdid(csr_result.ccdid)
        tls_utils.assert_valid_csr(csr_result, csr_nonce)

        server.sign_csr(csr_result.csr)
        device_cert_der = server.get_device_certificate()

        await tls_utils.send_provision_client_command(
            certificate=device_cert_der, ccdid=csr_result.ccdid
        )
        result = await tls_utils.send_provision_tls_endpoint_command(
            hostname=host_ip.encode('utf-8'),
            port=1234,
            expected_status=Status.Success,
            caid=prc_result.caid,
            ccdid=csr_result.ccdid,
        )
        return result.endpointID, host_ip

    async def postcondition_remove_tls_endpoint(self, tlsEndPoint):

        # Make sure the passed in values are not None.
        # This could happen in the unlikely event the test script fails prior to setting of any of these values (which would
        # be prior to the TLS EP being set).
        if tlsEndPoint is None:
            return

        tls_utils = TLSUtils(self, endpoint=0)
        await tls_utils.send_remove_tls_endpoint_command(
            endpoint_id=tlsEndPoint
        )