Source code for tests.utils.crypto_extension_utils

#
# SPDX-FileCopyrightText: <text>Copyright 2026 Arm Limited and/or its
# affiliates <open-source-office@arm.com></text>
#
# SPDX-License-Identifier: MIT

from __future__ import annotations

import re
import time




class CryptographicExtensionUtils:
    """
    Utility functions for Arm Cryptographic Extension tests.
    Provides helper methods to:

     - Generate self-signed certificates
     - Start and stop SSL servers
     - Perform HTTPS downloads with timing
    """



    def _generate_certificate(self, mgr, port, days) -> None:
        """
        Generate a self-signed certificate using OpenSSL

        :param mgr: Manager object to execute commands
        :param port: Console port to run commands on
        :param days: Number of days the certificate is valid
        :returns: None
        """
        cert_cmd = (
            "openssl req -x509 -newkey rsa:2048 -keyout server-key.pem -out "
            "server-cert.pem -days {days} -nodes -subj "
            "'/C=AU/ST=Some-State/O=Internet Widgits Pty Ltd/CN=localhost'"
        ).format(days=days)
        mgr.run_cmd(
            port,
            cert_cmd,
            timeout=120,
        )




    def _start_ssl_server(self, mgr, port) -> None:
        """
        Start SSL server serving 100MB random data file

        :param mgr: Manager object to execute commands
        :param port: Console port to run commands on
        :returns: None
        """
        server_cmd = (
            "sh -c '(while true; do "
            "dd if=/dev/urandom bs=1M count=10 2>/dev/null | "
            "openssl s_server "
            "-cert server-cert.pem "
            "-key server-key.pem "
            "-accept 4433 "
            "-quiet "
            "-naccept 1; "
            "done) &'"
        )
        mgr.run_cmd(
            port,
            server_cmd,
            timeout=30,
        )
        # Give the server a moment to start
        time.sleep(2)




    def _cleanup_server(self, mgr, port) -> None:
        """
        Cleanup the server process by bringing it to
        foreground and terminating

        :param mgr: Manager object to execute commands
        :param port: Console port to run commands on
        :returns: None
        """
        # Try to bring the background job to foreground and terminate
        try:
            mgr.run_cmd(
                port,
                "fg",
                timeout=5,
            )
        except (AssertionError, RuntimeError):
            # If fg fails, try to kill the openssl processes directly
            try:
                mgr.run_cmd(
                    port,
                    "pkill -f 's_server'",
                    timeout=10,
                )
            except AssertionError as exc:
                raise AssertionError(
                    "Failed to clean up openssl server processes"
                ) from exc




    def _cleanup_files(self, mgr, port) -> None:
        """
        Cleanup generated certificate files

        :param mgr: Manager object to execute commands
        :param port: Console port to run commands on
        :returns: None
        """

        # Cleanup certificate and key files
        for cmd in ("rm -f server-key.pem", "rm -f server-cert.pem"):
            mgr.run_cmd(
                port,
                cmd,
                timeout=10,
            )




    def _extract_time_fields(self, output: str) -> dict:
        """
        Extract timing information from the 'time' command output.
        Returns a dictionary with 'real', 'user', and 'sys' timing values.

        :param output: The output string from the 'time' command
        :returns: A dictionary with keys 'real', 'user', 'sys' and their
                      corresponding timing values in seconds
        """

        # Helper function to parse time fields from the output
        def parse(label: str, pattern: str) -> float:
            match = re.search(pattern, output)
            if not match:
                raise AssertionError(
                    f"Could not extract {label} time from output:\n{output}"
                )
            mins = int(match.group(1)) if match.group(1) else 0
            secs = float(match.group(2))
            return mins * 60 + secs

        # Define regex patterns for real, user, and sys time fields
        patterns = {
            "real": r"real\s+(?:(\d+)m)?(\d+(?:\.\d+)?)s",
            "user": r"user\s+(?:(\d+)m)?(\d+(?:\.\d+)?)s",
            "sys": r"sys\s+(?:(\d+)m)?(\d+(?:\.\d+)?)s",
        }
        return {
            label: parse(label, regex) for label, regex in patterns.items()
        }




    def _download(
        self,
        mgr,
        port,
        enable_crypto: bool,
        timeout: int = 60,
    ) -> dict:
        """Download data from the SSL server using openssl s_client with
        specified cipher and return timing information.


        :param mgr: Manager object to execute commands
        :param port: Console port to run commands on
        :param enable_crypto: Flag to enable or disable
                            cryptographic operations
        :param timeout: Timeout for the download command
        :returns: A dictionary with keys 'real', 'user', 'sys' and their
                  corresponding timing values in seconds
        """

        cipher_cmd = (
            "time sh -c 'echo -e "
            '"GET / HTTP/1.1\\r\\nHost: localhost\\r\\n'
            'Connection: close\\r\\n\\r\\n" | '
            "{env}openssl s_client "
            "-connect localhost:4433 "
            "-cipher AES256-GCM-SHA384 "
            "-servername localhost "
            "-quiet > /dev/null'"
        )
        # If enable_crypto is False, set OPENSSL_armcap=0x0
        # to disable crypto extension
        env_prefix = "" if enable_crypto else "OPENSSL_armcap=0x0 "
        output = mgr.run_cmd(
            port,
            cipher_cmd.format(env=env_prefix),
            timeout=timeout,
        )
        return self._extract_time_fields(output)




__all__ = ["CryptographicExtensionUtils"]