Spaces:
Configuration error
Configuration error
File size: 5,089 Bytes
f7c6670 88774d5 f7c6670 8d9926c f7c6670 8d9926c f7c6670 88774d5 f7c6670 88774d5 f7c6670 8d9926c f7c6670 8d9926c f7c6670 8d9926c f7c6670 8d9926c f7c6670 8d9926c f7c6670 8d9926c f7c6670 8d9926c f7c6670 8d9926c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | import os, joblib
import numpy as np
import pandas as pd
from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
SKIP_MODEL = os.getenv("SKIP_MODEL", "false") == "true"
app = Flask(__name__)
CORS(app, origins="*")
# ββ Load all model artifacts ββββββββββββββββββββββββββββββββββββββββββββββββ
MODEL_DIR = os.path.join(os.path.dirname(__file__), 'models')
if not SKIP_MODEL:
sentinel_brain = joblib.load(os.path.join(MODEL_DIR, 'sentinel_brain.joblib'))
le = joblib.load(os.path.join(MODEL_DIR, 'label_encoder.joblib'))
ohe = joblib.load(os.path.join(MODEL_DIR, 'ohe_encoder.joblib'))
freq_map = joblib.load(os.path.join(MODEL_DIR, 'freq_map.joblib'))
scaler = joblib.load(os.path.join(MODEL_DIR, 'scaler.joblib'))
selected_features = joblib.load(os.path.join(MODEL_DIR, 'selected_features.joblib'))
else:
sentinel_brain = None
le = None
ohe = None
freq_map = {}
scaler = None
selected_features = []
COLUMNS = [
'duration','protocol_type','service','flag','src_bytes','dst_bytes',
'land','wrong_fragment','urgent','hot','num_failed_logins','logged_in',
'num_compromised','root_shell','su_attempted','num_root','num_file_creations',
'num_shells','num_access_files','num_outbound_cmds','is_host_login',
'is_guest_login','count','srv_count','serror_rate','srv_serror_rate',
'rerror_rate','srv_rerror_rate','same_srv_rate','diff_srv_rate',
'srv_diff_host_rate','dst_host_count','dst_host_srv_count',
'dst_host_same_srv_rate','dst_host_diff_srv_rate','dst_host_same_src_port_rate',
'dst_host_srv_diff_host_rate','dst_host_serror_rate','dst_host_srv_serror_rate',
'dst_host_rerror_rate','dst_host_srv_rerror_rate','label','difficulty_level'
]
SEVERITY_MAP = {'normal':'None','DoS':'Critical','Probe':'Medium','R2L':'High','U2R':'Critical'}
# ββ Serve frontend ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
@app.route("/")
def index():
return send_from_directory("frontend", "index.html")
# ββ THIS IS THE KEY FIX: serve style.css, app.js, and any other static files
@app.route("/<path:filename>")
def static_files(filename):
return send_from_directory("frontend", filename)
# ββ Everything below is UNCHANGED ββββββββββββββββββββββββββββββββββββββββββ
def preprocess(df):
df = df.copy()
for col in COLUMNS:
if col not in df.columns:
df[col] = 0
if 'label' not in df.columns:
df['label'] = 'normal'
cats = ['protocol_type', 'flag']
enc_df = pd.DataFrame(
ohe.transform(df[cats]),
columns=ohe.get_feature_names_out(cats),
index=df.index
)
df = pd.concat([df, enc_df], axis=1)
df['service_freq'] = df['service'].map(freq_map).fillna(0)
for col in ['src_bytes', 'dst_bytes', 'duration']:
df[f'log_{col}'] = np.log1p(df[col].astype(float))
df['total_bytes'] = df['src_bytes'].astype(float) + df['dst_bytes'].astype(float)
df['src_bytes_ratio'] = df['src_bytes'].astype(float) / (df['total_bytes'] + 1e-5)
df['is_error_flag'] = df['flag'].isin(['S0','S1','S2','S3','REJ']).astype(int)
for f in selected_features:
if f not in df.columns:
df[f] = 0
feature_matrix = df[selected_features].values
feature_matrix = scaler.transform(feature_matrix)
return feature_matrix
@app.route('/health')
def health():
return jsonify({
'status': 'online',
'model_loaded': not SKIP_MODEL
})
@app.route('/predict', methods=['POST', 'OPTIONS'])
def predict():
if request.method == 'OPTIONS':
return jsonify({}), 200
if SKIP_MODEL:
return jsonify({
'status': 'error',
'message': 'Model not loaded (CI mode)'
}), 503
try:
data = request.get_json(force=True)
rows = data.get('rows', [])
df = pd.DataFrame(rows)
X = preprocess(df)
preds = sentinel_brain.predict(X)
proba = sentinel_brain.predict_proba(X)
classes = le.inverse_transform(preds)
results = [
{
'predicted_class': cls,
'severity': SEVERITY_MAP.get(cls, 'Unknown'),
'confidence': round(float(conf), 4),
'is_intrusion': cls != 'normal'
}
for cls, conf in zip(classes, proba.max(axis=1))
]
return jsonify({'status': 'ok', 'results': results})
except Exception as e:
import traceback
return jsonify({
'status': 'error',
'message': str(e),
'trace': traceback.format_exc()
}), 500
if __name__ == '__main__':
app.run(host='0.0.0.0', port=7860)
|