#!/usr/bin/env python3

import base64
import getpass
import json
import os
import ssl
import struct
import subprocess
import threading

import paho.mqtt.client as mqtt
from cryptography.hazmat.primitives.ciphers.aead import AESGCM


IMEI = "864011069739334"
BROKER = "127.0.0.1"
PORT = 8883

TOPIC_IN = f"leo/centrin/{IMEI}/s1"
TOPIC_OUT = f"leo/centrout/{IMEI}/s1"

APP_ID = "STS-TEST"

connected = threading.Event()
welcome_received = threading.Event()
response_received = threading.Event()

device_key = None
client_nonce_b64 = None
sid_hex = None


def get_device_key():
    php = r'''
require "/var/www/html/sts_secure/lib/db.php";
require "/var/www/html/sts_secure/lib/crypto.php";

$pdo = sts_db();

$s = $pdo->prepare(
    "SELECT device_key_ciphertext,
            device_key_nonce,
            device_key_tag
     FROM devices
     WHERE imei = :imei"
);

$s->execute([":imei" => "864011069739334"]);
$r = $s->fetch();

if (!$r) {
    fwrite(STDERR, "Device non trovato\n");
    exit(1);
}

$k = sts_decrypt_device_key_from_storage(
    $r["device_key_ciphertext"],
    $r["device_key_nonce"],
    $r["device_key_tag"]
);

echo bin2hex($k);
'''

    result = subprocess.check_output(
        ["php", "-r", php],
        text=True
    ).strip()

    key = bytes.fromhex(result)

    if len(key) != 32:
        raise RuntimeError("DeviceKey non valida")

    return key


def decrypt_s1o(payload):
    if not payload.startswith("S1O|"):
        return None

    parts = payload.split("|")

    if len(parts) != 4:
        return None

    nonce = base64.b64decode(parts[1])
    ciphertext = base64.b64decode(parts[2])
    tag = base64.b64decode(parts[3])

    aes = AESGCM(device_key)

    plain = aes.decrypt(
        nonce,
        ciphertext + tag,
        b"STS1:S2C"
    )

    return json.loads(
        plain.decode("utf-8")
    )


def send_secure_command(client, sid):
    seq = 1

    plain = json.dumps(
        {
            "Id": APP_ID,
            "Type": "GETEXCEPTCNT"
        },
        separators=(",", ":")
    ).encode("utf-8")

    sid_bytes = bytes.fromhex(sid)

    nonce = sid_bytes + struct.pack(
        ">I",
        seq
    )

    aes = AESGCM(device_key)

    encrypted = aes.encrypt(
        nonce,
        plain,
        b"STS1:C2S"
    )

    ciphertext = encrypted[:-16]
    tag = encrypted[-16:]

    envelope = (
        f"S1D|{sid}|{seq:08X}|"
        f"{base64.b64encode(ciphertext).decode()}|"
        f"{base64.b64encode(tag).decode()}"
    )

    print("TX comando cifrato GETEXCEPTCNT")

    client.publish(
        TOPIC_IN,
        envelope,
        qos=0
    )


def on_connect(client, userdata, flags, rc):
    if rc != 0:
        print(
            "Connessione MQTT fallita:",
            rc
        )
        return

    print("MQTT TLS connesso")
    print("SUB:", TOPIC_OUT)

    client.subscribe(
        TOPIC_OUT
    )

    connected.set()


def on_message(client, userdata, msg):
    global sid_hex

    text = msg.payload.decode(
        "utf-8",
        errors="replace"
    )

    print(
        "RX MQTT:",
        text[:80] +
        ("..." if len(text) > 80 else "")
    )

    try:
        data = decrypt_s1o(text)
    except Exception as e:
        print(
            "RX non decifrabile:",
            e
        )
        return

    if not data:
        return

    if data.get("Sec") == "WELCOME":


        if welcome_received.is_set():
            print("WELCOME duplicato ignorato")
            return


        if data.get("Nonce") != client_nonce_b64:
            print(
                "WELCOME con nonce errato"
            )
            return

        sid_hex = data.get("Sid")

        if not sid_hex:
            print(
                "WELCOME senza SID"
            )
            return

        print("WELCOME cifrato: OK")
        print("SID:", sid_hex)

        welcome_received.set()

        send_secure_command(
            client,
            sid_hex
        )

        return

    if data.get("Type") == "GETEXCEPTCNT":
        print(
            "Risposta applicativa cifrata: OK"
        )

        print(
            json.dumps(
                data,
                indent=2
            )
        )

        response_received.set()


def main():
    global device_key
    global client_nonce_b64

    print(
        "Recupero DeviceKey dal DB..."
    )

    device_key = get_device_key()

    print(
        "DeviceKey: OK (non visualizzata)"
    )

    password = getpass.getpass(
        "Password MQTT utente sts_secure_test: "
    )

    client_nonce = os.urandom(16)

    client_nonce_b64 = base64.b64encode(
        client_nonce
    ).decode()

    ctx = ssl.create_default_context()

    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE

    client = mqtt.Client(
        protocol=mqtt.MQTTv311
    )

    client.username_pw_set(
        "sts_secure_test",
        password
    )

    client.tls_set_context(ctx)

    client.on_connect = on_connect
    client.on_message = on_message

    print(
        "Connessione broker..."
    )

    client.connect(
        BROKER,
        PORT,
        30
    )

    client.loop_start()

    if not connected.wait(5):
        raise RuntimeError(
            "Timeout connessione MQTT"
        )

    hello = (
        f"S1H|{APP_ID}|"
        f"{client_nonce_b64}"
    )

    print(
        "TX HELLO sicuro"
    )

    client.publish(
        TOPIC_IN,
        hello,
        qos=0
    )

    if not welcome_received.wait(8):
        raise RuntimeError(
            "WELCOME sicuro non ricevuto"
        )

    if not response_received.wait(8):
        raise RuntimeError(
            "Risposta applicativa non ricevuta"
        )

    print()
    print(
        "=============================="
    )
    print(
        " MQTT SICURO END-TO-END: OK"
    )
    print(
        "=============================="
    )

    client.disconnect()
    client.loop_stop()


if __name__ == "__main__":
    main()