chain.py

import hashlib
import json
import requests
import os
import sys
from time import time
from uuid import uuid4
from urllib.parse import urlparse
from flask import Flask, jsonify, request
from flask_cors import CORS 
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import serialization, hashes

class Blockchain:
    # --- CONFIGURATION CONSTANTS ---
    MIN_TRANSACTION_AMOUNT = 0.001
    MAX_BLOCK_SIZE = 10  # Max transactions per block
    DIFFICULTY_ADJUSTMENT_INTERVAL = 5 # Adjust difficulty every 5 blocks
    BLOCK_GENERATION_TARGET = 10 # Target time (seconds) to mine a block

    def __init__(self, port):
        self.chain = []
        self.current_transactions = []
        self.nodes = set()
        self.port = port
        self.node_identifier = str(uuid4()).replace('-', '')
        self.chain_file = f'chain_{port}.json'
        
        # PERSISTENCE: Load state on startup
        if os.path.exists(self.chain_file):
            print(f"[*] Loading chain from {self.chain_file}...")
            self.load_chain()
        else:
            print("[*] Creating Genesis Block...")
            # Genesis block has fixed difficulty of 4
            self.new_block(previous_hash='1', proof=100, difficulty=4)

    def register_node(self, address):
        """Add a new node to the list of nodes"""
        parsed_url = urlparse(address)
        if parsed_url.netloc:
            self.nodes.add(parsed_url.netloc)
        elif parsed_url.path:
            # Accepts 'localhost:5000' without http://
            self.nodes.add(parsed_url.path)
        else:
            raise ValueError('Invalid URL')

    def save_chain(self):
        """Save the chain to disk (Fault Tolerance)"""
        try:
            with open(self.chain_file, 'w') as f:
                json.dump(self.chain, f, indent=4)
        except Exception as e:
            print(f"[!] Error saving chain: {e}")

    def load_chain(self):
        """Load the chain from disk"""
        try:
            with open(self.chain_file, 'r') as f:
                self.chain = json.load(f)
        except Exception as e:
            print(f"[!] Error loading chain: {e}")
            # If load fails, start fresh
            self.chain = []
            self.new_block(previous_hash='1', proof=100, difficulty=4)

    def get_difficulty(self):
        """
        DYNAMIC DIFFICULTY ALGORITHM
        Adjusts the difficulty based on how fast the last few blocks were mined.
        """
        if len(self.chain) < self.DIFFICULTY_ADJUSTMENT_INTERVAL:
            return 4 # Default start difficulty
        
        last_block = self.chain[-1]
        
        # If we just hit an adjustment interval
        if last_block['index'] % self.DIFFICULTY_ADJUSTMENT_INTERVAL == 0:
            # Compare timestamp of current block vs block X ago
            prev_index = max(0, len(self.chain) - self.DIFFICULTY_ADJUSTMENT_INTERVAL)
            prev_block = self.chain[prev_index]
            
            time_taken = last_block['timestamp'] - prev_block['timestamp']
            expected_time = self.DIFFICULTY_ADJUSTMENT_INTERVAL * self.BLOCK_GENERATION_TARGET
            
            current_difficulty = last_block.get('difficulty', 4)
            
            if time_taken < expected_time / 2:
                print(f"[+] Mining too fast! Increasing difficulty to {current_difficulty + 1}")
                return current_difficulty + 1
            elif time_taken > expected_time * 2:
                new_diff = max(1, current_difficulty - 1)
                print(f"[-] Mining too slow. Decreasing difficulty to {new_diff}")
                return new_diff
            else:
                return current_difficulty
        
        return last_block.get('difficulty', 4)

    def valid_chain(self, chain):
        """
        DEEP VALIDATION
        Checks hashes, proofs, signatures, and recalculates the ledger history.
        """
        last_block = chain[0]
        current_index = 1
        
        # Temporary ledger to track balances during validation
        temp_balances = {} 

        while current_index < len(chain):
            block = chain[current_index]
            print(f"Verifying Block {block['index']}...")

            # 1. Check Hash Link
            if block['previous_hash'] != self.hash(last_block):
                print(f"Invalid Previous Hash in block {block['index']}")
                return False

            # 2. Check Proof of Work (using the difficulty recorded in THAT block)
            # Note: Genesis block usually has fixed proof, we skip strict check for block 1 if desired
            difficulty = block.get('difficulty', 4)
            if not self.valid_proof(last_block['proof'], block['proof'], difficulty):
                print(f"Invalid Proof in block {block['index']}")
                return False

            # 3. Verify Transactions (The "History Rewrite" Defense)
            for tx in block['transactions']:
                sender = tx['sender']
                recipient = tx['recipient']
                amount = tx['amount']
                signature = tx.get('signature')

                # Handle System Reward
                if sender == "0":
                    # EXCEPTION: Allow the "GENESIS_FUNDING" signature to break the 1-coin rule
                    # This allows us to fund Alice for testing.
                    if signature == "GENESIS_FUNDING":
                         temp_balances[recipient] = temp_balances.get(recipient, 0) + amount
                         continue

                    # Standard Rule: Mining rewards must be exactly 1
                    if amount != 1: 
                        print(f"Invalid Reward Amount: {amount}")
                        return False
                    
                    # Credit the miner
                    temp_balances[recipient] = temp_balances.get(recipient, 0) + amount
                    continue

                # Verify Cryptographic Signature
                message = f"{sender}{recipient}{amount}"
                if not self.verify_signature(sender, signature, message):
                    print("Invalid Transaction Signature Detected!")
                    return False

                # Verify Balance (Prevent spending money you didn't have AT THAT TIME)
                sender_balance = temp_balances.get(sender, 0)
                
                # Note: In a real sync, we should start temp_balances from the genesis
                # For this simplified demo, we assume the chain being validated includes logic
                # to build the state. We calculate the state change:
                temp_balances[sender] = sender_balance - amount
                temp_balances[recipient] = temp_balances.get(recipient, 0) + amount

            last_block = block
            current_index += 1
        
        return True

    def resolve_conflicts(self):
        """
        CONSENSUS ALGORITHM
        Resolve conflicts by replacing our chain with the longest VALID chain in the network.
        """
        neighbours = self.nodes
        new_chain = None
        max_length = len(self.chain)

        for node in neighbours:
            try:
                response = requests.get(f'http://{node}/chain')
                if response.status_code == 200:
                    length = response.json()['length']
                    chain = response.json()['chain']

                    # Check if their chain is longer AND valid
                    if length > max_length and self.valid_chain(chain):
                        max_length = length
                        new_chain = chain
            except requests.exceptions.ConnectionError:
                continue

        if new_chain:
            self.chain = new_chain
            self.save_chain()
            return True
        return False

    def new_block(self, proof, previous_hash=None, difficulty=4):
        """Create a new Block in the Blockchain"""
        block = {
            'index': len(self.chain) + 1,
            'timestamp': time(),
            'transactions': self.current_transactions,
            'proof': proof,
            'previous_hash': previous_hash or self.hash(self.chain[-1]),
            'difficulty': difficulty # Store difficulty so it can be verified later
        }
        self.current_transactions = []
        self.chain.append(block)
        self.save_chain()
        return block

    @staticmethod
    def verify_signature(public_key_pem, signature_hex, message):
        """Verifies an RSA signature"""
        try:
            public_key = serialization.load_pem_public_key(public_key_pem.encode())
            signature = bytes.fromhex(signature_hex)
            public_key.verify(
                signature,
                message.encode(),
                padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH),
                hashes.SHA256()
            )
            return True
        except Exception:
            return False

    def new_transaction(self, sender, recipient, amount, signature):
        """Creates a new transaction to go into the next mined Block"""
        
        # SPAM PROTECTION: Dust check
        if amount < self.MIN_TRANSACTION_AMOUNT:
            return False, "Amount too low"

        # SPAM PROTECTION: Block size limit
        if len(self.current_transactions) >= self.MAX_BLOCK_SIZE:
            return False, "Block is full, please wait"

        # SECURITY: Verify Signature immediately before adding to pool
        if sender != "0":
            message = f"{sender}{recipient}{amount}"
            if not self.verify_signature(sender, signature, message):
                return False, "Invalid Signature"
            
            # LOGIC: Check balance (Simplified global check)
            current_balance = self.get_balance(sender)
            if current_balance < amount:
                return False, "Insufficient Funds"

        self.current_transactions.append({
            'sender': sender,
            'recipient': recipient,
            'amount': amount,
            'signature': signature
        })
        return True, self.last_block['index'] + 1

    def get_balance(self, address):
        """Calculates balance by replaying the entire ledger"""
        balance = 0
        for block in self.chain:
            for tx in block['transactions']:
                if tx['recipient'] == address:
                    balance += tx['amount']
                if tx['sender'] == address:
                    balance -= tx['amount']
        
        # Deduct amounts in the current mempool (pending transactions)
        for tx in self.current_transactions:
            if tx['sender'] == address:
                balance -= tx['amount']
                
        return balance

    @property
    def last_block(self):
        return self.chain[-1]

    @staticmethod
    def hash(block):
        """Creates a SHA-256 hash of a Block"""
        # We must make sure the dictionary is ordered, or we'll get inconsistent hashes
        block_string = json.dumps(block, sort_keys=True).encode()
        return hashlib.sha256(block_string).hexdigest()

    def proof_of_work(self, last_proof, difficulty):
        """
        Proof of Work Algorithm:
        Find a number p' such that hash(pp') contains 'difficulty' leading zeroes
        """
        proof = 0
        while self.valid_proof(last_proof, proof, difficulty) is False:
            proof += 1
        return proof

    @staticmethod
    def valid_proof(last_proof, proof, difficulty):
        guess = f'{last_proof}{proof}'.encode()
        guess_hash = hashlib.sha256(guess).hexdigest()
        # Dynamic Difficulty: Check for N leading zeros
        return guess_hash[:difficulty] == "0" * difficulty

