"""Security hardening for the College KT Portal.

Everything here is switched on by `security.register(app)` (called at the end of app.py):

  * response headers   - Content-Security-Policy (script nonces, no inline handlers), clickjacking,
                         MIME-sniffing, referrer, permissions and cache protection
  * sessions           - HttpOnly + SameSite cookies, idle time-out, sign-out of other sessions on a
                         password change
  * brute force        - sign-in lock-out for the admin password, a cap on POSTs per address
  * uploads            - content inspection (real file type, macros, scripts, executables, zip bombs)
  * safe redirects     - `next=` may only point back into the portal
  * error pages        - no stack traces or internals shown to visitors

Run `python security.py --scan` to re-inspect files that were uploaded before these checks existed.
"""
import hmac
import io
import os
import re
import secrets
import threading
import time
import zipfile
from pathlib import Path

from flask import abort, g, render_template_string, request, session

# ---------------------------------------------------------------------------------------------
# settings
# ---------------------------------------------------------------------------------------------

IDLE_TIMEOUT = 2 * 60 * 60          # sign out after 2 hours without a request
LOGIN_FAILS = 5                     # wrong admin passwords (per address + username) before a lock-out
LOGIN_FAILS_PER_IP = 20             # ... or across all usernames from one address
LOGIN_WINDOW = 15 * 60              # failures are counted over, and a lock-out lasts, 15 minutes
POST_LIMIT = (120, 60)              # at most 120 form submissions per minute from one address
MAX_REQUEST_BYTES = 32 * 1024 * 1024
MIN_PASSWORD_LENGTH = 10
DEFAULT_ADMIN_PASSWORD = "admin123"

# ---------------------------------------------------------------------------------------------
# small in-memory throttles
# ---------------------------------------------------------------------------------------------


class Throttle:
    """Counts events per key inside a sliding window. Thread-safe; state is per process."""

    def __init__(self, limit, window):
        self.limit, self.window = limit, window
        self._hits, self._lock = {}, threading.Lock()

    def _recent(self, key, now):
        hits = [t for t in self._hits.get(key, ()) if now - t < self.window]
        if hits:
            self._hits[key] = hits
        else:
            self._hits.pop(key, None)
        return hits

    def blocked(self, key):
        with self._lock:
            return len(self._recent(key, time.time())) >= self.limit

    def retry_after(self, key):
        with self._lock:
            hits = self._recent(key, time.time())
            return max(1, int(self.window - (time.time() - hits[0]))) if hits else 0

    def hit(self, key):
        with self._lock:
            now = time.time()
            self._recent(key, now)
            self._hits.setdefault(key, []).append(now)

    def clear(self, key):
        with self._lock:
            self._hits.pop(key, None)

    def reset(self):
        with self._lock:
            self._hits.clear()


login_by_user = Throttle(LOGIN_FAILS, LOGIN_WINDOW)
login_by_ip = Throttle(LOGIN_FAILS_PER_IP, LOGIN_WINDOW)
posts_by_ip = Throttle(*POST_LIMIT)


def client_ip():
    return request.remote_addr or "unknown"


def login_locked(username):
    """Seconds the admin sign-in is locked for this address/username (0 = allowed)."""
    ip = client_ip()
    keys = ((login_by_ip, ip), (login_by_user, f"{ip}|{(username or '').strip().lower()}"))
    return max((t.retry_after(k) for t, k in keys if t.blocked(k)), default=0)


def login_failed(username):
    ip = client_ip()
    login_by_ip.hit(ip)
    login_by_user.hit(f"{ip}|{(username or '').strip().lower()}")


def login_succeeded(username):
    login_by_user.clear(f"{client_ip()}|{(username or '').strip().lower()}")


# ---------------------------------------------------------------------------------------------
# redirects, passwords, CSRF
# ---------------------------------------------------------------------------------------------


