ping.py

"""Provide a mechanism to ping the server to see if everything is
working as expected.

ping.py allows you to query a self-signed websocket endpoint. Simply
provide the certificate file by pasting it into the `PUBLIC_KEY`
variable.
"""

import asyncio
import logging
import pathlib
import ssl
import tempfile
import time
from typing import Final

import websockets

# Set up logging
logging.basicConfig(
    format="%(asctime)-15s %(levelname)s :: %(filename)s:%(lineno)s:%(funcName)s() :: %(message)s",  # noqa: E501
    datefmt="%Y-%m-%d %H:%M:%S",
    level="INFO",
    handlers=[
        logging.StreamHandler(),
    ],
)

logging.Formatter.converter = time.gmtime


logger = logging.getLogger(__name__)

PING_URI: Final[str] = "wss://127.0.0.1:8001/ws/ping/"

PUBLIC_KEY: Final[
    str
] = """
-----BEGIN CERTIFICATE-----

-----END CERTIFICATE-----
"""


def write_certificate() -> str:
    """Output the certificate to a named temporary file.

    NB. Linux only.
    """
    # pylint: disable=R1732
    tmp_file = tempfile.NamedTemporaryFile("w", encoding="utf8", delete=False)
    tmp_file.write(PUBLIC_KEY.strip())
    tmp_file.close()
    return tmp_file.name


async def ping():
    """Ping the server."""
    cert_file = write_certificate()
    logging.info("Temporary cert file: %s", cert_file)
    ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    localhost_pem = pathlib.Path(cert_file)
    ssl_context.load_verify_locations(localhost_pem)
    validator_connection = f"{PING_URI}"
    # pylint: disable=E1101
    async with websockets.connect(validator_connection, ssl=ssl_context) as websocket:
        await websocket.send("PING")
        msg = await websocket.recv()
        logger.info(msg)


async def main():
    """Primary entry point of this script."""
    await ping()


if __name__ == "__main__":
    asyncio.run(main())