app.py

# app.py
from flask import Flask, request, jsonify, render_template, abort, send_file, make_response, redirect, url_for, Response
from werkzeug.utils import secure_filename
from waf_engine import waf
from database import init_db, get_stats, unban_ip, get_all_logs, get_all_bans, clear_database, export_logs_csv, get_log_by_id, verify_admin, set_telegram_sync_token, link_telegram_account, get_telegram_status, clear_ai_knowledge, get_all_ai_rules, delete_ai_rule, disconnect_telegram_account
from config import Config
import json
import io
import jwt
import datetime
import os
from functools import wraps
import threading
import requests
import time
import uuid

# --- NEW: Import Waitress for Production Serving ---
try:
    from waitress import serve
except ImportError:
    serve = None

# Import our new PDF generators
try:
    from report_generator import generate_incident_pdf, generate_global_pdf
except ImportError:
    generate_incident_pdf = None
    generate_global_pdf = None

app = Flask(__name__)

# --- NEW: WALLPAPER UPLOAD CONFIG ---
app.config['MAX_CONTENT_LENGTH'] = Config.MAX_CONTENT_LENGTH
os.makedirs(Config.UPLOAD_FOLDER, exist_ok=True)
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'webp'}

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

# Initialize DB on start
init_db()

# --- TELEGRAM BACKGROUND POLLER ---
def telegram_poller():
    if not hasattr(Config, 'TELEGRAM_BOT_TOKEN') or not Config.TELEGRAM_BOT_TOKEN:
        return
    
    last_update_id = 0
    while True:
        try:
            url = f"https://api.telegram.org/bot{Config.TELEGRAM_BOT_TOKEN}/getUpdates?offset={last_update_id}&timeout=30"
            res = requests.get(url, timeout=35).json()
            
            if res.get('ok'):
                for item in res['result']:
                    last_update_id = item['update_id'] + 1
                    msg = item.get('message', {})
                    text = msg.get('text', '')
                    chat_id = msg.get('chat', {}).get('id')
                    
                    if text.startswith('/start ') and chat_id:
                        token = text.split(' ')[1].strip()
                        success = link_telegram_account(token, chat_id)
                        
                        if success:
                            msg_url = f"https://api.telegram.org/bot{Config.TELEGRAM_BOT_TOKEN}/sendMessage"
                            requests.post(msg_url, json={
                                "chat_id": chat_id, 
                                "text": "✅ *SentinelShield SOC* \n\nYour account has been successfully verified and linked! You will now receive critical threat alerts here.",
                                "parse_mode": "Markdown"
                            })
        except Exception as e:
            pass
        
        time.sleep(2)

threading.Thread(target=telegram_poller, daemon=True).start()

# --- NEW: IP WHITELIST CONFIGURATION ---
# IPs in this list will completely bypass WAF scanning
WHITELISTED_IPS = ['127.0.0.1']

# --- WAF MIDDLEWARE ---
@app.before_request
def waf_middleware():
    # 1. Cloud IP Extraction
    if request.headers.get('X-Forwarded-For'):
        client_ip = request.headers.get('X-Forwarded-For').split(',')[0].strip()
    else:
        client_ip = request.remote_addr

    # 2. God Mode Check
    if client_ip in WHITELISTED_IPS:
        return

    # 3. THE FIX: Only ignore harmless static files (CSS/JS/Images)
    # The WAF will now aggressively scan the Dashboard, APIs, and the Root path!
    if request.path.startswith('/static') or request.path == '/favicon.ico':
        return

    # 4. WAF Engine Scanning
    method = request.method
    url = request.url
    headers = dict(request.headers)
    body = request.get_data(as_text=True)

    decision = waf.inspect_request(client_ip, method, url, headers, body)

    if decision['action'] == 'BLOCKED':
        return jsonify({
            "error": "Request Blocked by SentinelShield",
            "reason": decision['reason'],
            "ip": client_ip
        }), 403