def safe_next(target, fallback):
    """Only allow a redirect back into this portal. Rejects //host, /\\host, http://host, control characters."""
    t = target or ""
    if (t.startswith("/") and not t.startswith("//") and "\\" not in t
            and not any(ord(c) < 32 for c in t) and "://" not in t.split("?", 1)[0]):
        return t
    return fallback


COMMON_PASSWORDS = {
    "password", "password1", "password123", "passw0rd", "admin", "admin123", "admin1234", "administrator",
    "qwerty123", "letmein123", "welcome123", "changeme", "123456789", "1234567890", "iloveyou1", "abc123456",
}


def password_problem(new, username=""):
    """Why a new password is unacceptable, or None."""
    if len(new) < MIN_PASSWORD_LENGTH:
        return f"The new password must be at least {MIN_PASSWORD_LENGTH} characters."
    if len(new) > 128:
        return "The new password is too long (128 characters at most)."
    if new.lower() in COMMON_PASSWORDS or (username and new.lower() == username.lower()):
        return "That password is too easy to guess. Choose a different one."
    if not (re.search(r"[A-Za-z]", new) and re.search(r"[\d\W_]", new)):
        return "Use letters plus at least one number or symbol."
    return None


def csrf_matches(sent):
    token = session.get("csrf")
    return bool(token) and isinstance(sent, str) and hmac.compare_digest(token.encode(), sent.encode())


def password_stamp(password_hash):
    """A short fingerprint of the stored password hash; a session stops working when it changes."""
    return (password_hash or "")[-16:]


# ---------------------------------------------------------------------------------------------
# upload inspection
# ---------------------------------------------------------------------------------------------

_EICAR = b"EICAR-STANDARD-" + b"ANTIVIRUS-TEST-FILE"
_OLE = b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1"
_PROGRAM_MAGIC = (b"\x7fELF", b"\xca\xfe\xba\xbe", b"\xcf\xfa\xed\xfe", b"\xce\xfa\xed\xfe",
                  b"\xfe\xed\xfa\xce", b"\xfe\xed\xfa\xcf", b"MSCF", b"\x4c\x00\x00\x00\x01\x14\x02\x00")


def _is_program(d):
    """Windows/Linux/macOS executables and shortcuts. (A text file that merely starts with "MZ" is not one.)"""
    return d.startswith(_PROGRAM_MAGIC) or (d[:2] == b"MZ" and b"PE\x00\x00" in d[:4096])
DANGEROUS_EXT = {
    "exe", "dll", "com", "scr", "pif", "msi", "msp", "bat", "cmd", "ps1", "psm1", "vbs", "vbe", "js", "jse", "wsf", "wsh",
    "hta", "cpl", "jar", "lnk", "reg", "sh", "bash", "py", "pyw", "php", "phtml", "pl", "rb", "asp", "aspx", "jsp",
    "html", "htm", "xhtml", "svg", "swf", "apk", "app", "dmg", "iso", "gadget", "inf", "scf", "url", "chm", "docm",
    "xlsm", "pptm", "dotm", "xlam", "ppam", "sldm",
}
_ZIP_MAX_ENTRIES = 3000
_ZIP_MAX_TOTAL = 300 * 1024 * 1024
_ZIP_MAX_RATIO = 200

OOXML = {"docx", "dotx", "xlsx", "xlsm", "pptx", "ppsx"}
ODF = {"odt", "ods", "odp"}
OLE_DOCS = {"doc", "dot", "xls", "ppt", "pps"}
PLAIN_TEXT = {"txt", "md", "log", "csv", "tsv"}
IMAGES = {"png", "jpg", "jpeg", "jpe", "gif", "webp", "bmp", "tif", "tiff", "heic", "heif", "avif", "ico"}


