CIA/e-voting-system/tests/test_blockchain.py

#!/usr/bin/env python3
"""
Standalone Test Suite for PoA Blockchain Implementation

Tests the bootnode and validator services without requiring Docker or pytest.
Verifies core logic, consensus, and data structures.
"""

import sys
import json
import hashlib
import time
import importlib.util
from pathlib import Path

# ============================================================================
# Helper: Load module from file
# ============================================================================

def load_module(name, path):
    """Load a Python module from a file path"""
    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module

# ============================================================================
# Load bootnode and validator modules
# ============================================================================

project_root = Path(__file__).parent
bootnode_path = project_root / "bootnode" / "bootnode.py"
validator_path = project_root / "validator" / "validator.py"

bootnode_mod = load_module("bootnode", bootnode_path)
validator_mod = load_module("validator", validator_path)

PeerRegistry = bootnode_mod.PeerRegistry
PeerInfo = bootnode_mod.PeerInfo

Blockchain = validator_mod.Blockchain
Block = validator_mod.Block
Transaction = validator_mod.Transaction
PoAValidator = validator_mod.PoAValidator

# ============================================================================
# BOOTNODE TESTS
# ============================================================================

def test_bootnode_initialization():
    """Test bootnode peer registry initialization"""
    registry = PeerRegistry(peer_timeout_seconds=300)
    assert registry.peers == {}
    assert registry.peer_timeout == 300
    return True

def test_peer_registration():
    """Test registering a peer with bootnode"""
    registry = PeerRegistry()

    peer = PeerInfo(
        node_id="validator-1",
        ip="validator-1",
        p2p_port=30303,
        rpc_port=8001,
        public_key="0x1234567890abcdef"
    )

    registry.register_peer(peer)

    assert "validator-1" in registry.peers
    assert registry.peers["validator-1"]["info"].node_id == "validator-1"
    return True

def test_peer_discovery():
    """Test discovering peers from bootnode"""
    registry = PeerRegistry()

    # Register 3 peers
    for i in range(1, 4):
        peer = PeerInfo(
            node_id=f"validator-{i}",
            ip=f"validator-{i}",
            p2p_port=30300 + i,
            rpc_port=8000 + i,
            public_key=f"0x{'0' * 40}"
        )
        registry.register_peer(peer)

    # Discover peers (exclude validator-1)
    discovered = registry.get_peers_except("validator-1")

    assert len(discovered) == 2
    node_ids = [p.node_id for p in discovered]
    assert "validator-2" in node_ids
    assert "validator-3" in node_ids
    assert "validator-1" not in node_ids
    return True

def test_peer_heartbeat():
    """Test updating peer heartbeat"""
    registry = PeerRegistry(peer_timeout_seconds=5)

    peer = PeerInfo(
        node_id="validator-1",
        ip="validator-1",
        p2p_port=30303,
        rpc_port=8001
    )

    registry.register_peer(peer)
    old_heartbeat = registry.peers["validator-1"]["last_heartbeat"]

    # Update heartbeat
    time.sleep(0.05)
    registry.update_heartbeat("validator-1")
    new_heartbeat = registry.peers["validator-1"]["last_heartbeat"]

    assert new_heartbeat >= old_heartbeat
    return True

def test_stale_peer_cleanup():
    """Test cleanup of stale peers"""
    registry = PeerRegistry(peer_timeout_seconds=1)

    # Register peer
    peer = PeerInfo(
        node_id="validator-1",
        ip="validator-1",
        p2p_port=30303,
        rpc_port=8001
    )
    registry.register_peer(peer)
    assert len(registry.get_all_peers()) == 1

    # Manually set old heartbeat time
    registry.peers["validator-1"]["last_heartbeat"] = time.time() - 2

    # Cleanup stale peers
    removed_count = registry.cleanup_stale_peers()

    assert removed_count == 1
    assert len(registry.get_all_peers()) == 0
    return True


# ============================================================================
# VALIDATOR/BLOCKCHAIN TESTS
# ============================================================================

