node.py

#!/usr/bin/env python3
"""
node.py

Peer node implementation for the blockchain network:
handles incoming messages, mining, tracker interactions,
and graceful shutdown on LEAVE, with Merkle-proof support
and automatic fork resolution on re-join or missing-parent NEW_BLOCK.
"""
import argparse
import socket
import threading
import json
import time
import signal
import sys

from blockchain import Blockchain, BlockHeader
from tx         import VoteTx
from p2p        import send, broadcast

def handle_connection(conn, bc: Blockchain, peers: dict, node_id: str, port: int, tracker_addr: str):
    try:
        data = conn.recv(100_000).decode()
        if not data:
            return

        msg = json.loads(data)
        typ = msg.get("type") or msg.get("cmd")

        # --- 1) Tracker → updated peer list + longest-chain sync ----
        if typ == "PEER_LIST_UPDATE":
            peers.clear()
            for p in msg["peers"]:
                if p["peer_id"] != node_id:
                    peers[p["peer_id"]] = p["addr"]
            print(f"[{node_id}] Peers updated: {peers}")

            # Immediately sync to the longest chain among peers
            for peer_id, addr in peers.items():
                host, prt = addr.split(":")
                try:
                    s = socket.socket()
                    s.settimeout(2)
                    s.connect((host, int(prt)))
                    s.sendall(json.dumps({"type": "GET_CHAIN"}).encode())
                    raw = s.recv(100_000).decode()
                    s.close()

                    other_chain_json = json.loads(raw)
                    other_chain = []
                    for blk in other_chain_json:
                        hdr = BlockHeader(
                            index=blk["index"],
                            timestamp=blk["timestamp"],
                            prev_hash=bytes.fromhex(blk["prev_hash"]),
                            merkle_root=bytes.fromhex(blk["merkle_root"]),
                            difficulty=blk["difficulty"],
                            nonce=blk["nonce"]
                        )
                        txs = [
                            VoteTx(
                                voter_pubkey=bytes.fromhex(t["voter_pubkey"]),
                                choice_id=t["choice_id"],
                                timestamp=t["timestamp"],
                                signature=bytes.fromhex(t["signature"])
                            )
                            for t in blk["txs"]
                        ]
                        other_chain.append((hdr, txs))

                    print(f"[{node_id}] trying to replace from {peer_id}: new_chain_len={len(other_chain)}, current_len={len(bc.chain)}")
                    swapped = bc.replace_chain(other_chain)
                    print(f"[{node_id}] replace_chain returned {swapped} (new length = {len(bc.chain)})")
                    if swapped:
                        print(f"[{node_id}] 👉 Adopted longer chain from {peer_id}")
                        break
                except Exception:
                    continue
            return

        # --- 2) NEW_TX from CLI or peer ----
        if typ == "NEW_TX":
            txd = msg["tx"]
            tx = VoteTx(
                voter_pubkey=bytes.fromhex(txd["voter_pubkey"]),
                choice_id=txd["choice_id"],
                timestamp=txd["timestamp"],
                signature=bytes.fromhex(txd["signature"])
            )

            forwarded = msg.get("forwarded", False)
            accepted = bc.add_transaction(tx)

            # reply to client
            if not forwarded:
                if accepted:
                    conn.sendall(json.dumps({"status": "OK"}).encode())
                else:
                    conn.sendall(json.dumps({
                        "status": "ERROR",
                        "reason": "duplicate or invalid"
                    }).encode())

            # rebroadcast once
            if accepted and not forwarded:
                fwd = dict(msg)
                fwd["forwarded"] = True
                broadcast(peers, fwd)
            return

        # --- 3) NEW_BLOCK from peer ----
        if typ == "NEW_BLOCK":
            header = BlockHeader(**msg["header"])
            txs = [
                VoteTx(
                    voter_pubkey=bytes.fromhex(t["voter_pubkey"]),
                    choice_id=t["choice_id"],
                    timestamp=t["timestamp"],
                    signature=bytes.fromhex(t["signature"])
                )
                for t in msg["txs"]
            ]

            if bc.add_block(header, txs):
                print(f"[{node_id}] Added block #{header.index} with {len(txs)} tx(s)")
                with bc.lock:
                    confirmed = {VoteTx.txid(tx) for tx in txs}
                    bc.mempool = [p for p in bc.mempool if VoteTx.txid(p) not in confirmed]
            else:
                print(f"[{node_id}] ⛔ Rejected invalid/tampered block #{header.index}")
                # attempt fork resolution if it was due to missing parent
                print(f"[{node_id}] Couldn't add block #{header.index}; fetching full chains...")
                for peer_id, addr in peers.items():
                    host, prt = addr.split(":")
                    try:
                        s = socket.socket()
                        s.settimeout(2)
                        s.connect((host, int(prt)))
                        s.sendall(json.dumps({"type": "GET_CHAIN"}).encode())
                        raw = s.recv(100_000).decode()
                        s.close()

                        other_chain_json = json.loads(raw)
                        other_chain = []
                        for blk in other_chain_json:
                            hdr = BlockHeader(
                                index=blk["index"],
                                timestamp=blk["timestamp"],
                                prev_hash=bytes.fromhex(blk["prev_hash"]),
                                merkle_root=bytes.fromhex(blk["merkle_root"]),
                                difficulty=blk["difficulty"],
                                nonce=blk["nonce"]
                            )
                            txs_peer = [
                                VoteTx(
                                    voter_pubkey=bytes.fromhex(t["voter_pubkey"]),
                                    choice_id=t["choice_id"],
                                    timestamp=t["timestamp"],
                                    signature=bytes.fromhex(t["signature"])
                                )
                                for t in blk["txs"]
                            ]
                            other_chain.append((hdr, txs_peer))

                        swapped = bc.replace_chain(other_chain)
                        if swapped:
                            print(f"[{node_id}] ⚡️ Adopted chain from {peer_id} after missing-parent block")
                            break
                    except Exception:
                        continue
            return

        # --- 4) CLI → TALLY ----
        if typ == "TALLY":
            counts = bc.tally()
            conn.sendall(json.dumps(counts).encode())
            return

        # --- 5) CLI → VERIFY ----
        if typ == "VERIFY":
            txid = msg.get("txid")
            # confirmed?
            for header, txs in bc.chain:
                for tx in txs:
                    if VoteTx.txid(tx) == txid:
                        resp = {
                            "status": "confirmed",
                            "block": header.index,
                            "timestamp": header.timestamp
                        }
                        conn.sendall(json.dumps(resp).encode())
                        return
            # pending?
            for tx in bc.mempool:
                if VoteTx.txid(tx) == txid:
                    conn.sendall(json.dumps({"status": "pending"}).encode())
                    return
            conn.sendall(json.dumps({"error": "txid not found"}).encode())
            return

        # --- 6) Web UI → GET_CHAIN ----
        if typ == "GET_CHAIN":
            chain_data = []
            for header, txs in bc.chain:
                chain_data.append({
                    "index":       header.index,
                    "timestamp":   header.timestamp,
                    "prev_hash":   header.prev_hash.hex(),
                    "merkle_root": header.merkle_root.hex(),
                    "difficulty":  header.difficulty,
                    "nonce":       header.nonce,
                    "txs": [
                        {
                            "voter_pubkey": t.voter_pubkey.hex(),
                            "choice_id":    t.choice_id,
                            "timestamp":    t.timestamp,
                            "signature":    t.signature.hex()
                        }
                        for t in txs
                    ]
                })
            conn.sendall(json.dumps(chain_data).encode())
            return

        # --- 7) Web UI → GET_PEERS ----
        if typ == "GET_PEERS":
            peers_data = [{"id": node_id, "address": f"127.0.0.1:{port}"}]
            peers_data += [{"id": pid, "address": addr} for pid, addr in peers.items()]
            conn.sendall(json.dumps(peers_data).encode())
            return

        # --- 8) CLI → GET_PROOF (Merkle inclusion) ----
        if typ == "GET_PROOF":
            blk = msg.get("block")
            txid = msg.get("txid")
            try:
                result = bc.get_proof(blk, txid)
                _, txs = bc.chain[blk]
                tx_json = None
                for t in txs:
                    if VoteTx.txid(t) == txid:
                        tx_json = {
                            "voter_pubkey": t.voter_pubkey.hex(),
                            "choice_id":    t.choice_id,
                            "timestamp":    t.timestamp,
                            "signature":    t.signature.hex()
                        }
                        break
                if tx_json is None:
                    raise ValueError("txid not found in block")
                resp = {
                    "status": "OK",
                    "root":   result["root"],
                    "proof":  result["proof"],
                    "tx":     tx_json
                }
                conn.sendall(json.dumps(resp).encode())
            except Exception as e:
                conn.sendall(json.dumps({
                    "status": "ERROR",
                    "reason": str(e)
                }).encode())
            return

        # --- 9) Web UI → RESET ----
        if typ == "RESET":
            with bc.lock:
                # Clear mempool
                bc.mempool = []
                # Reset chain to genesis block
                genesis = BlockHeader(
                    index=0,
                    timestamp=time.time(),
                    prev_hash=b'\x00'*32,
                    merkle_root=b'\x00'*32,
                    difficulty=bc.difficulty,
                    nonce=0
                )
                ghash = genesis.hash()
                bc.blocks = {ghash: (genesis, [])}
                bc.parent = {ghash: None}
                bc.children = {ghash: []}
                bc.height = {ghash: 0}
                bc.best_tip = ghash
                bc.chain = [(genesis, [])]
            conn.sendall(json.dumps({"status": "OK"}).encode())
            return

    except Exception as e:
        print(f"[{node_id}] Error handling message: {e}")
    finally:
        conn.close()