def _is_image(ext, d):
    if ext == "png":
        return d.startswith(b"\x89PNG\r\n\x1a\n")
    if ext in ("jpg", "jpeg", "jpe"):
        return d.startswith(b"\xff\xd8\xff")
    if ext == "gif":
        return d[:6] in (b"GIF87a", b"GIF89a")
    if ext == "webp":
        return d[:4] == b"RIFF" and d[8:12] == b"WEBP"
    if ext == "bmp":
        return d[:2] == b"BM"
    if ext in ("tif", "tiff"):
        return d[:4] in (b"II*\x00", b"MM\x00*")
    if ext in ("heic", "heif", "avif"):
        return d[4:8] == b"ftyp"
    if ext == "ico":
        return d[:4] == b"\x00\x00\x01\x00"
    return False


def _pdf_problem(d):
    if b"%PDF-" not in d[:1024]:
        return "it is not a real PDF file"
    # PDF names may be written with #xx escapes (/J#61vaScript); undo them before looking
    text = re.sub(rb"#([0-9A-Fa-f]{2})", lambda m: bytes([int(m.group(1), 16)]), d)
    if re.search(rb"/(JavaScript|JS|Launch|RichMedia|EmbeddedFile|XFA)\b", text):
        return "the PDF contains scripts, launch actions or embedded files"
    return None


def _zip_problem(ext, d):
    try:
        zf = zipfile.ZipFile(io.BytesIO(d))
    except zipfile.BadZipFile:
        return "it is not a valid archive/document"
    infos = zf.infolist()
    if len(infos) > _ZIP_MAX_ENTRIES:
        return "it contains too many files"
    total = 0
    names = []
    for i in infos:
        name = i.filename.replace("\\", "/")
        names.append(name)
        if name.startswith("/") or ".." in name.split("/") or re.match(r"^[A-Za-z]:", name):
            return "it contains unsafe file paths"
        if i.flag_bits & 0x1:
            return "it is password-protected, so it cannot be checked"
        total += i.file_size
        if total > _ZIP_MAX_TOTAL or (i.compress_size and i.file_size / i.compress_size > _ZIP_MAX_RATIO and i.file_size > 1024 * 1024):
            return "it expands to an unreasonable size"
        low = name.lower()
        if low.endswith("vbaproject.bin") or low.startswith(("basic/", "scripts/")) or "/vbaproject" in low or "activex/" in low:
            return "it contains macros or ActiveX controls"
        entry_ext = low.rsplit(".", 1)[-1] if "." in low.rsplit("/", 1)[-1] else ""
        if ext == "zip" and entry_ext in DANGEROUS_EXT:
            return f"it contains a program or script file ({name.rsplit('/', 1)[-1]})"
        if ext != "zip" and entry_ext in {"exe", "dll", "bat", "cmd", "scr", "com", "ps1", "vbs", "js", "jar", "msi", "hta", "lnk"}:
            return "it has an embedded program or script"
    if ext in OOXML and "[Content_Types].xml" not in names:
        return "it is not a real Office document"
    if ext in ODF and "mimetype" not in names:
        return "it is not a real OpenDocument file"
    return None


_SVG_BAD = re.compile(rb"<\s*script|[\s\"'/]on[a-z]+\s*=|javascript:|<\s*foreignObject|<\s*iframe|<\s*embed|<\s*object|<!ENTITY|xlink:href\s*=\s*[\"']\s*(?!#)", re.I)


def inspect_upload(ext, data):
    """Check that a file really is what its extension says and carries nothing active.
    Returns None when it looks fine, otherwise a short reason for the person uploading.
    This is a content check, not an antivirus engine: run real anti-malware over the uploads folder as well."""
    ext = (ext or "").lower().lstrip(".")
    d = bytes(data)
    if _EICAR in d:
        return "it is flagged as a virus test file"
    if _is_program(d):
        return "it is a program, not a document"
    if ext in DANGEROUS_EXT and ext != "svg":
        return f"the .{ext} type can run code and is not accepted"
    if ext in IMAGES:
        return None if _is_image(ext, d) else "its contents are not a real image"
    if ext == "svg":
        if b"<svg" not in d[:4096].lower():
            return "it is not a real SVG image"
        return "the SVG contains scripts or active content" if _SVG_BAD.search(d) else None
    if ext == "pdf":
        return _pdf_problem(d)
    if ext in OLE_DOCS:
        if not d.startswith(_OLE):
            return "it is not a real Office document"
        if "_VBA_PROJECT".encode("utf-16-le") in d or b"_VBA_PROJECT" in d:
            return "it contains macros"
        return None
    if ext == "rtf":
        if not d.lstrip()[:5] == b"{\\rtf":
            return "it is not a real RTF document"
        return "it contains embedded objects" if re.search(rb"\\obj(data|update|autlink)", d, re.I) else None
    if ext in OOXML or ext in ODF or ext == "zip":
        if not d.startswith((b"PK\x03\x04", b"PK\x05\x06")):
            return "it is not a real document or archive"
        return _zip_problem(ext, d)
    if ext in PLAIN_TEXT:
        return "it is not a text file" if b"\x00" in d[:8192] else None
    return "that file type is not accepted"


