CIA/e-voting-system/blockchain-worker/worker.py

"""
Blockchain Worker Service

A simple HTTP service that handles blockchain operations for the main API.
This allows the main backend to delegate compute-intensive blockchain tasks
to dedicated worker nodes.

The worker exposes HTTP endpoints for:
- Adding blocks to a blockchain
- Verifying blockchain integrity
- Retrieving blockchain data
"""

from fastapi import FastAPI, HTTPException, status
from pydantic import BaseModel
from typing import Optional, Dict, Any
import logging
import json
from dataclasses import dataclass, asdict
import time
import sys
import os

# Add parent directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from backend.crypto.hashing import SecureHash
from backend.crypto.signatures import DigitalSignature

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="Blockchain Worker",
    description="Dedicated worker for blockchain operations",
    version="1.0.0"
)


# ============================================================================
# Models (duplicated from backend for worker independence)
# ============================================================================

@dataclass
class Block:
    """Block in the blockchain containing encrypted votes"""
    index: int
    prev_hash: str
    timestamp: float
    encrypted_vote: str
    transaction_id: str
    block_hash: str
    signature: str


class AddBlockRequest(BaseModel):
    """Request to add a block to blockchain"""
    election_id: int
    encrypted_vote: str
    transaction_id: str


class AddBlockResponse(BaseModel):
    """Response after adding block"""
    index: int
    block_hash: str
    signature: str
    timestamp: float


class VerifyBlockchainRequest(BaseModel):
    """Request to verify blockchain integrity"""
    election_id: int
    blockchain_data: Dict[str, Any]


class VerifyBlockchainResponse(BaseModel):
    """Response of blockchain verification"""
    valid: bool
    total_blocks: int
    total_votes: int


# ============================================================================
# In-Memory Blockchain Storage (for this worker instance)
# ============================================================================

class Blockchain:
    """
    In-memory blockchain for vote storage.

    This is duplicated from the backend but kept in-memory for performance.
    Actual persistent storage should be in the main backend's database.
    """

    def __init__(self, authority_sk: Optional[str] = None, authority_vk: Optional[str] = None):
        """Initialize blockchain"""
        self.chain: list = []
        self.authority_sk = authority_sk
        self.authority_vk = authority_vk
        self.signature_verifier = DigitalSignature()
        self._create_genesis_block()

    def _create_genesis_block(self) -> None:
        """Create the genesis block"""
        genesis_hash = "0" * 64
        genesis_block_content = self._compute_block_content(
            index=0,
            prev_hash=genesis_hash,
            timestamp=time.time(),
            encrypted_vote="",
            transaction_id="genesis"
        )
        genesis_block_hash = SecureHash.sha256_hex(genesis_block_content.encode())
        genesis_signature = self._sign_block(genesis_block_hash) if self.authority_sk else ""

        genesis_block = Block(
            index=0,
            prev_hash=genesis_hash,
            timestamp=time.time(),
            encrypted_vote="",
            transaction_id="genesis",
            block_hash=genesis_block_hash,
            signature=genesis_signature
        )
        self.chain.append(genesis_block)

    def _compute_block_content(
        self,
        index: int,
        prev_hash: str,
        timestamp: float,
        encrypted_vote: str,
        transaction_id: str
    ) -> str:
        """Compute deterministic block content for hashing"""
        content = {
            "index": index,
            "prev_hash": prev_hash,
            "timestamp": timestamp,
            "encrypted_vote": encrypted_vote,
            "transaction_id": transaction_id
        }
        return json.dumps(content, sort_keys=True, separators=(',', ':'))

    def _sign_block(self, block_hash: str) -> str:
        """Sign a block with authority's private key"""
        if not self.authority_sk:
            return ""

        try:
            signature = self.signature_verifier.sign(
                block_hash.encode(),
                self.authority_sk
            )
            return signature.hex()
        except Exception:
            # Fallback to simple hash-based signature
            return SecureHash.sha256_hex((block_hash + self.authority_sk).encode())

    def add_block(self, encrypted_vote: str, transaction_id: str) -> Block:
        """Add a new block to the blockchain"""
        if not self.verify_chain_integrity():
            raise ValueError("Blockchain integrity compromised. Cannot add block.")

        new_index = len(self.chain)
        prev_block = self.chain[-1]
        prev_hash = prev_block.block_hash
        timestamp = time.time()

        block_content = self._compute_block_content(
            index=new_index,
            prev_hash=prev_hash,
            timestamp=timestamp,
            encrypted_vote=encrypted_vote,
            transaction_id=transaction_id
        )
        block_hash = SecureHash.sha256_hex(block_content.encode())
        signature = self._sign_block(block_hash)

        new_block = Block(
            index=new_index,
            prev_hash=prev_hash,
            timestamp=timestamp,
            encrypted_vote=encrypted_vote,
            transaction_id=transaction_id,
            block_hash=block_hash,
            signature=signature
        )

        self.chain.append(new_block)
        return new_block

    def verify_chain_integrity(self) -> bool:
        """Verify blockchain integrity"""
        for i in range(1, len(self.chain)):
            current_block = self.chain[i]
            prev_block = self.chain[i - 1]

            # Check chain link
            if current_block.prev_hash != prev_block.block_hash:
                return False

            # Check block hash
            block_content = self._compute_block_content(
                index=current_block.index,
                prev_hash=current_block.prev_hash,
                timestamp=current_block.timestamp,
                encrypted_vote=current_block.encrypted_vote,
                transaction_id=current_block.transaction_id
            )
            expected_hash = SecureHash.sha256_hex(block_content.encode())

            if current_block.block_hash != expected_hash:
                return False

            # Check signature if available
            if self.authority_vk and current_block.signature:
                if not self._verify_block_signature(current_block):
                    return False

        return True

    def _verify_block_signature(self, block: Block) -> bool:
        """Verify a block's signature"""
        if not self.authority_vk or not block.signature:
            return True

        try:
            return self.signature_verifier.verify(
                block.block_hash.encode(),
                bytes.fromhex(block.signature),
                self.authority_vk
            )
        except Exception:
            expected_sig = SecureHash.sha256_hex((block.block_hash + self.authority_vk).encode())
            return block.signature == expected_sig

    def get_blockchain_data(self) -> dict:
        """Get complete blockchain state"""
        blocks_data = []
        for block in self.chain:
            blocks_data.append({
                "index": block.index,
                "prev_hash": block.prev_hash,
                "timestamp": block.timestamp,
                "encrypted_vote": block.encrypted_vote,
                "transaction_id": block.transaction_id,
                "block_hash": block.block_hash,
                "signature": block.signature
            })

        return {
            "blocks": blocks_data,
            "verification": {
                "chain_valid": self.verify_chain_integrity(),
                "total_blocks": len(self.chain),
                "total_votes": len(self.chain) - 1
            }
        }

    def get_vote_count(self) -> int:
        """Get number of votes recorded (excludes genesis block)"""
        return len(self.chain) - 1


