TLS-PSK 密钥认证

python.exe -c "import ssl; print(ssl.OPENSSL_VERSION); print('HAS_PSK =', ssl.HAS_PSK)"

HAS_PSK = True

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import base64
import binascii
import getpass
import secrets
import socket
import ssl
import string
import struct
import sys
import time


MIN_PYTHON_VERSION = (3, 13)
DEFAULT_PORT = 8883
DEFAULT_KEEPALIVE = 60
DEFAULT_TIMEOUT = 8
DEFAULT_EXPIRY_SECONDS = 3600
SDK_APP_ID = "12010126"
TLS_PSK_CIPHER = "PSK-AES128-CBC-SHA"


def check_runtime():
    """在读取设备凭证前检查TLS-PSK运行环境。"""
    print("正在检查运行环境。")
    print(f"Python版本：{sys.version.split()[0]}")
    print(f"Python路径：{sys.executable}")
    print(f"OpenSSL版本：{ssl.OPENSSL_VERSION}")

    if sys.version_info < MIN_PYTHON_VERSION:
        raise RuntimeError(
            "当前Python版本不支持TLS-PSK客户端接口。\\n"
            "请使用Python 3.13及以上版本运行本程序。\\n"
            "例如：D:\\\\tools\\\\python313\\\\python.exe "
            "mqtt_tls_psk_connect.py"
        )

    if not getattr(ssl, "HAS_PSK", False):
        raise RuntimeError(
            "当前Python所使用的OpenSSL未启用TLS-PSK能力。"
        )

    if not hasattr(ssl.SSLContext, "set_psk_client_callback"):
        raise RuntimeError(
            "当前ssl模块缺少set_psk_client_callback接口。"
        )

    print("TLS-PSK能力：支持")
    print("运行环境检查通过。\\n")


def required_input(prompt):
    """读取必填参数并清除首尾空格。"""
    value = input(prompt).strip()
    if not value:
        raise ValueError("输入内容不能为空")
    return value


def get_port():
    """读取并校验MQTT端口。"""
    value = input(
        f"MQTT服务器端口（直接回车使用{DEFAULT_PORT}）："
    ).strip()

    if not value:
        return DEFAULT_PORT

    try:
        port = int(value)
    except ValueError as exc:
        raise ValueError("端口号必须是整数") from exc

    if not 1 <= port <= 65535:
        raise ValueError("端口号必须在1～65535范围内")

    return port


def decode_device_secret(device_secret):
    """将Base64格式的DeviceSecret解码为TLS PSK。"""
    try:
        psk = base64.b64decode(
            device_secret,
            validate=True,
        )
    except (binascii.Error, ValueError) as exc:
        raise ValueError(
            "DeviceSecret不是有效的Base64字符串，"
            "请检查是否完整复制了控制台中的设备密钥"
        ) from exc

    if not psk:
        raise ValueError("DeviceSecret解码结果为空")

    return psk


def random_connid(length=5):
    """生成MQTT Username使用的随机connid。"""
    alphabet = string.ascii_letters + string.digits
    return "".join(
        secrets.choice(alphabet)
        for _ in range(length)
    )


def encode_utf8(value):
    """编码MQTT UTF-8字符串。"""
    data = value.encode("utf-8")

    if len(data) > 65535:
        raise ValueError("MQTT字符串长度超过65535字节")

    return struct.pack("!H", len(data)) + data


def encode_remaining_length(length):
    """编码MQTT Remaining Length字段。"""
    result = bytearray()

    while True:
        encoded = length % 128
        length //= 128

        if length > 0:
            encoded |= 0x80

        result.append(encoded)

        if length == 0:
            break

    return bytes(result)