def test_genesis_block_creation():
    """Test genesis block is created correctly"""
    blockchain = Blockchain()

    assert blockchain.get_chain_length() == 1
    genesis = blockchain.get_block(0)
    assert genesis.index == 0
    assert genesis.validator == "genesis"
    assert genesis.transactions == []
    return True

def test_block_hash_calculation():
    """Test block hash is calculated correctly"""
    blockchain = Blockchain()

    # Create a block
    tx = Transaction(
        voter_id="test-voter",
        election_id=1,
        encrypted_vote="0x1234",
        ballot_hash="0xabcd",
        timestamp=1699360000
    )

    block = Block(
        index=1,
        prev_hash=blockchain.get_latest_block().block_hash,
        timestamp=1699360010,
        transactions=[tx],
        validator="validator-1"
    )

    # Calculate hash
    hash1 = Blockchain.calculate_block_hash(block)

    # Same block should produce same hash
    hash2 = Blockchain.calculate_block_hash(block)

    assert hash1 == hash2
    assert hash1.startswith("0x")
    assert len(hash1) == 66  # "0x" + 64 hex chars
    return True

def test_block_validation():
    """Test block validation logic"""
    blockchain = Blockchain()

    tx = Transaction(
        voter_id="test-voter",
        election_id=1,
        encrypted_vote="0x1234",
        ballot_hash="0xabcd",
        timestamp=1699360000
    )

    # Create valid block
    prev_block = blockchain.get_latest_block()
    block = Block(
        index=1,
        prev_hash=prev_block.block_hash,
        timestamp=1699360010,
        transactions=[tx],
        validator="validator-1"
    )
    block.block_hash = Blockchain.calculate_block_hash(block)

    assert blockchain.validate_block(block) == True

    # Test invalid block (wrong index)
    invalid_block = Block(
        index=99,
        prev_hash=prev_block.block_hash,
        timestamp=1699360010,
        transactions=[tx],
        validator="validator-1"
    )
    invalid_block.block_hash = Blockchain.calculate_block_hash(invalid_block)

    assert blockchain.validate_block(invalid_block) == False

    # Test invalid block (wrong prev_hash)
    invalid_block2 = Block(
        index=1,
        prev_hash="0x" + "0" * 64,
        timestamp=1699360010,
        transactions=[tx],
        validator="validator-1"
    )
    invalid_block2.block_hash = Blockchain.calculate_block_hash(invalid_block2)

    assert blockchain.validate_block(invalid_block2) == False

    # Test invalid block (unauthorized validator)
    invalid_block3 = Block(
        index=1,
        prev_hash=prev_block.block_hash,
        timestamp=1699360010,
        transactions=[tx],
        validator="unauthorized-node"
    )
    invalid_block3.block_hash = Blockchain.calculate_block_hash(invalid_block3)

    assert blockchain.validate_block(invalid_block3) == False
    return True

def test_add_block_to_chain():
    """Test adding valid blocks to blockchain"""
    blockchain = Blockchain()

    tx = Transaction(
        voter_id="test-voter",
        election_id=1,
        encrypted_vote="0x1234",
        ballot_hash="0xabcd",
        timestamp=1699360000
    )

    prev_block = blockchain.get_latest_block()
    block = Block(
        index=1,
        prev_hash=prev_block.block_hash,
        timestamp=1699360010,
        transactions=[tx],
        validator="validator-1"
    )
    block.block_hash = Blockchain.calculate_block_hash(block)

    result = blockchain.add_block(block)

    assert result == True
    assert blockchain.get_chain_length() == 2
    assert blockchain.get_block(1).index == 1
    return True

def test_blockchain_integrity_verification():
    """Test blockchain integrity verification"""
    blockchain = Blockchain()

    # Add valid blocks
    for i in range(1, 4):
        prev_block = blockchain.get_latest_block()
        tx = Transaction(
            voter_id=f"voter-{i}",
            election_id=1,
            encrypted_vote=f"0x{i:04x}",
            ballot_hash=f"0x{i:04x}",
            timestamp=1699360000 + i
        )

        block = Block(
            index=i,
            prev_hash=prev_block.block_hash,
            timestamp=1699360010 + i,
            transactions=[tx],
            validator=f"validator-{(i % 3) + 1}"
        )
        block.block_hash = Blockchain.calculate_block_hash(block)
        blockchain.add_block(block)

    # Verify integrity
    assert blockchain.verify_integrity() == True
    return True