class BlockchainManager:
    """Manages blockchain instances per election"""

    def __init__(self):
        self.blockchains: Dict[int, Blockchain] = {}

    def get_or_create_blockchain(
        self,
        election_id: int,
        authority_sk: Optional[str] = None,
        authority_vk: Optional[str] = None
    ) -> Blockchain:
        """Get or create blockchain for an election"""
        if election_id not in self.blockchains:
            self.blockchains[election_id] = Blockchain(authority_sk, authority_vk)
        return self.blockchains[election_id]


# Global blockchain manager
blockchain_manager = BlockchainManager()


# ============================================================================
# Health Check
# ============================================================================

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {"status": "healthy", "service": "blockchain-worker"}


# ============================================================================
# Blockchain Operations
# ============================================================================

@app.post("/blockchain/add-block", response_model=AddBlockResponse)
async def add_block(request: AddBlockRequest):
    """
    Add a block to an election's blockchain.

    This performs the compute-intensive blockchain operations:
    - Hash computation
    - Digital signature
    - Chain integrity verification
    """
    try:
        blockchain = blockchain_manager.get_or_create_blockchain(request.election_id)
        block = blockchain.add_block(
            encrypted_vote=request.encrypted_vote,
            transaction_id=request.transaction_id
        )

        logger.info(
            f"Block added - Election: {request.election_id}, "
            f"Index: {block.index}, Hash: {block.block_hash[:16]}..."
        )

        return AddBlockResponse(
            index=block.index,
            block_hash=block.block_hash,
            signature=block.signature,
            timestamp=block.timestamp
        )

    except ValueError as e:
        logger.error(f"Invalid blockchain state: {e}")
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail=str(e)
        )
    except Exception as e:
        logger.error(f"Error adding block: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Failed to add block to blockchain"
        )


@app.post("/blockchain/verify", response_model=VerifyBlockchainResponse)
async def verify_blockchain(request: VerifyBlockchainRequest):
    """
    Verify blockchain integrity.

    This performs cryptographic verification:
    - Chain hash integrity
    - Digital signature verification
    - Block consistency
    """
    try:
        blockchain = blockchain_manager.get_or_create_blockchain(request.election_id)

        # Verify the blockchain
        is_valid = blockchain.verify_chain_integrity()

        logger.info(
            f"Blockchain verification - Election: {request.election_id}, "
            f"Valid: {is_valid}, Blocks: {len(blockchain.chain)}"
        )

        return VerifyBlockchainResponse(
            valid=is_valid,
            total_blocks=len(blockchain.chain),
            total_votes=blockchain.get_vote_count()
        )

    except Exception as e:
        logger.error(f"Error verifying blockchain: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Failed to verify blockchain"
        )


@app.get("/blockchain/{election_id}")
async def get_blockchain(election_id: int):
    """
    Get complete blockchain state for an election.
    """
    try:
        blockchain = blockchain_manager.get_or_create_blockchain(election_id)
        return blockchain.get_blockchain_data()

    except Exception as e:
        logger.error(f"Error retrieving blockchain: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Failed to retrieve blockchain"
        )


@app.get("/blockchain/{election_id}/stats")
async def get_blockchain_stats(election_id: int):
    """Get blockchain statistics for an election"""
    try:
        blockchain = blockchain_manager.get_or_create_blockchain(election_id)

        return {
            "election_id": election_id,
            "total_blocks": len(blockchain.chain),
            "total_votes": blockchain.get_vote_count(),
            "is_valid": blockchain.verify_chain_integrity()
        }

    except Exception as e:
        logger.error(f"Error retrieving blockchain stats: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Failed to retrieve blockchain stats"
        )


if __name__ == "__main__":
    import uvicorn

    port = int(os.getenv("WORKER_PORT", "8001"))

    logger.info(f"Starting blockchain worker on port {port}")
    uvicorn.run(app, host="0.0.0.0", port=port, log_level="info")