# Практикум 8. Безпека IoT: шифрування, цілісність, OWASP
# Студент: Шейко М.П., група ІА-44, варіант N = 13
#   Закладена вразливість (13 mod 3 = 1): відсутність TLS -> переведено на порт 8883 з ssl=True (OWASP I2)
#   Алгоритм шифрування payload (13 mod 2 = 1): AES-256-CBC (ключ 32 байти)
#   Розмір HMAC-ключа (16 + 8 * (13 mod 3)): 24 байти

import network, time, machine, dht, uos, ubinascii, uhashlib, ujson
try:
    import ussl
except ImportError:
    import ssl as ussl
try:
    import ucryptolib
except ImportError:
    import cryptolib as ucryptolib
from machine import Pin
from umqtt.simple import MQTTClient

# ---------------- Крок 4. Винесення секретів у secrets.py ----------------
try:
    from secrets import WIFI_PASS, AES_KEY, HMAC_KEY
except ImportError:
    raise RuntimeError("Немає secrets.py — секрети не зашиваємо в код!")

assert len(AES_KEY) == 32, "Для варіанта 13 (AES-256-CBC) ключ AES має бути 32 байти"
assert len(HMAC_KEY) == 24, "Для варіанта 13 ключ HMAC має бути 24 байти"

# ---------------- Крок 1Б. Підключення до Wi-Fi ----------------
def wifi_connect(ssid="Wokwi-GUEST", password=WIFI_PASS):
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    if not wlan.isconnected():
        print("Wi-Fi: під'єднуюся до", ssid, "...")
        wlan.connect(ssid, password)
        for _ in range(40):                     # чекаємо до ~20 с
            if wlan.isconnected():
                break
            time.sleep(0.5)
    if not wlan.isconnected():
        raise RuntimeError("Wi-Fi не під'єднався")
    print("Wi-Fi OK, IP:", wlan.ifconfig()[0])
    return wlan

# ---------------- Крок 1В. TLS-брокер (порт 8883) ----------------
# Виправлено вразливість варіанта 13 (OWASP I2: Insecure Network Services / Lack of Transport Encryption):
# замість відкритого порту 1883 використовуємо захищений порт 8883 та ssl=True.
BROKER    = "test.mosquitto.org"
PORT      = 8883                                # TLS-порт замість 1883
CLIENT_ID = "kpi-secure-" + machine.unique_id().hex()

def mqtt_connect():
    client = MQTTClient(CLIENT_ID, BROKER, port=PORT,
                        keepalive=30, ssl=True, ssl_params={})
    client.connect()
    print("TLS-з'єднання встановлено (порт 8883)")
    return client

# ---------------- Крок 1Г. Читання датчика DHT22 ----------------
sensor = dht.DHT22(Pin(15))

def read_sensor():
    sensor.measure()                            # опитуємо датчик
    return sensor.temperature()                 # температура, °C

# ---------------- Крок 2. Симетричне шифрування payload (AES-256-CBC) ----------------
def pad(data):                                  # PKCS7-подібне вирівнювання до 16 байт
    n = 16 - (len(data) % 16)
    return data + bytes([n]) * n

def encrypt(plaintext):
    iv = uos.urandom(16)                        # НОВИЙ випадковий IV на КОЖНЕ повідомлення
    cipher = ucryptolib.aes(AES_KEY, 2, iv)     # 2 = режим CBC, len(AES_KEY)=32 -> AES-256
    ct = cipher.encrypt(pad(plaintext))
    return ubinascii.b2a_base64(iv + ct).strip()  # IV попереду, все в base64

# ---------------- Крок 3. Додати HMAC-SHA256 для цілісності (Encrypt-then-MAC) ----------------
def hmac_sha256(key, msg):
    block = 64
    if len(key) > block:
        key = uhashlib.sha256(key).digest()
    key = key + b"\x00" * (block - len(key))
    o_pad = bytes(b ^ 0x5c for b in key)
    i_pad = bytes(b ^ 0x36 for b in key)
    inner = uhashlib.sha256(i_pad + msg).digest()
    return uhashlib.sha256(o_pad + inner).digest()

def secure_message(plaintext):
    payload = encrypt(plaintext)                # base64(IV+CT) з кроку 2
    tag = ubinascii.hexlify(hmac_sha256(HMAC_KEY, payload))
    return payload + b"." + tag                 # "<payload>.<hmac>"

# ---------------- Збираємо все разом ----------------
TOPIC = b"iis/ia-44/node-13/secure"             # топік групи ia-44, вузол node-13

wlan   = wifi_connect()                         # 1) спершу мережа
client = mqtt_connect()                         # 2) потім TLS-брокер
print("Захищений вузол запущено ->", TOPIC)

while True:
    t = read_sensor()                           # 3) реальний DHT22
    body = ujson.dumps({"temp": t, "ts": time.time()}).encode()
    packet = secure_message(body)               # 4) "<base64(IV+CT)>.<hmac>"
    client.publish(TOPIC, packet)
    print("PUB", TOPIC, "(зашифровано + HMAC):", packet.decode())
    time.sleep(2)                               # DHT22: не частіше ніж раз на ~2 с