# --- AUTH DECORATORS ---
def token_required(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        token = request.cookies.get('auth_token')
        if not token:
            return jsonify({'message': 'Authentication Token is missing!', 'status': 401}), 401
        try:
            decoded = jwt.decode(token, Config.JWT_SECRET, algorithms=["HS256"])
            request.current_user = decoded['user'] 
        except:
            return jsonify({'message': 'Token is invalid or expired!', 'status': 401}), 401
        return f(*args, **kwargs)
    return decorated

def admin_page_required(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        token = request.cookies.get('auth_token')
        if not token:
            return redirect(url_for('admin_login_page'))
        try:
            jwt.decode(token, Config.JWT_SECRET, algorithms=["HS256"])
        except:
            return redirect(url_for('admin_login_page'))
        return f(*args, **kwargs)
    return decorated

# --- AUTH ENDPOINTS ---
@app.route('/admin-login', methods=['GET'])
def admin_login_page():
    return render_template('admin_login.html')

@app.route('/api/auth/login', methods=['POST'])
def api_auth_login():
    data = request.json
    if not data or not data.get('username') or not data.get('password'):
        return jsonify({'message': 'Missing credentials'}), 400
    
    user = verify_admin(data['username'], data['password'])
    if not user:
        return jsonify({'message': 'Invalid credentials'}), 401
        
    token = jwt.encode({
        'user': user['username'],
        'role': user['role'],
        'exp': datetime.datetime.utcnow() + datetime.timedelta(hours=Config.JWT_EXPIRATION_HOURS)
    }, Config.JWT_SECRET, algorithm="HS256")
    
    resp = make_response(jsonify({'status': 'success', 'message': 'Logged in successfully'}))
    resp.set_cookie('auth_token', token, httponly=True, samesite='Strict')
    return resp

@app.route('/api/auth/logout', methods=['POST'])
def api_auth_logout():
    resp = make_response(jsonify({'status': 'success'}))
    resp.set_cookie('auth_token', '', expires=0)
    return resp

# --- REVERSE PROXY MODE ---
@app.route('/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH'])
def reverse_proxy(path):
    """
    Acts as a Gateway. If the WAF allows the request, it gets forwarded here.
    """
    # NOW READS FROM DYNAMIC CONFIG MEMORY
    proxy_url = Config.REVERSE_PROXY_URL
    
    if not proxy_url:
        return jsonify({
            "message": "SentinelShield WAF Active.",
            "status": "Protected",
            "endpoint": f"/{path}",
            "note": "Configure a Target URL in the Dashboard Settings to forward traffic."
        }), 200

    target_url = f"{proxy_url}/{path}"
    if request.query_string:
        target_url += f"?{request.query_string.decode('utf-8')}"

    try:
        headers = {key: value for (key, value) in request.headers if key.lower() != 'host'}
        
        resp = requests.request(
            method=request.method,
            url=target_url,
            headers=headers,
            data=request.get_data(),
            cookies=request.cookies,
            allow_redirects=False,
            timeout=15
        )

        excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
        resp_headers = [(name, value) for (name, value) in resp.raw.headers.items() if name.lower() not in excluded_headers]
        
        return Response(resp.content, resp.status_code, resp_headers)
        
    except requests.exceptions.RequestException as e:
        return jsonify({"error": "502 Bad Gateway", "message": f"SentinelShield could not reach backend at {proxy_url}"}), 502

# --- DASHBOARD PAGE ---
@app.route('/')
@admin_page_required
def index():
    token = request.cookies.get('auth_token')
    exp_time = 0
    if token:
        try:
            decoded = jwt.decode(token, Config.JWT_SECRET, algorithms=["HS256"])
            exp_time = decoded.get('exp', 0)
        except Exception:
            pass
            
    return render_template('dashboard.html', exp_time=exp_time)

# --- API ROUTES ---
@app.route('/api/stats')
@token_required
def api_stats():
    stats = get_stats()
    return jsonify(stats)

@app.route('/api/logs')
@token_required
def api_logs():
    logs = get_all_logs()
    data = []
    for r in logs:
        country = r[10] if len(r) > 10 else "Unknown"
        data.append({
            "id": r[0], "time": r[1], "ip": r[2], "method": r[3], 
            "url": r[4], "attack": r[7], "score": r[8], "action": r[9], 
            "country": country
        })
    return jsonify(data)

@app.route('/api/logs/<int:log_id>')
@token_required
def api_log_detail(log_id):
    log = get_log_by_id(log_id)
    if log:
        return jsonify(log)
    return jsonify({'error': 'Log not found'}), 404

@app.route('/api/bans')
@token_required
def api_bans():
    bans = get_all_bans()
    data = [{"ip": r[0], "banned_at": r[1], "expires": r[2], "reason": r[3]} for r in bans]
    return jsonify(data)

@app.route('/api/unban/<ip>', methods=['POST'])
@token_required
def api_unban(ip):
    unban_ip(ip)
    return jsonify({"status": "success", "message": f"IP {ip} unbanned"})

@app.route('/api/settings', methods=['GET', 'POST'])
@token_required
def api_settings():
    if request.method == 'POST':
        data = request.json
        if 'block_threshold' in data: Config.BLOCK_THRESHOLD = int(data['block_threshold'])
        if 'rate_limit' in data: Config.MAX_REQUESTS_PER_WINDOW = int(data['rate_limit'])
        if 'ban_duration' in data: Config.BAN_DURATION = int(data['ban_duration'])
        if 'reverse_proxy_url' in data: Config.REVERSE_PROXY_URL = data['reverse_proxy_url']
        
        return jsonify({"status": "updated", "message": "Configuration Saved Successfully", "config": {
            "block_threshold": Config.BLOCK_THRESHOLD,
            "rate_limit": Config.MAX_REQUESTS_PER_WINDOW,
            "ban_duration": Config.BAN_DURATION,
            "reverse_proxy_url": Config.REVERSE_PROXY_URL
        }})
    
    return jsonify({
        "block_threshold": Config.BLOCK_THRESHOLD,
        "rate_limit": Config.MAX_REQUESTS_PER_WINDOW,
        "ban_duration": Config.BAN_DURATION,
        "reverse_proxy_url": Config.REVERSE_PROXY_URL or ""
    })

@app.route('/api/settings/wallpaper', methods=['POST'])
@token_required
def api_upload_wallpaper():
    if 'wallpaper' not in request.files:
        return jsonify({"status": "error", "message": "No file provided"}), 400
    
    file = request.files['wallpaper']
    if file.filename == '':
        return jsonify({"status": "error", "message": "No file selected"}), 400
        
    if file and allowed_file(file.filename):
        # Generate a safe, unique filename to bypass browser cache
        filename = secure_filename(file.filename)
        ext = filename.rsplit('.', 1)[1].lower()
        new_filename = f"custom_bg_{int(time.time())}.{ext}"
        filepath = os.path.join(Config.UPLOAD_FOLDER, new_filename)
        
        # Cleanup old wallpapers so we don't fill the hard drive
        for old_file in os.listdir(Config.UPLOAD_FOLDER):
            if old_file.startswith('custom_bg_'):
                try: os.remove(os.path.join(Config.UPLOAD_FOLDER, old_file))
                except: pass
                
        file.save(filepath)
        
        return jsonify({
            "status": "success", 
            "message": "Wallpaper updated successfully! Applying now...", 
            "url": url_for('static', filename=f'uploads/{new_filename}')
        })
        
    return jsonify({"status": "error", "message": "Invalid file type. Allowed: PNG, JPG, GIF, WEBP"}), 400

@app.route('/api/telegram/status')
@token_required
def api_telegram_status():
    status = get_telegram_status(request.current_user)
    return jsonify(status)

@app.route('/api/telegram/generate', methods=['POST'])
@token_required
def api_telegram_generate():
    sync_token = uuid.uuid4().hex[:8]
    set_telegram_sync_token(request.current_user, sync_token)
    bot_username = "SentinelShieldAlerts_bot" 
    link = f"https://t.me/{bot_username}?start={sync_token}"
    return jsonify({"status": "success", "link": link})

@app.route('/api/telegram/disconnect', methods=['POST'])
@token_required
def api_telegram_disconnect():
    """Wipes the Telegram Chat ID and Sync Token from the database."""
    success = disconnect_telegram_account(request.current_user)
    if success:
        return jsonify({"status": "success", "message": "Telegram disconnected successfully!"})
    return jsonify({"status": "error", "message": "Failed to disconnect Telegram."}), 400

# --- ADAPTIVE DEFENSE ENDPOINTS ---
@app.route('/api/ai/summary')
@token_required
def api_get_ai_summary():
    """Fetches the active AI knowledge base for the Settings panel."""
    rules = get_all_ai_rules()
    data = []
    for r in rules:
        data.append({
            "id": r[0], "pattern": r[1], "attack_type": r[2], 
            "confidence": r[3], "created_at": r[4]
        })
    return jsonify(data)

@app.route('/api/ai/delete/<int:rule_id>', methods=['POST'])
@token_required
def api_delete_ai_rule(rule_id):
    """Surgically removes a false-positive rule from the AI's memory and hot-reloads the WAF."""
    success = delete_ai_rule(rule_id)
    if success:
        from rules import load_custom_rules
        load_custom_rules()
        return jsonify({"status": "success", "message": "Rule surgically removed from live memory."})
    return jsonify({"status": "error", "message": "Failed to delete rule."}), 400

# --- MAINTENANCE & REPORTS ---
@app.route('/api/database/clear', methods=['POST'])
@token_required
def api_clear_db():
    clear_database()
    return jsonify({"status": "success", "message": "Database Logs Cleared Successfully"})

@app.route('/api/ai/clear', methods=['POST'])
@token_required
def api_clear_ai():
    clear_ai_knowledge()
    from rules import load_custom_rules
    load_custom_rules()
    return jsonify({"status": "success", "message": "AI Learning Data Cleared Successfully"})

@app.route('/api/report/download')
@token_required
def api_download_report():
    if not generate_global_pdf:
        return jsonify({"error": "ReportLab missing", "message": "Run: pip install reportlab"}), 501
        
    stats = get_stats()
    raw_logs = get_all_logs()
    logs_data = []
    for r in raw_logs:
        logs_data.append({
            "id": r[0], "time": r[1], "ip": r[2], "method": r[3], 
            "url": r[4], "attack": r[7], "score": r[8]
        })
        
    pdf_buffer = generate_global_pdf(stats, logs_data)
    
    return send_file(
        pdf_buffer,
        mimetype='application/pdf',
        as_attachment=True,
        download_name='Sentinel_Global_Report.pdf'
    )

@app.route('/api/report/pdf/<int:log_id>')
@token_required
def api_download_pdf_report(log_id):
    if not generate_incident_pdf:
        return jsonify({"error": "ReportLab missing", "message": "Run: pip install reportlab"}), 501
        
    log = get_log_by_id(log_id)
    if not log:
        return jsonify({"error": "Not Found", "message": "Log entry does not exist."}), 404
        
    pdf_buffer = generate_incident_pdf(log)
    return send_file(
        pdf_buffer,
        mimetype='application/pdf',
        as_attachment=True,
        download_name=f'Sentinel_Incident_{log_id}.pdf'
    )

if __name__ == '__main__':
    print("\n=======================================================")
    print("🛡️  SENTINELSHIELD WAF IS ONLINE")
    print("🚀 Running on PRODUCTION WSGI Server (Waitress)")
    print("🌐 Accessible at: http://127.0.0.1:5000/")
    print("=======================================================\n")
    
    if serve:
        serve(app, host='0.0.0.0', port=5000, threads=6) 
    else:
        print("⚠️ Waitress not installed. Falling back to development server.")
        app.run(host='0.0.0.0', port=5000, debug=True)