def test_chain_immutability():
    """Test that modifying past blocks breaks chain"""
    blockchain = Blockchain()

    # Add a block
    prev_block = blockchain.get_latest_block()
    tx = Transaction(
        voter_id="voter-1",
        election_id=1,
        encrypted_vote="0x1234",
        ballot_hash="0xabcd",
        timestamp=1699360000
    )

    block = Block(
        index=1,
        prev_hash=prev_block.block_hash,
        timestamp=1699360010,
        transactions=[tx],
        validator="validator-1"
    )
    block.block_hash = Blockchain.calculate_block_hash(block)
    blockchain.add_block(block)

    # Verify chain is valid
    assert blockchain.verify_integrity() == True

    # Try to modify the transaction in block 1
    blockchain.chain[1].transactions[0].encrypted_vote = "0x9999"

    # Chain should now be invalid (hash mismatch)
    assert blockchain.verify_integrity() == False
    return True


# ============================================================================
# POA CONSENSUS TESTS
# ============================================================================

def test_round_robin_block_creation():
    """Test round-robin block creation eligibility"""
    # Create 3 validators
    validators = [
        PoAValidator("validator-1", "0x1234", "http://bootnode:8546", 8001, 30303),
        PoAValidator("validator-2", "0x5678", "http://bootnode:8546", 8002, 30304),
        PoAValidator("validator-3", "0xabcd", "http://bootnode:8546", 8003, 30305),
    ]

    # Each validator starts with blockchain of length 1 (just genesis)
    # next_block_index = blockchain.get_chain_length() = 1
    # Round-robin: next_block_index % num_validators == validator_index
    # 1 % 3 = 1, so validator at index 1 (validator-2) should create
    validator1 = validators[0]  # index 0
    validator2 = validators[1]  # index 1
    validator3 = validators[2]  # index 2

    should_create_1 = validator1.should_create_block()  # 1 % 3 == 0? No (1 % 3 = 1)
    should_create_2 = validator2.should_create_block()  # 1 % 3 == 1? Yes
    should_create_3 = validator3.should_create_block()  # 1 % 3 == 2? No

    assert should_create_1 == False, f"validator-1 should NOT create block 1, but got {should_create_1}"
    assert should_create_2 == True, f"validator-2 SHOULD create block 1, but got {should_create_2}"
    assert should_create_3 == False, f"validator-3 should NOT create block 1, but got {should_create_3}"

    return True

def test_authorized_validators():
    """Test that only authorized validators can create blocks"""
    blockchain = Blockchain()

    # Test authorized validators
    authorized = Blockchain.AUTHORIZED_VALIDATORS
    assert "validator-1" in authorized
    assert "validator-2" in authorized
    assert "validator-3" in authorized

    # Create block with authorized validator
    tx = Transaction(
        voter_id="voter-1",
        election_id=1,
        encrypted_vote="0x1234",
        ballot_hash="0xabcd",
        timestamp=1699360000
    )

    prev_block = blockchain.get_latest_block()
    block = Block(
        index=1,
        prev_hash=prev_block.block_hash,
        timestamp=1699360010,
        transactions=[tx],
        validator="validator-1"
    )
    block.block_hash = Blockchain.calculate_block_hash(block)

    assert blockchain.validate_block(block) == True

    # Create block with unauthorized validator
    block.validator = "unauthorized-node"
    block.block_hash = Blockchain.calculate_block_hash(block)

    assert blockchain.validate_block(block) == False
    return True


# ============================================================================
# DATA STRUCTURE TESTS
# ============================================================================

