#!/usr/bin/python3
"""agent.ping over TLS with a pre-shared key, for data/probe.

Reads identity, key file and port from the environment (PSK_IDENTITY,
PSK_FILE, PSK_PORT). The key is read from the file here and never printed.
Exit 0 when the agent answers "1", otherwise one line with the reason and
exit 1. Needs python3 >= 3.13 (ssl.SSLContext.set_psk_client_callback).
"""
import os
import socket
import ssl
import struct
import sys


def main():
    if not hasattr(ssl.SSLContext, "set_psk_client_callback"):
        print("not ready (no PSK client: zabbix_get or python3 >= 3.13)")
        return 1
    identity = os.environ["PSK_IDENTITY"]
    port = int(os.environ.get("PSK_PORT") or 10050)
    try:
        with open(os.environ["PSK_FILE"]) as f:
            key = bytes.fromhex(f.read().strip())
    except (OSError, ValueError):
        print("PSK file is not readable or not hex")
        return 1

    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    # TLS 1.2 PSK suites are not in OpenSSL's default list; TLS 1.3 uses the
    # callback without a cipher change.
    ctx.set_ciphers("PSK")
    ctx.set_psk_client_callback(lambda hint: (identity, key))

    request = b"agent.ping"
    try:
        with socket.create_connection(("127.0.0.1", port), timeout=3) as raw:
            with ctx.wrap_socket(raw) as s:
                s.sendall(b"ZBXD\x01" + struct.pack("<Q", len(request)) + request)
                data = b""
                while True:
                    chunk = s.recv(4096)
                    if not chunk:
                        break
                    data += chunk
    except ConnectionRefusedError:
        print(f"port {port} not reachable")
        return 1
    except (OSError, ssl.SSLError) as e:
        print(f"agent.ping via TLS-PSK failed: {type(e).__name__}: {getattr(e, 'reason', None) or e}")
        return 1

    answer = data[13:].decode(errors="replace")
    if answer == "1":
        return 0
    print(f"agent.ping via TLS-PSK answers '{answer}' instead of 1 (127.0.0.1 in Server=?)")
    return 1


if __name__ == "__main__":
    sys.exit(main())
