app.py

from flask import Flask, request, jsonify
from flask_cors import CORS
import joblib
import numpy as np
import pandas as pd
from web3 import Web3
import json
import logging
import pickle
from crypto_utils import extract_features
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.sequence import pad_sequences

app = Flask(__name__)
CORS(app)

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

# Load all ML models for both classical and deep learning
models = {
    'Classical': {
        'SHA256': joblib.load('models/sha256_model.pkl'),
        'RSA': joblib.load('models/rsa_model.pkl'),
        'DES': joblib.load('models/des_model.pkl')
    },
    'DL': {
        'Blowfish': load_model('models/blowfish_model.h5'),
        'AES': load_model('models/aes_model.h5')
    }
}

with open('models/aes_tokenizer.pkl', 'rb') as f:
    aes_tokenizer = pickle.load(f)
with open('models/blowfish_tokenizer.pkl', 'rb') as f:
    blowfish_tokenizer = pickle.load(f)

# Define expected features for each model
model_features = {
    'SHA256': ['length', 'entropy', 'mean', 'std_dev', 'min_value', 'max_value'],
    'RSA': ['length', 'entropy', 'high_byte_ratio', 'low_byte_ratio', 'mean', 'std_dev'],
    'Blowfish': ['length', 'block_count', 'entropy', 'mean', 'median', 'q1', 'q3'],
    'AES': ['length', 'entropy', 'mean', 'aes_specific'],
    'DES': ['length', 'entropy', 'des_specific']
}

# Blockchain setup
w3 = Web3(Web3.HTTPProvider('http://127.0.0.1:7545'))  
contract_address = '0x8D69F970B7e1C5A24B8144665cC02E261E390950'
with open('blockchain/contract_abi.json') as f:
    contract_abi = json.load(f)
contract = w3.eth.contract(address=contract_address, abi=contract_abi)

# Model node addresses 
node_addresses = {
    'SHA256': '0x37BA38E02A6157abd3804a1e87BDbc1E49D47b38',
    'RSA': '0x5856dC54d5A036f6CF6c30b28A55d41D2c44C53e',
    'Blowfish': '0x4f3B9A4E2C9B78b617687398063308FAB8F66D58',
    'AES': '0x8A23d95E52Cb269Dec4C7FDE433624021F17F202',
    'DES': '0xedB88af25Efe6616917e663294D066A67B91e098'
}

def get_dl_features(algo, text):
    if algo == 'AES':
        seq = aes_tokenizer.texts_to_sequences([text])
        pad = pad_sequences(seq, maxlen=128, padding='post')
        return pad
    elif algo == 'Blowfish':
        seq = blowfish_tokenizer.texts_to_sequences([text])
        pad = pad_sequences(seq, maxlen=128, padding='post')
        return pad
    else:
        return None

@app.route('/predict', methods=['POST'])
def predict():
    """Get predictions from both Classical and DL models and combine results"""
    encrypted_text = request.json.get('text', '')
    if not encrypted_text or not isinstance(encrypted_text, str):
        logger.error("Invalid input: 'text' is required and must be a string.")
        return jsonify({'error': "Invalid input: 'text' is required and must be a string."}), 400

    features = extract_features(encrypted_text)
    if not features:
        logger.error("Failed to extract features from input.")
        return jsonify({'error': 'Failed to extract features'}), 400

    logger.info(f"Extracted features: {features}")

    combined_predictions = {}
    for model_type in ['Classical', 'DL']:
        for algo, model in models[model_type].items():
            try:
                if model_type == 'DL' and algo in ['AES', 'Blowfish']:
                    feature_input = get_dl_features(algo, encrypted_text)
                    proba = model.predict(feature_input, verbose=0)[0]
                else:
                    expected_feats = model_features.get(algo, list(features.keys()))
                    selected_features = {feat: features[feat] for feat in expected_feats}
                    feature_df = pd.DataFrame([selected_features])
                    if model_type == 'DL':
                        proba = model.predict(feature_df.values, verbose=0)[0]
                    else:
                        proba = model.predict_proba(feature_df)[0]

                confidence = float(np.max(proba)) * 100
                calibrated_confidence = confidence * 0.9
                logger.info(f"Model: {algo}, Type: {model_type}, Raw Confidence: {confidence:.2f}%, Calibrated Confidence: {calibrated_confidence:.2f}%")
                if algo not in combined_predictions:
                    combined_predictions[algo] = {
                        'algorithm': algo,
                        'confidence_sum': 0.0,
                        'count': 0,
                        'nodeAddress': node_addresses[algo]
                    }
                combined_predictions[algo]['confidence_sum'] += calibrated_confidence
                combined_predictions[algo]['count'] += 1
            except Exception as e:
                logger.error(f"Error in {algo} model ({model_type}): {str(e)}")
                if algo not in combined_predictions:
                    combined_predictions[algo] = {
                        'algorithm': algo,
                        'confidence_sum': 0.0,
                        'count': 0,
                        'nodeAddress': node_addresses[algo],
                        'error': str(e)
                    }

    # Average confidence from both models
    predictions = []
    for algo, data in combined_predictions.items():
        avg_confidence = data['confidence_sum'] / data['count'] if data['count'] > 0 else 0.0
        pred = {
            'algorithm': algo,
            'confidence': round(avg_confidence, 2),
            'nodeAddress': data['nodeAddress']
        }
        if 'error' in data:
            pred['error'] = data['error']
        predictions.append(pred)

    consensus = max(predictions, key=lambda x: x['confidence']) if predictions else None

    scores = {}
    try:
        model_count = contract.functions.getModelCount().call()
        for i in range(model_count):
            model_info = contract.functions.models(i).call()
            scores[model_info[1]] = model_info[2] / 100
    except Exception as e:
        logger.error(f"Error fetching trust scores: {str(e)}")

    trust_score = scores.get(consensus['algorithm'], 0) if consensus else 0

    return jsonify({
        'predictions': predictions,
        'consensus': consensus,
        'trustScore': trust_score,
        'model_type': 'combined'
    })

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)