def listen(node_id, port, bc, peers, tracker_addr):
    sock = socket.socket()
    sock.bind(("0.0.0.0", port))
    sock.listen()
    print(f"[{node_id}] Listening on port {port}")
    while True:
        conn, _ = sock.accept()
        threading.Thread(
            target=handle_connection,
            args=(conn, bc, peers, node_id, port, tracker_addr),
            daemon=True
        ).start()

def print_chain_periodically(bc, node_id):
    while True:
        print(f"\n[{node_id}] Current blockchain:")
        for header, txs in bc.chain:
            print(f"  Block #{header.index} @ {header.timestamp} - {len(txs)} tx(s)")
        print(f"  Mempool: {len(bc.mempool)} pending tx(s)\n")
        time.sleep(20)  # Adjust interval as needed

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--id",      required=True, help="Unique node ID")
    parser.add_argument("--port",    type=int, required=True, help="Listening port")
    parser.add_argument("--tracker", default="127.0.0.1:9000", help="Tracker address")
    args = parser.parse_args()

    node_id = args.id
    addr    = f"127.0.0.1:{args.port}"
    peers   = {}
    bc      = Blockchain()

    # Graceful shutdown handler
    def on_shutdown(sig, frame):
        send(args.tracker, {"cmd": "LEAVE", "peer_id": node_id, "addr": addr})
        sys.exit(0)
    signal.signal(signal.SIGINT,  on_shutdown)
    signal.signal(signal.SIGTERM, on_shutdown)

    # Start RPC listener before joining, so we receive updates
    threading.Thread(target=listen, args=(node_id, args.port, bc, peers, args.tracker), daemon=True).start()

    # Announce to tracker (will trigger PEER_LIST_UPDATE)
    send(args.tracker, {"cmd": "JOIN", "peer_id": node_id, "addr": addr})

    def mine_loop():
        min_txs = 3 # minimum number of transactions to mine for a block
        last_mine = time.time() # we also mine a new block if the last block was mined more than 10 seconds ago and there is no new transaction during this time
        max_wait = 45 # maximum time to wait for new transactions before mining a block

        while True:
            now = time.time()
            # get mempool size under lock
            pool_size = len(bc.mempool)

            # decide whether to mine
            should_mine = (
                    pool_size >= min_txs or
                    (pool_size > 0 and (now - last_mine) >= max_wait)
            )

            if should_mine:
                print(f"Mining condition satisfied, [{node_id}] Mining block with {pool_size} tx(s) in mempool, time since last mine: {now - last_mine:.2f}s")
                header, txs, elapsed = bc.mine_block()
                last_mine = time.time()

                # prepare JSON
                txs_json = [{
                    "voter_pubkey": t.voter_pubkey.hex(),
                    "choice_id": t.choice_id,
                    "timestamp": t.timestamp,
                    "signature": t.signature.hex()
                } for t in txs]
                msg = {
                    "type": "NEW_BLOCK",
                    "header": header.__dict__,
                    "txs": txs_json
                }

                broadcast(peers, msg)
                print(
                    f"[{header.index}] Mined {len(txs)} tx(s) in {elapsed:.2f}s; "
                    f"next diff = {bc.difficulty}"
                )

            time.sleep(0.5)

    threading.Thread(target=mine_loop, daemon=True).start()

    threading.Thread(
        target=print_chain_periodically,
        args=(bc, node_id),
        daemon=True
    ).start()
    print(f"[{node_id}] Node running on {addr}. Waiting for peers and txs…")
    while True:
        time.sleep(10)


if __name__ == "__main__":
    main()