def build_connect_packet(client_id, username, keepalive):
    """
    构造MQTT 3.1.1 CONNECT报文。

    TLS-PSK模式下：
    - MQTT Username仍用于标识设备；
    - MQTT Password不填写；
    - DeviceSecret在TLS握手层作为PSK使用。
    """
    connect_flags = 0x82
    # 0x80：包含Username
    # 0x02：Clean Session
    # 不设置Password标志位

    variable_header = (
        encode_utf8("MQTT")
        + bytes([4])
        + bytes([connect_flags])
        + struct.pack("!H", keepalive)
    )

    payload = (
        encode_utf8(client_id)
        + encode_utf8(username)
    )

    remaining_length = len(variable_header) + len(payload)

    return (
        bytes([0x10])
        + encode_remaining_length(remaining_length)
        + variable_header
        + payload
    )


def recv_exact(sock, length):
    """从Socket读取指定长度的数据。"""
    chunks = bytearray()

    while len(chunks) < length:
        data = sock.recv(length - len(chunks))

        if not data:
            raise ConnectionError("服务器已关闭连接")

        chunks.extend(data)

    return bytes(chunks)


def receive_mqtt_packet(sock):
    """接收一个完整MQTT报文。"""
    first_byte = recv_exact(sock, 1)[0]

    multiplier = 1
    remaining_length = 0

    while True:
        encoded = recv_exact(sock, 1)[0]
        remaining_length += (encoded & 0x7F) * multiplier

        if encoded & 0x80 == 0:
            break

        multiplier *= 128

        if multiplier > 128 * 128 * 128:
            raise ValueError("MQTT Remaining Length格式错误")

    payload = recv_exact(sock, remaining_length)
    return first_byte, payload


def parse_connack(first_byte, payload):
    """解析MQTT 3.1.1 CONNACK报文。"""
    if first_byte != 0x20 or len(payload) != 2:
        raise ValueError("服务器返回的不是有效CONNACK报文")

    return_code = payload[1]

    messages = {
        0: "连接成功",
        1: "不支持的协议版本",
        2: "ClientId不合法",
        3: "服务不可用",
        4: "Username或认证信息错误",
        5: "未授权",
    }

    return return_code, messages.get(
        return_code,
        f"未知返回码：{return_code}",
    )


def create_tls_psk_context(psk_identity, psk):
    """创建TLS-PSK客户端上下文。"""
    context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)

    # TLS-PSK路径不使用X.509服务器证书。
    context.check_hostname = False
    context.verify_mode = ssl.CERT_NONE

    # 当前已验证的平台TLS-PSK路径使用TLS 1.2。
    context.minimum_version = ssl.TLSVersion.TLSv1_2
    context.maximum_version = ssl.TLSVersion.TLSv1_2

    try:
        context.set_ciphers(TLS_PSK_CIPHER)
    except ssl.SSLError as exc:
        raise RuntimeError(
            f"当前OpenSSL不支持密码套件：{TLS_PSK_CIPHER}"
        ) from exc

    def psk_callback(identity_hint):
        return psk_identity, psk

    context.set_psk_client_callback(psk_callback)
    return context


