#!/usr/bin/python3
from flask import Flask, request, redirect, abort
import requests as r
app = Flask(__name__)
proxy_node_port = 5000
static_node_port = 5001
logic_node_port = 5003
content_node_port = 5004
static_node = "static:" + str(static_node_port)
logic_node = "logic:" + str(logic_node_port)
content_node = "content:" + str(content_node_port)

# Только реальные страницы. /db/users/pwn.html больше не уходит в content (V1).
ALLOWED_HTML = {"index.html", "login.html", "register.html", "specific.html"}
ALLOWED_API = {"login", "register", "update_personal", "send_msg", "post_comment"}
STATIC_EXT = {"css", "png", "jpg", "ico", "js", "otf", "eot", "svg", "ttf", "woff", "woff2"}


def safe_path(path: str):
    if not path or ".." in path or "\\" in path or "\x00" in path:
        return False
    if path.startswith("db/") or "/db/" in path:
        return False
    return True


def fwd_headers():
    """PATCH V4 (defense in depth): не пропускаем клиентский X-Internal внутрь —
    иначе клиент мог бы попытаться выдать себя за внутренний узел. Также режем
    hop-by-hop заголовки, которые ломают проксирование."""
    drop = {"x-internal", "host", "content-length", "connection",
            "transfer-encoding", "keep-alive", "te", "upgrade"}
    return {k: v for k, v in request.headers.items() if k.lower() not in drop}


@app.route('/', defaults={'path': 'index.html'})
@app.route('/<path:path>', methods=["GET", "POST"])
def catch_all(path: str):
    if not safe_path(path):
        abort(404)
    if request.content_length and request.content_length > 1_000_000:
        abort(413)
    ext = path.split('.')[-1]
    if ext == "html":
        if path not in ALLOWED_HTML:
            abort(404)
        qs = request.full_path.split('?')[1] if '?' in request.full_path else ''
        try:
            resp = r.get(f"http://{content_node}/{path}?{qs}", headers=fwd_headers(),
                         timeout=10)
        except r.RequestException:
            abort(502)
        return resp.content, resp.status_code, resp.headers.items()
    if ext in STATIC_EXT:
        try:
            resp = r.get(f"http://{static_node}/{path}", stream=True, timeout=5)
        except r.RequestException:
            abort(502)
        ct = resp.headers.get("Content-Type", "application/octet-stream")
        headers = [("Cache-Control", "max-age=36000;"), ("Content-Type", ct)]
        return resp.content, resp.status_code, headers
    if path.startswith("api/"):
        action = path[4:].split("/", 1)[0]
        if action not in ALLOWED_API:
            abort(404)
        if request.method.lower() not in ("get", "post"):
            return redirect("/", code=302)
        passcall = getattr(r, request.method.lower())
        # PATCH V3: allow_redirects=False. requests по умолчанию ХОДИТ по Location,
        # который отдаёт logic, а Location раньше строился из заголовка Origin =>
        # proxy делал запрос во внутреннюю сеть за атакующего (SSRF).
        try:
            resp = passcall(f"http://{logic_node}/{path[4:]}", json=request.get_json(silent=True),
                            headers=fwd_headers(), allow_redirects=False, timeout=8)
        except r.RequestException:
            abort(502)
        return resp.content, resp.status_code, resp.headers.items()
    return redirect("/", code=302)


if __name__ == '__main__':
    app.run(host="0.0.0.0", port=proxy_node_port)