# --- FLASK SERVER SETUP ---

app = Flask(__name__)
CORS(app)

# Basic Argument Parsing for Port
port = 5000
if len(sys.argv) > 1:
    port = int(sys.argv[1])

blockchain = Blockchain(port)

@app.route('/mine', methods=['GET'])
def mine():
    # 1. Calculate Difficulty
    difficulty = blockchain.get_difficulty()
    
    # 2. Run Proof of Work
    last_block = blockchain.last_block
    last_proof = last_block['proof']
    proof = blockchain.proof_of_work(last_proof, difficulty)

    # 3. Receive Reward
    # FIX: Check if a specific address requested the mining reward (from the Web Wallet)
    miner_address = request.args.get('miner_address') 
    if not miner_address:
        miner_address = blockchain.node_identifier # Default to node if no address provided

    blockchain.new_transaction(
        sender="0",
        recipient=miner_address,
        amount=1,
        signature="REWARD"
    )

    # 4. Forge Block
    previous_hash = blockchain.hash(last_block)
    block = blockchain.new_block(proof, previous_hash, difficulty)

    # 5. Broadcast
    for node in blockchain.nodes:
        try:
            requests.get(f'http://{node}/nodes/resolve', timeout=1) 
        except:
            pass

    response = {
        'message': "New Block Forged",
        'index': block['index'],
        'transactions': block['transactions'],
        'proof': block['proof'],
        'difficulty': block['difficulty'],
        'previous_hash': block['previous_hash'],
    }
    return jsonify(response), 200