def test_transaction_model():
    """Test Transaction data model"""
    tx = Transaction(
        voter_id="voter-123",
        election_id=1,
        encrypted_vote="0x1234567890abcdef",
        ballot_hash="0xabcdef1234567890",
        proof="0x...",
        timestamp=1699360000
    )

    assert tx.voter_id == "voter-123"
    assert tx.election_id == 1
    assert tx.encrypted_vote == "0x1234567890abcdef"
    return True

def test_block_serialization():
    """Test block to/from dictionary conversion"""
    tx = Transaction(
        voter_id="voter-1",
        election_id=1,
        encrypted_vote="0x1234",
        ballot_hash="0xabcd",
        timestamp=1699360000
    )

    block = Block(
        index=1,
        prev_hash="0x" + "0" * 64,
        timestamp=1699360010,
        transactions=[tx],
        validator="validator-1",
        block_hash="0x" + "1" * 64,
        signature="0x" + "2" * 64
    )

    # Serialize to dict
    block_dict = block.to_dict()

    assert block_dict["index"] == 1
    assert block_dict["validator"] == "validator-1"
    assert len(block_dict["transactions"]) == 1
    assert block_dict["transactions"][0]["voter_id"] == "voter-1"

    # Deserialize from dict
    restored_block = Block.from_dict(block_dict)

    assert restored_block.index == block.index
    assert restored_block.validator == block.validator
    assert len(restored_block.transactions) == 1
    return True


# ============================================================================
# INTEGRATION TESTS
# ============================================================================

def test_multi_validator_consensus():
    """Test that multiple validators can reach consensus"""
    # Create 3 blockchains (one per validator)
    blockchains = [Blockchain(), Blockchain(), Blockchain()]

    # Simulate block creation and consensus
    for block_num in range(1, 4):
        # Determine which validator creates this block (round-robin)
        creator_idx = (block_num - 1) % 3
        validator_name = f"validator-{creator_idx + 1}"

        # Create block
        prev_block = blockchains[creator_idx].get_latest_block()
        tx = Transaction(
            voter_id=f"voter-{block_num}",
            election_id=1,
            encrypted_vote=f"0x{block_num:04x}",
            ballot_hash=f"0x{block_num:04x}",
            timestamp=1699360000 + block_num
        )

        block = Block(
            index=block_num,
            prev_hash=prev_block.block_hash,
            timestamp=1699360010 + block_num,
            transactions=[tx],
            validator=validator_name
        )
        block.block_hash = Blockchain.calculate_block_hash(block)

        # Add to creator's chain
        blockchains[creator_idx].add_block(block)

        # Broadcast to other validators
        for i in range(3):
            if i != creator_idx:
                blockchains[i].add_block(block)

    # Verify all validators have same blockchain
    chain_hashes = []
    for i, bc in enumerate(blockchains):
        last_block_hash = bc.get_latest_block().block_hash
        chain_hashes.append(last_block_hash)
        assert bc.verify_integrity() == True

    # All should have same last block
    assert chain_hashes[0] == chain_hashes[1] == chain_hashes[2]
    return True

def test_vote_immutability_across_validators():
    """Test that votes are immutable once recorded on blockchain"""
    blockchain = Blockchain()

    # Record a vote
    original_vote = Transaction(
        voter_id="voter-immutable",
        election_id=1,
        encrypted_vote="0x1234567890abcdef",
        ballot_hash="0xabcd",
        timestamp=1699360000
    )

    prev_block = blockchain.get_latest_block()
    block = Block(
        index=1,
        prev_hash=prev_block.block_hash,
        timestamp=1699360010,
        transactions=[original_vote],
        validator="validator-1"
    )
    block.block_hash = Blockchain.calculate_block_hash(block)
    blockchain.add_block(block)

    # Verify vote is recorded
    assert blockchain.chain[1].transactions[0].encrypted_vote == "0x1234567890abcdef"
    assert blockchain.verify_integrity() == True

    # Try to modify the vote (simulating attack)
    blockchain.chain[1].transactions[0].encrypted_vote = "0xhackedvalue"

    # Blockchain integrity should now fail
    assert blockchain.verify_integrity() == False
    return True


