import paho.mqtt.client as mqtt
import threading
import os

# MQTT Credentials (Hardcoded)
MQTT_BROKER = "6328b5345e964f8aac70312eacd3ee65.s1.eu.hivemq.cloud"
MQTT_PORT = 8883  # Secure MQTT over TLS
MQTT_TOPIC = "sensorGain"
MQTT_USERNAME = "pipinstall"
MQTT_PASSWORD = "S1111s123"

def generate(type="AES"):
    """
    Generates a cryptographic key. Currently supports only AES key generation.

    Parameters:
        type (str): Type of key to generate (default is "AES").

    Returns:
        str: A 256-bit AES key in hexadecimal format.

    Raises:
        ValueError: If an unsupported type is requested.
        Exception: If no seed is received within the specified timeout.
    """

    if type.upper() != "AES":
        raise ValueError("Only 'AES' key generation is supported at this time.")

    seed = None
    event = threading.Event()

    def on_message(client, userdata, message):
        nonlocal seed
        if seed is None:
            seed = message.payload.decode()
            event.set()

    # Use MQTT v3.1.1 to avoid deprecation warnings
    client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2, protocol=mqtt.MQTTv311)
    client.username_pw_set(MQTT_USERNAME, MQTT_PASSWORD)  # Set credentials
    client.tls_set()  # Enable TLS encryption
    client.on_message = on_message

    try:
        client.connect(MQTT_BROKER, MQTT_PORT, 60)
        client.subscribe(MQTT_TOPIC)
        client.loop_start()

        # Wait for a seed
        event.wait(timeout=15)

        client.loop_stop()
        client.disconnect()

        if seed is None:
            raise Exception("No random seed available at the moment. Inquire developer.")

        # Generate a truly random 256-bit AES key
        random_key = os.urandom(32)  # 32 bytes = 256 bits

        # Mix seed entropy into the key (XOR combination)
        mixed_key = bytes(a ^ b for a, b in zip(random_key, seed.encode()[:32].ljust(32, b'\x00')))

        # Return as a hexadecimal string
        return mixed_key.hex()

    except Exception as e:
        client.loop_stop()
        raise Exception(f"Error in TRNG: {e}")

if __name__ == "__main__":
    try:
        key = generate()
        print("Generated AES Key:", key)
    except Exception as e:
        print("Error:", e)