@app.route('/transactions/new', methods=['POST'])
def new_transaction():
    values = request.get_json()
    required = ['sender', 'recipient', 'amount', 'signature']
    if not all(k in values for k in required):
        return 'Missing values', 400

    success, message = blockchain.new_transaction(
        values['sender'], 
        values['recipient'], 
        values['amount'], 
        values['signature']
    )
    
    if success:
        return jsonify({'message': f'Transaction added to Block {message}'}), 201
    else:
        return jsonify({'message': f'Error: {message}'}), 400

@app.route('/chain', methods=['GET'])
def full_chain():
    response = {
        'chain': blockchain.chain,
        'length': len(blockchain.chain),
    }
    return jsonify(response), 200

@app.route('/nodes/register', methods=['POST'])
def register_nodes():
    values = request.get_json()
    nodes = values.get('nodes')
    if nodes is None:
        return "Error: Please supply a valid list of nodes", 400

    for node in nodes:
        blockchain.register_node(node)

    return jsonify({
        'message': 'New nodes have been added',
        'total_nodes': list(blockchain.nodes)
    }), 201

@app.route('/nodes/resolve', methods=['GET'])
def consensus():
    replaced = blockchain.resolve_conflicts()
    if replaced:
        response = {'message': 'Our chain was replaced', 'new_chain': blockchain.chain}
    else:
        response = {'message': 'Our chain is authoritative', 'chain': blockchain.chain}
    return jsonify(response), 200

# Helper to generate wallets locally for testing
@app.route('/wallet/new', methods=['GET'])
def create_wallet():
    private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    public_key = private_key.public_key()
    
    pem_priv = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption()
    ).decode('utf-8')
    
    pem_pub = public_key.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo
    ).decode('utf-8')
    
    return jsonify({'private_key': pem_priv, 'public_key': pem_pub}), 200

if __name__ == '__main__':
    # Threaded=True allows handling multiple requests (like sync + mining) better
    app.run(host='0.0.0.0', port=port, threaded=True)