import network, time, hashlib, ujson as json
from umqtt.simple import MQTTClient

WIFI_SSID = "Wokwi-GUEST"
WIFI_PASSWORD = ""
MQTT_BROKER = "test.mosquitto.org"
MQTT_PORT = 1883
NODE_ID = "NODE_C"

TOPIC_FORWARD = b"iot/cncp/routing/forward"
TOPIC_ALERT = b"iot/cncp/security/alert"

hop_filter = set()

def sha256b(data):
    return hashlib.sha256(data).digest()

def to_hex(b):
    return "".join("{:02x}".format(x) for x in b)

def filtered_body(d, exact_exclude=(), prefix_exclude=()):
    out = {}
    for k in d:
        if k in exact_exclude:
            continue
        skip = False
        for p in prefix_exclude:
            if k.startswith(p):
                skip = True
                break
        if not skip:
            out[k] = d[k]
    return out

def canonical_json(d):
    parts = []
    for k in sorted(d.keys()):
        parts.append('"%s":%s' % (k, json.dumps(d[k])))
    return "{" + ",".join(parts) + "}"

class Crypto:
    def __init__(self, node_id):
        self.priv = sha256b((node_id + "_PRIVATE_SEED_BLOCKCHAIN").encode())
        x = sha256b(self.priv)
        self.pub = bytes([0x04]) + x + sha256b(self.priv + x)

    def sign(self, d, exact_exclude=(), prefix_exclude=()):
        body = canonical_json(filtered_body(d, exact_exclude, prefix_exclude)).encode()
        return to_hex(sha256b(sha256b(self.pub[1:33]) + body))

    def verify(self, d, sig, pub_hex, exact_exclude=(), prefix_exclude=()):
        try:
            pub = bytes(int(pub_hex[i:i+2], 16) for i in range(0, len(pub_hex), 2))
            vkey = sha256b(pub[1:33])
            body = canonical_json(filtered_body(d, exact_exclude, prefix_exclude)).encode()
            return to_hex(sha256b(vkey + body)) == sig
        except Exception:
            return False

    def pub_hex(self):
        return to_hex(self.pub)

def connect_wifi():
    w = network.WLAN(network.STA_IF)
    w.active(True)
    w.connect(WIFI_SSID, WIFI_PASSWORD)
    print("[NODE_C] Connecting WiFi...")
    for _ in range(20):
        if w.isconnected():
            print("[NODE_C] WiFi:", w.ifconfig()[0])
            return True
        time.sleep(0.5)
    return False

def connect_mqtt(client):
    while True:
        try:
            client.connect()
            client.subscribe(TOPIC_FORWARD)
            print("[NODE_C] MQTT connected")
            return
        except Exception as e:
            print("[NODE_C] MQTT retry:", e)
            time.sleep(2)

crypto = Crypto(NODE_ID)
client = None

def on_message(topic, payload):
    global client
    try:
        msg = json.loads(payload)
        rid = msg.get("route_id", "")

        if rid in hop_filter:
            return
        hop_filter.add(rid)

        if msg.get("relay_1_node") != "NODE_B":
            return

        print("[NODE_C] Processing", rid)

        ok_a = crypto.verify(
            msg,
            msg.get("signature", ""),
            msg.get("public_key", ""),
            exact_exclude=("signature", "public_key"),
            prefix_exclude=("relay_",)
        )

        if not ok_a:
            print("[NODE_C] NODE_A INVALID")
            alert = json.dumps({
                "attack_type": "MITM",
                "attacker_id": "UNKNOWN",
                "evidence": "node_a_sig_invalid",
                "timestamp": time.ticks_ms()
            })
            client.publish(TOPIC_ALERT, alert)
            return

        ok_b = crypto.verify(
            msg,
            msg.get("relay_1_sig", ""),
            msg.get("relay_1_pub", ""),
            exact_exclude=("relay_1_sig", "relay_1_pub")
        )

        if not ok_b:
            print("[NODE_C] NODE_B INVALID")
            alert = json.dumps({
                "attack_type": "MITM",
                "attacker_id": "NODE_B",
                "evidence": "node_b_sig_invalid",
                "timestamp": time.ticks_ms()
            })
            client.publish(TOPIC_ALERT, alert)
            return

        print("[NODE_C] Signatures VALID")

        msg["relay_2_node"] = NODE_ID
        msg["relay_2_hop"] = 2
        msg["relay_2_time"] = time.ticks_ms()
        msg["relay_2_sig"] = crypto.sign(msg, exact_exclude=("relay_2_sig", "relay_2_pub"))
        msg["relay_2_pub"] = crypto.pub_hex()

        client.publish(TOPIC_FORWARD, json.dumps(msg))
        print("[NODE_C] Forwarded:", rid)

    except Exception as e:
        print("[NODE_C] Error:", e)

if not connect_wifi():
    raise Exception("WiFi failed")

client = MQTTClient(NODE_ID + "_client", MQTT_BROKER, MQTT_PORT, keepalive=60)
client.set_callback(on_message)
connect_mqtt(client)
print("[NODE_C] Subscribed")

while True:
    try:
        client.check_msg()
    except Exception as e:
        print("[NODE_C] MQTT Error:", e)
        connect_mqtt(client)
    time.sleep(0.5)