def main():
    # 必须先检查运行环境，再读取设备信息和凭证。
    check_runtime()

    print("密钥认证设备TLS-PSK连接验证工具")
    print("本工具仅验证设备认证和上线。")
    print("输入的DeviceSecret仅保留在程序内存中，"
          "不会显示或写入文件。\\n")

    product_id = required_input("ProductID：")
    device_name = required_input("DeviceName：")

    default_broker = (
        f"{product_id}.iotcloud.tencentdevices.com"
    )

    broker = input(
        "MQTT服务器地址"
        f"（直接回车使用{default_broker}）："
    ).strip() or default_broker

    if any(char in broker for char in "${}/"):
        raise ValueError(
            "MQTT服务器地址应为不包含协议和路径的完整域名"
        )

    port = get_port()

    device_secret = getpass.getpass(
        "DeviceSecret（输入内容不会显示）："
    ).strip()

    if not device_secret:
        raise ValueError("DeviceSecret不能为空")

    psk = decode_device_secret(device_secret)

    client_id = f"{product_id}{device_name}"
    psk_identity = client_id
    connid = random_connid()
    expiry = int(time.time()) + DEFAULT_EXPIRY_SECONDS
    username = (
        f"{client_id};{SDK_APP_ID};{connid};{expiry}"
    )

    print("\\n连接配置：")
    print(f"MQTT服务器：{broker}:{port}")
    print(f"ClientId：{client_id}")
    print(f"PSK Identity：{psk_identity}")
    print(f"PSK长度：{len(psk)}字节")
    print(f"TLS密码套件：{TLS_PSK_CIPHER}")
    print("MQTT Username：已自动生成")
    print("MQTT Password：不使用")
    print("DeviceSecret不会输出。")

    confirm = input(
        "\\n是否开始TLS-PSK连接？[y/N]："
    ).strip().lower()

    if confirm != "y":
        print("已取消连接。")
        return 0

    context = create_tls_psk_context(
        psk_identity,
        psk,
    )

    raw_socket = None
    tls_socket = None

    try:
        print(f"\\n正在连接：{broker}:{port}")

        raw_socket = socket.create_connection(
            (broker, port),
            timeout=DEFAULT_TIMEOUT,
        )

        tls_socket = context.wrap_socket(
            raw_socket,
            server_hostname=broker,
        )

        print("TLS-PSK握手成功。")
        print(f"TLS版本：{tls_socket.version()}")

        cipher = tls_socket.cipher()
        if cipher:
            print(f"协商密码套件：{cipher[0]}")

        connect_packet = build_connect_packet(
            client_id,
            username,
            DEFAULT_KEEPALIVE,
        )

        tls_socket.sendall(connect_packet)

        first_byte, payload = receive_mqtt_packet(
            tls_socket
        )

        return_code, message = parse_connack(
            first_byte,
            payload,
        )

        if return_code != 0:
            raise RuntimeError(
                f"MQTT认证失败：{message}"
            )

        print("\\nTLS-PSK连接及MQTT认证成功。")
        print("设备已连接物联网开发平台。")
        print("请前往控制台查看设备在线状态。")
        print("程序将保持连接，按Ctrl+C退出。")

        # 每30秒发送一次MQTT PINGREQ，保持设备在线。
        tls_socket.settimeout(40)

        while True:
            time.sleep(30)
            tls_socket.sendall(b"\\xC0\\x00")

            packet_type, packet_payload = (
                receive_mqtt_packet(tls_socket)
            )

            if packet_type != 0xD0:
                raise RuntimeError(
                    "未收到预期的MQTT PINGRESP"
                )

    finally:
        if tls_socket is not None:
            try:
                # MQTT DISCONNECT
                tls_socket.sendall(b"\\xE0\\x00")
            except Exception:
                pass

            try:
                tls_socket.close()
            except Exception:
                pass

        elif raw_socket is not None:
            try:
                raw_socket.close()
            except Exception:
                pass

    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        print("\\n用户终止连接。")
        sys.exit(0)
    except socket.gaierror:
        print("\\n运行失败：无法解析MQTT服务器地址。")
        sys.exit(1)
    except socket.timeout:
        print("\\n运行失败：连接或响应超时。")
        sys.exit(1)
    except ssl.SSLError as exc:
        print(f"\\nTLS-PSK连接失败：{exc}")
        sys.exit(1)
    except Exception as exc:
        print(f"\\n运行失败：{exc}")
        sys.exit(1)

# Windows PowerShell
python mqtt_tls_psk_connect.py

# Linux && macOS Shell
python3 mqtt_tls_psk_connect.py

# 运行成功后依次输入
ProductID：
DeviceName：
MQTT服务器地址：
MQTT服务器端口（直接回车使用8883）：
DeviceSecret（输入内容不会显示）：

是否开始TLS-PSK连接？[y/N]：y

正在连接：******.iotcloud.tencentdevices.com:8883
TLS-PSK握手成功。
TLS版本：TLSv1.2
协商密码套件：

TLS-PSK连接及MQTT认证成功。
设备已连接物联网开发平台。
请前往控制台查看设备在线状态。
程序将保持连接，按Ctrl+C退出。