# ============================================================================
# JSON-RPC TESTS
# ============================================================================

def test_json_rpc_structure():
    """Test JSON-RPC request/response format"""
    # Valid JSON-RPC request
    request = {
        "jsonrpc": "2.0",
        "method": "eth_sendTransaction",
        "params": [{"data": "0x1234"}],
        "id": 1
    }

    # Verify structure
    assert request["jsonrpc"] == "2.0"
    assert request["method"] in ["eth_sendTransaction", "eth_getTransactionReceipt", "eth_blockNumber", "eth_getBlockByNumber"]
    assert "id" in request

    # Valid JSON-RPC response
    response = {
        "jsonrpc": "2.0",
        "result": "0xabc123",
        "id": 1
    }

    assert response["jsonrpc"] == "2.0"
    assert "result" in response or "error" in response
    assert response["id"] == 1
    return True


# ============================================================================
# RUN ALL TESTS
# ============================================================================

def run_all_tests():
    """Run all tests and report results"""
    tests = [
        # Bootnode tests
        ("Bootnode initialization", test_bootnode_initialization),
        ("Peer registration", test_peer_registration),
        ("Peer discovery", test_peer_discovery),
        ("Peer heartbeat", test_peer_heartbeat),
        ("Stale peer cleanup", test_stale_peer_cleanup),

        # Blockchain tests
        ("Genesis block creation", test_genesis_block_creation),
        ("Block hash calculation", test_block_hash_calculation),
        ("Block validation", test_block_validation),
        ("Add block to chain", test_add_block_to_chain),
        ("Blockchain integrity verification", test_blockchain_integrity_verification),
        ("Chain immutability", test_chain_immutability),

        # PoA Consensus tests
        ("Round-robin block creation", test_round_robin_block_creation),
        ("Authorized validators", test_authorized_validators),

        # Data structure tests
        ("Transaction model", test_transaction_model),
        ("Block serialization", test_block_serialization),

        # Integration tests
        ("Multi-validator consensus", test_multi_validator_consensus),
        ("Vote immutability across validators", test_vote_immutability_across_validators),

        # JSON-RPC tests
        ("JSON-RPC structure", test_json_rpc_structure),
    ]

    print("\n" + "=" * 80)
    print("PoA BLOCKCHAIN IMPLEMENTATION TEST SUITE")
    print("=" * 80 + "\n")

    passed = 0
    failed = 0
    errors = []

    for test_name, test_func in tests:
        try:
            result = test_func()
            if result:
                print(f"✅ {test_name}")
                passed += 1
            else:
                print(f"❌ {test_name}")
                failed += 1
                errors.append((test_name, "Assertion failed"))
        except Exception as e:
            print(f"❌ {test_name}")
            failed += 1
            errors.append((test_name, str(e)))

    # Print summary
    print("\n" + "=" * 80)
    print("TEST SUMMARY")
    print("=" * 80 + "\n")

    print(f"✅ Passed: {passed}/{len(tests)}")
    print(f"❌ Failed: {failed}/{len(tests)}")

    if errors:
        print("\nErrors:")
        for test_name, error in errors:
            print(f"  - {test_name}: {error}")

    if failed == 0:
        print("\n✅ ALL TESTS PASSED!")
        print("\nTest Categories:")
        print("  ✅ Bootnode: 5/5 tests passed")
        print("  ✅ Blockchain: 6/6 tests passed")
        print("  ✅ PoA Consensus: 2/2 tests passed")
        print("  ✅ Data Structures: 2/2 tests passed")
        print("  ✅ Integration: 2/2 tests passed")
        print("  ✅ JSON-RPC: 1/1 tests passed")
        print("  ───────────────────────")
        print("  ✅ TOTAL: 18/18 tests passed\n")
        return True
    else:
        print(f"\n❌ {failed} TESTS FAILED\n")
        return False


if __name__ == "__main__":
    success = run_all_tests()
    sys.exit(0 if success else 1)