def scan_uploads(folder):
    """Re-inspect stored files (used by --scan). Yields (filename, reason)."""
    for p in sorted(Path(folder).iterdir()):
        if p.is_file():
            reason = inspect_upload(p.suffix, p.read_bytes())
            if reason:
                yield p.name, reason


# ---------------------------------------------------------------------------------------------
# Jinja helper for the few places rich text is shown
# ---------------------------------------------------------------------------------------------

def register(app, sanitize_rich_text=None, secure_cookies=None):
    """Turn the protections on for `app`."""
    from markupsafe import Markup

    secure = os.environ.get("KT_HTTPS") == "1" if secure_cookies is None else secure_cookies
    app.config.update(
        SESSION_COOKIE_HTTPONLY=True,
        SESSION_COOKIE_SAMESITE="Lax",
        SESSION_COOKIE_SECURE=secure,
        MAX_CONTENT_LENGTH=MAX_REQUEST_BYTES,
        MAX_FORM_MEMORY_SIZE=2 * 1024 * 1024,
        MAX_FORM_PARTS=2000,
    )
    app.config.setdefault("SECURITY_RATE_LIMIT", True)

    if sanitize_rich_text:
        # {{ value|rich }} - re-sanitises before marking as safe, so stored or re-displayed HTML can never carry script
        app.jinja_env.filters["rich"] = lambda v: Markup(sanitize_rich_text(v or "", max_len=50000))

    # ---- request gate: nonce, POST cap, idle time-out, session validity ------------------------
    @app.before_request
    def _security_gate():
        g.csp_nonce = secrets.token_urlsafe(18)
        if request.method == "POST" and app.config["SECURITY_RATE_LIMIT"]:
            if posts_by_ip.blocked(client_ip()):
                abort(429)
            posts_by_ip.hit(client_ip())
        if session.get("uid"):
            now = int(time.time())
            if now - int(session.get("seen", now)) > IDLE_TIMEOUT:
                session.clear()
                g.user = None
                return
            session["seen"] = now
            user = g.get("user")
            if user is not None:
                stamp = password_stamp(user["password_hash"])
                if "pwv" not in session:
                    session["pwv"] = stamp
                elif not hmac.compare_digest(str(session["pwv"]), stamp):
                    session.clear()          # the password was changed: every older session ends
                    g.user = None

    class _Nonce:
        """{{ csp_nonce }} - a Jinja *global* (not a context variable), so it also works inside imported macros."""
        def __str__(self):
            return g.get("csp_nonce", "")
    app.jinja_env.globals["csp_nonce"] = _Nonce()

    # ---- response headers -----------------------------------------------------------------------
    @app.after_request
    def _security_headers(resp):
        h = resp.headers
        nonce = g.get("csp_nonce", "")
        if resp.direct_passthrough or resp.mimetype != "text/html":
            # downloads and images: never run anything, never sniff
            h.setdefault("Content-Security-Policy", "default-src 'none'; sandbox; frame-ancestors 'none'")
        else:
            h["Content-Security-Policy"] = (
                "default-src 'self'; "
                f"script-src 'self' 'nonce-{nonce}'; "
                "style-src 'self' 'unsafe-inline'; "
                "img-src 'self' data: blob: https:; font-src 'self' data:; "
                "connect-src 'self'; media-src 'self'; "
                "frame-src https: http:; "                 # KT "Add Embed" shows pages from other sites (sandboxed in the template)
                "object-src 'none'; base-uri 'self'; form-action 'self'; frame-ancestors 'none'")
        h["X-Content-Type-Options"] = "nosniff"
        h["X-Frame-Options"] = "DENY"
        h["Referrer-Policy"] = "same-origin"
        h["Permissions-Policy"] = "camera=(), microphone=(), geolocation=(), payment=(), usb=(), serial=(), bluetooth=()"
        h["Cross-Origin-Opener-Policy"] = "same-origin"
        h["Cross-Origin-Resource-Policy"] = "same-origin"
        h["X-Permitted-Cross-Domain-Policies"] = "none"
        h["X-Robots-Tag"] = "noindex, nofollow"
        age = re.search(r"max-age=(\d+)", h.get("Cache-Control", ""))
        if age and int(age.group(1)) > 0:
            h["Cache-Control"] = h["Cache-Control"].replace("public", "private")    # logos/screenshots: this browser only
        else:
            h["Cache-Control"] = "no-store"           # pages hold private data: never keep them in a shared/back-button cache
            h["Pragma"] = "no-cache"
        if request.is_secure:
            h["Strict-Transport-Security"] = "max-age=31536000; includeSubDomains"
        return resp

    # ---- error pages ----------------------------------------------------------------------------
    messages = {
        400: ("Request rejected", "That request could not be processed. Go back, reload the page and try again."),
        403: ("Not allowed", "You do not have permission to do that."),
        404: ("Page not found", "That page does not exist."),
        405: ("Not allowed", "That action is not allowed here."),
        413: ("File too large", "That upload is too big. Try a smaller file."),
        429: ("Too many requests", "You are going too fast. Wait a minute and try again."),
        500: ("Something went wrong", "The portal hit an error. It has been logged - please try again."),
    }

    def _page(code):
        title, text = messages.get(code, messages[500])
        body = render_template_string(
            "<!doctype html><meta charset=utf-8><meta name=viewport content='width=device-width,initial-scale=1'>"
            "<title>{{ t }} - College KT Portal</title>"
            "<body style=\"font-family:Segoe UI,Arial,sans-serif;background:#f4f7fb;color:#1c2a3e;display:grid;place-items:center;min-height:100vh;margin:0\">"
            "<main style=\"background:#fff;border:1px solid #d5dbe5;border-radius:14px;padding:32px 36px;max-width:460px\">"
            "<h1 style=\"margin:0 0 8px;font-size:1.4rem\">{{ t }}</h1><p style=\"margin:0 0 18px;color:#5b6b82\">{{ m }}</p>"
            "<a href=\"/\" style=\"color:#006dff;font-weight:600\">Back to the portal</a></main>", t=title, m=text)
        return body, code

    for code in messages:
        app.register_error_handler(code, lambda e, c=code: _page(c))

    @app.errorhandler(Exception)
    def _unhandled(e):
        from werkzeug.exceptions import HTTPException
        if isinstance(e, HTTPException):
            return _page(e.code if e.code in messages else 500) if e.code >= 400 else e
        app.logger.exception("Unhandled error on %s %s", request.method, request.path)
        return _page(500)


if __name__ == "__main__":
    import sys
    if "--scan" in sys.argv:
        folder = Path(__file__).with_name("uploads")
        bad = list(scan_uploads(folder))
        print(f"Scanned {sum(1 for p in folder.iterdir() if p.is_file())} stored file(s) in {folder}")
        for name, why in bad:
            print(f"  SUSPICIOUS  {name}: {why}")
        print("No problems found." if not bad else f"{len(bad)} file(s) need attention (nothing was deleted).")
    else:
        print(__doc__)