import dht
import machine
from time import sleep
import ucryptolib
import ubinascii
from ssd1306 import SSD1306_I2C

# Pin Definitions
DHT_PIN = 15
LED_PIN = 2
I2C_SDA_PIN = 21
I2C_SCL_PIN = 22

# Access Control Key
ACCESS_KEY = "secureKey123"

# AES Encryption Key and IV (16 bytes each)
AES_KEY = b"1234567890123456"
AES_IV = b"1234567890123456"

# Setup I2C and OLED Display
i2c = machine.I2C(0, scl=machine.Pin(I2C_SCL_PIN), sda=machine.Pin(I2C_SDA_PIN))
oled = SSD1306_I2C(128, 64, i2c)  # Ensure width=128 and height=64 are provided

# Setup DHT Sensor and LED
dht_sensor = dht.DHT22(machine.Pin(DHT_PIN))
led = machine.Pin(LED_PIN, machine.Pin.OUT)

def encrypt_data(data):
    """Encrypt the given data using AES."""
    aes = ucryptolib.aes(AES_KEY, 2, AES_IV)  # AES mode 2 = CBC
    padded_data = data + ' ' * (16 - len(data) % 16)  # Pad to multiple of 16 bytes
    encrypted = aes.encrypt(padded_data)
    return ubinascii.hexlify(encrypted).decode('utf-8')

def display_message(lines):
    """Display a message on the OLED screen."""
    oled.fill(0)  # Clear screen
    for i, line in enumerate(lines):
        oled.text(line, 0, i * 10)
    oled.show()

def send_data(key, temperature, humidity):
    """Send encrypted data if access key is correct."""
    if key != ACCESS_KEY:
        # Unauthorized Access
        led.on()
        display_message(["Unauthorized Access", "Alert Triggered!"])
        print("Unauthorized access attempt detected!")
        sleep(2)
        led.off()
        return

    # Authorized Access
    data = f"Temp:{temperature},Hum:{humidity}"
    encrypted_data = encrypt_data(data)
    display_message(["Data Sent:", f"Temp: {temperature}", f"Hum: {humidity}"])
    print("Encrypted Data:", encrypted_data)

def main():
    display_message(["System Booting...", "Initializing..."])
    sleep(2)

    while True:
        try:
            # Read from DHT Sensor
            dht_sensor.measure()
            temp = dht_sensor.temperature()
            hum = dht_sensor.humidity()

            if temp is None or hum is None:
                raise ValueError("Failed to read from DHT sensor")

            # Send data with valid key
            send_data(ACCESS_KEY, temp, hum)
            sleep(5)

            # Send data with invalid key
            send_data("wrongKey", temp, hum)
            sleep(5)

        except Exception as e:
            display_message(["Error:", str(e)])
            print("Error:", e)
            sleep(2)

if __name__ == "__main__":
    main()