/* SPDX-License-Identifier: GPL-3.0-only */
#include "web_cookie_auth.h"
#include <stdio.h>
#include <string.h>
#include "esp_timer.h"
#include "freertos/FreeRTOS.h"
#include "mbedtls/sha256.h"
#include "secure_random.h"
#include "web_auth_parse.h"
#include "web_httpd_adapter.h"
#include "web_login_ui.h"
#include "web_serial_transport.h"

#define SESSION_COOKIE "__Host-sak-session"
#define PRELOGIN_COOKIE "__Host-sak-prelogin"
#define COOKIE_FLAGS "; Secure; HttpOnly; SameSite=Strict; Path=/; Max-Age="
#define CHALLENGE_US 120000000LL
#define WINDOW_US 60000000LL

typedef struct {
    int64_t expiry;
    uint8_t token_digest[32];
    uint8_t origin_digest[32];
    char csrf[65];
} challenge_t;

static portMUX_TYPE s_lock = portMUX_INITIALIZER_UNLOCKED;
static challenge_t s_challenges[4];
static bool s_ready;
static uint64_t s_epoch;
static int64_t s_window;
static unsigned s_attempts;
static web_cookie_auth_snapshot_t s_counts;

static bool equal(const void *a, const void *b, size_t size)
{
    const uint8_t *x = a, *y = b;
    unsigned difference = 0;
    for (size_t i = 0; i < size; ++i) difference |= x[i] ^ y[i];
    return difference == 0;
}

static bool header(httpd_req_t *r, const char *name, char *out, size_t size)
{
    size_t length = httpd_req_get_hdr_value_len(r, name);
    out[0] = 0;
    return length < size && httpd_req_get_hdr_value_str(r, name, out, size) == ESP_OK;
}

static bool origin(httpd_req_t *r, bool required, char canonical[129])
{
    char host[129] = {0}, supplied[137] = {0}, site[16] = {0};
    if (!header(r, "Host", host, sizeof(host))) return false;
    if (header(r, "Sec-Fetch-Site", site, sizeof(site))) {
        if (strcmp(site, "same-origin") && strcmp(site, "none")) return false;
    } else if (httpd_req_get_hdr_value_len(r, "Sec-Fetch-Site")) return false;
    if (!header(r, "Origin", supplied, sizeof(supplied))) {
        if (required || httpd_req_get_hdr_value_len(r, "Origin")) return false;
        int length = snprintf(supplied, sizeof(supplied), "https://%s", host);
        if (length < 0 || (size_t)length >= sizeof(supplied)) return false;
    }
    return web_auth_parse_origin(host, strlen(host), supplied, strlen(supplied), canonical);
}

static bool cookie(httpd_req_t *r, const char *name, char token[65])
{
    char cookies[1025] = {0};
    bool valid = header(r, "Cookie", cookies, sizeof(cookies)) &&
        web_auth_parse_cookie(cookies, strlen(cookies), name, token);
    secure_wipe(cookies, sizeof(cookies));
    return valid;
}

static bool cookies_valid(httpd_req_t *r)
{
    char cookies[1025] = {0}, token[65] = {0};
    bool present;
    esp_err_t error = httpd_req_get_hdr_value_str(r, "Cookie", cookies, sizeof(cookies));
    bool valid = error == ESP_ERR_NOT_FOUND ||
        (error == ESP_OK && httpd_req_get_hdr_value_len(r, "Cookie") < sizeof(cookies) &&
         web_auth_parse_optional_cookie(cookies, strlen(cookies), SESSION_COOKIE, token, &present) &&
         web_auth_parse_optional_cookie(cookies, strlen(cookies), PRELOGIN_COOKIE, token, &present));
    secure_wipe(cookies, sizeof(cookies));
    secure_wipe(token, sizeof(token));
    return valid;
}

static esp_err_t response(httpd_req_t *r, const char *status, const char *body)
{
    esp_err_t error = httpd_resp_set_status(r, status);
    if (error == ESP_OK) error = httpd_resp_set_type(r, "application/json; charset=utf-8");
    if (error == ESP_OK) error = httpd_resp_set_hdr(r, "Cache-Control", "no-store");
    if (error == ESP_OK) error = httpd_resp_set_hdr(r, "X-Content-Type-Options", "nosniff");
    if (error == ESP_OK) error = httpd_resp_set_hdr(r, "Referrer-Policy", "no-referrer");
    if (error == ESP_OK) error = httpd_resp_set_hdr(r, "X-Frame-Options", "DENY");
    if (error == ESP_OK) error = httpd_resp_sendstr(r, body);
    /* Never let HTTPD drain an attacker-controlled rejected request body. */
    return web_httpd_unread_body(r) ? ESP_FAIL : error;
}

static esp_err_t failure(httpd_req_t *r, const char *status, const char *code)
{
    bool security = !strcmp(status, "403 Forbidden");
    bool credentials = !strcmp(code, "invalid_credentials");
    bool throttled = !strcmp(code, "throttled");
    bool full = !strcmp(code, "capacity");
    taskENTER_CRITICAL(&s_lock);
    s_counts.security_rejections += security;
    s_counts.login_failures += credentials;
    s_counts.throttled += throttled;
    s_counts.capacity_rejections += full;
    taskEXIT_CRITICAL(&s_lock);
    char body[80];
    snprintf(body, sizeof(body), "{\"error\":\"%s\"}", code);
    return response(r, status, body);
}

static esp_err_t capacity(httpd_req_t *r)
{
    if (httpd_resp_set_hdr(r, "Retry-After", "5") != ESP_OK) return ESP_FAIL;
    return failure(r, "503 Service Unavailable", "capacity");
}

esp_err_t web_cookie_auth_start(void)
{
    esp_err_t error = web_session_store_init();
    taskENTER_CRITICAL(&s_lock);
    if (error == ESP_OK && s_epoch != UINT64_MAX) {
        ++s_epoch;
        secure_wipe(s_challenges, sizeof(s_challenges));
        s_window = 0;
        s_attempts = 0;
        s_ready = true;
    } else {
        s_ready = false;
        error = ESP_ERR_INVALID_STATE;
    }
    taskEXIT_CRITICAL(&s_lock);
    return error;
}

void web_cookie_auth_stop(void)
{
    taskENTER_CRITICAL(&s_lock);
    s_ready = false;
    if (s_epoch != UINT64_MAX) ++s_epoch;
    secure_wipe(s_challenges, sizeof(s_challenges));
    s_window = 0;
    s_attempts = 0;
    taskEXIT_CRITICAL(&s_lock);
    web_session_store_stop();
}

void web_cookie_auth_get_snapshot(web_cookie_auth_snapshot_t *snapshot)
{
    if (!snapshot) return;
    int64_t now = esp_timer_get_time();
    taskENTER_CRITICAL(&s_lock);
    *snapshot = s_counts;
    snapshot->ready = s_ready;
    snapshot->active_challenges = 0;
    for (unsigned i = 0; i < 4; ++i) {
        if (s_challenges[i].expiry <= now) secure_wipe(&s_challenges[i], sizeof(s_challenges[i]));
        else ++snapshot->active_challenges;
    }
    taskEXIT_CRITICAL(&s_lock);
}

void web_cookie_auth_clear_counters(void)
{
    taskENTER_CRITICAL(&s_lock);
    memset(&s_counts, 0, sizeof(s_counts));
    taskEXIT_CRITICAL(&s_lock);
}

static esp_err_t require(httpd_req_t *r, bool mutation, bool upgrade, size_t body_limit,
                         web_session_view_t *view, bool *allowed)
{
    char canonical[129] = {0}, token[65] = {0}, csrf[65] = {0};
    *allowed = false;
    memset(view, 0, sizeof(*view));
    if (!web_httpd_headers_valid(r) || !cookies_valid(r) || (!upgrade && strchr(r->uri, '?')) ||
        r->content_len > body_limit || r->method != (mutation ? HTTP_POST : HTTP_GET))
        return failure(r, "400 Bad Request", "invalid_request");
    if (!origin(r, mutation || upgrade, canonical))
        return failure(r, "403 Forbidden", "origin");
    taskENTER_CRITICAL(&s_lock);
    bool ready = s_ready;
    taskEXIT_CRITICAL(&s_lock);
    if (!ready) return failure(r, "503 Service Unavailable", "unavailable");
    esp_err_t error = ESP_ERR_NOT_FOUND;
    if (cookie(r, SESSION_COOKIE, token))
        error = web_session_store_lookup(token, strlen(token), canonical, strlen(canonical), view);
    secure_wipe(token, sizeof(token));
    if (error != ESP_OK) {
        if (error != ESP_ERR_NOT_FOUND)
            return failure(r, "503 Service Unavailable", "unavailable");
        if (!strcmp(r->uri, "/")) {
            if (httpd_resp_set_hdr(r, "Location", "/login") != ESP_OK) return ESP_FAIL;
            return response(r, "303 See Other", "");
        }
        return failure(r, "401 Unauthorized", "authentication_required");
    }
    if (mutation && (!header(r, "X-CSRF-Token", csrf, sizeof(csrf)) ||
        strlen(csrf) != 64U || !equal(csrf, view->csrf, 64U))) {
        secure_wipe(csrf, sizeof(csrf));
        secure_wipe(view, sizeof(*view));
        return failure(r, "403 Forbidden", "csrf");
    }
    secure_wipe(csrf, sizeof(csrf));
    *allowed = true;
    return ESP_OK;
}

esp_err_t web_cookie_auth_require(httpd_req_t *r, bool mutation, bool upgrade,
                                  web_session_view_t *view, bool *allowed)
{
    return require(r, mutation, upgrade, 0, view, allowed);
}

esp_err_t web_cookie_auth_require_json(httpd_req_t *r, size_t body_limit,
                                      web_session_view_t *view, bool *allowed)
{
    return require(r, true, false, body_limit, view, allowed);
}

static bool secret(char out[65])
{
    uint8_t bytes[32];
    bool ok = secure_random_fill(bytes, sizeof(bytes)) == ESP_OK;
    if (ok) {
        static const char hex[] = "0123456789abcdef";
        for (size_t i = 0; i < sizeof(bytes); ++i) {
            out[2*i] = hex[bytes[i] >> 4];
            out[2*i+1] = hex[bytes[i] & 15];
        }
        out[64] = 0;
    }
    secure_wipe(bytes, sizeof(bytes));
    return ok;
}

esp_err_t web_cookie_auth_handler(httpd_req_t *r)
{
    bool login = !strcmp(r->uri, "/api/login");
    bool bootstrap = !strcmp(r->uri, "/api/login-challenge");
    bool logout = !strcmp(r->uri, "/api/logout");
    bool document = !strcmp(r->uri, "/login");
    web_session_view_t view = {0};
    char canonical[129] = {0}, token[65] = {0}, csrf[65] = {0};
    char set_cookie[180] = {0}, body[513] = {0};
    web_auth_credentials_t credentials = {0};
    challenge_t candidate = {0};
    uint8_t digest[32] = {0}, origin_digest[32] = {0};
    esp_err_t result = ESP_FAIL;
    const char *status = "400 Bad Request", *code = "invalid_request";
    bool consumed = false, allowed = false;
    bool challenge_published = false;
    uint64_t epoch;
    int selected = -1;
    int64_t now = esp_timer_get_time();
    taskENTER_CRITICAL(&s_lock);
    bool ready = s_ready;
    epoch = s_epoch;
    taskEXIT_CRITICAL(&s_lock);
    if (!ready) { status = "503 Service Unavailable"; code = "unavailable"; goto deny; }
    if (!web_httpd_headers_valid(r) || !cookies_valid(r) || strchr(r->uri, '?') ||
        r->method != ((login || logout) ? HTTP_POST : HTTP_GET) ||
        (!login && r->content_len)) goto deny;
    if (document) { result = web_login_ui_send_response(r); goto cleanup; }
    if (!login && !bootstrap) {
        result = web_cookie_auth_require(r, logout, false, &view, &allowed);
        if (!allowed) goto cleanup;
        if (logout) {
            web_serial_transport_revoke_web_session(view.id);
            taskENTER_CRITICAL(&s_lock);
            ++s_counts.logouts;
            taskEXIT_CRITICAL(&s_lock);
            result = httpd_resp_set_hdr(r, "Set-Cookie", SESSION_COOKIE "=" COOKIE_FLAGS "0");
            if (result != ESP_OK) goto cleanup;
            result = response(r, "204 No Content", "");
        } else {
            /* Database usernames are restricted ASCII; encode nevertheless. */
            char username[97] = {0};
            size_t used = 0;
            for (size_t i = 0; i < view.principal.username_length && i < 16; ++i)
                used += (size_t)snprintf(username + used, sizeof(username) - used,
                                         "\\u%04x", (unsigned char)view.principal.username[i]);
            int64_t remaining = (view.expires_at_us - esp_timer_get_time()) / 1000000LL;
            snprintf(body, sizeof(body), "{\"username\":\"%s\",\"role\":\"%s\",\"csrf\":\"%s\",\"expires_in\":%lld}",
                     username, view.principal.role == USER_ROLE_ADMIN ? "admin" : "user", view.csrf,
                     (long long)(remaining > 0 ? remaining : 0));
            result = response(r, "200 OK", body);
        }
        goto cleanup;
    }
    if (!origin(r, login, canonical)) { status = "403 Forbidden"; code = "origin"; goto deny; }
    if (mbedtls_sha256((const uint8_t *)canonical, strlen(canonical), origin_digest, 0)) goto deny;
    if (bootstrap) {
        char flag[2];
        if (!header(r, "X-Login-Bootstrap", flag, sizeof(flag)) || strcmp(flag, "1")) {
            status = "403 Forbidden"; code = "csrf"; goto deny;
        }
    } else {
        char type[40];
        if (!header(r, "Content-Type", type, sizeof(type)) ||
            (strcmp(type, "application/json") && strcmp(type, "application/json; charset=utf-8"))) {
            status = "415 Unsupported Media Type"; code = "content_type"; goto deny;
        }
        if (!r->content_len || r->content_len > 512U) {
            status = "413 Payload Too Large"; code = "body_size"; goto deny;
        }
        if (cookie(r, SESSION_COOKIE, token) &&
            web_session_store_lookup(token, 64, canonical, strlen(canonical), &view) == ESP_OK) {
            status = "409 Conflict"; code = "already_authenticated"; goto deny;
        }
        if (!header(r, "X-CSRF-Token", csrf, sizeof(csrf)) || strlen(csrf) != 64) {
            status = "403 Forbidden"; code = "csrf"; goto deny;
        }
    }
    bool has_cookie = cookie(r, PRELOGIN_COOKIE, token);
    if (has_cookie && mbedtls_sha256((const uint8_t *)token, 64, digest, 0)) goto deny;
    taskENTER_CRITICAL(&s_lock);
    for (int i = 0; i < 4; ++i) {
        challenge_t *entry = &s_challenges[i];
        if (entry->expiry <= now) secure_wipe(entry, sizeof(*entry));
        if (s_ready && epoch == s_epoch && entry->expiry && has_cookie &&
            equal(entry->token_digest, digest, 32) && equal(entry->origin_digest, origin_digest, 32)) {
            if (bootstrap || equal(entry->csrf, csrf, 64)) {
                candidate = *entry;
                selected = i;
                if (login) { secure_wipe(entry, sizeof(*entry)); consumed = true; }
            }
        }
    }
    taskEXIT_CRITICAL(&s_lock);
    if (bootstrap) {
        bool fresh = selected < 0;
        if (fresh) {
            if (!secret(token) || !secret(candidate.csrf) ||
                mbedtls_sha256((const uint8_t *)token, 64, candidate.token_digest, 0)) {
                status = "503 Service Unavailable"; code = "unavailable"; goto deny;
            }
            memcpy(candidate.origin_digest, origin_digest, 32);
            candidate.expiry = now + CHALLENGE_US;
            taskENTER_CRITICAL(&s_lock);
            if (s_ready && epoch == s_epoch) for (int i = 0; i < 4; ++i) {
                if (!s_challenges[i].expiry) {
                    s_challenges[i] = candidate; selected = i; challenge_published = true; break;
                }
            }
            taskEXIT_CRITICAL(&s_lock);
            if (selected < 0) { result = capacity(r); goto cleanup; }
            snprintf(set_cookie, sizeof(set_cookie), PRELOGIN_COOKIE "=%s" COOKIE_FLAGS "120", token);
            if (httpd_resp_set_hdr(r, "Set-Cookie", set_cookie) != ESP_OK) goto cleanup;
        }
        snprintf(body, sizeof(body), "{\"csrf\":\"%s\",\"expires_in\":%lld}", candidate.csrf,
                 (long long)((candidate.expiry - now) / 1000000LL));
        result = response(r, "200 OK", body);
        goto cleanup;
    }
    if (!consumed) { status = "403 Forbidden"; code = "challenge_expired"; goto deny; }
    if (httpd_resp_set_hdr(r, "Set-Cookie", PRELOGIN_COOKIE "=" COOKIE_FLAGS "0") != ESP_OK) goto cleanup;
    size_t received = 0;
    int64_t deadline = now + 3000000LL;
    while (received < r->content_len && esp_timer_get_time() < deadline) {
        int count = httpd_req_recv(r, body + received, r->content_len - received);
        if (count <= 0) goto deny;
        received += (size_t)count;
    }
    if (received != r->content_len || !web_auth_parse_login(body, received, &credentials)) goto deny;
    now = esp_timer_get_time();
    unsigned attempts;
    int64_t retry;
    taskENTER_CRITICAL(&s_lock);
    ready = s_ready && epoch == s_epoch;
    if (now - s_window >= WINDOW_US) { s_window = now; s_attempts = 0; }
    attempts = s_attempts;
    if (ready && attempts < 5) { ++s_attempts; ++s_counts.login_attempts; }
    retry = (s_window + WINDOW_US - now + 999999LL) / 1000000LL;
    taskEXIT_CRITICAL(&s_lock);
    if (!ready) { status = "503 Service Unavailable"; code = "unavailable"; goto deny; }
    if (attempts >= 5) {
        char seconds[16];
        snprintf(seconds, sizeof(seconds), "%lld", (long long)retry);
        if (httpd_resp_set_hdr(r, "Retry-After", seconds) != ESP_OK) goto cleanup;
        result = failure(r, "429 Too Many Requests", "throttled");
        goto cleanup;
    }
    bool authenticated = false;
    user_principal_t principal = {0};
    esp_err_t error = user_database_authenticate_password(credentials.username, credentials.username_length,
        credentials.password, credentials.password_length, &principal, &authenticated);
    secure_wipe(body, sizeof(body));
    secure_wipe(&credentials, sizeof(credentials));
    if (error == ESP_OK && authenticated)
        error = web_session_store_issue(&principal, canonical, strlen(canonical), token, &view);
    secure_wipe(&principal, sizeof(principal));
    if (error == ESP_ERR_NO_MEM) { result = capacity(r); goto cleanup; }
    if (error != ESP_OK) { status = "503 Service Unavailable"; code = "unavailable"; goto deny; }
    if (!authenticated) { status = "401 Unauthorized"; code = "invalid_credentials"; goto deny; }
    snprintf(set_cookie, sizeof(set_cookie), SESSION_COOKIE "=%s" COOKIE_FLAGS "3600", token);
    if (httpd_resp_set_hdr(r, "Set-Cookie", set_cookie) != ESP_OK) {
        web_session_store_invalidate(view.id); goto cleanup;
    }
    result = response(r, "200 OK", "{\"authenticated\":true}");
    if (result != ESP_OK) web_session_store_invalidate(view.id);
    goto cleanup;
deny:
    result = failure(r, status, code);
cleanup:
    if (challenge_published && result != ESP_OK) {
        taskENTER_CRITICAL(&s_lock);
        if (epoch == s_epoch && selected >= 0 &&
            equal(s_challenges[selected].token_digest, candidate.token_digest, 32))
            secure_wipe(&s_challenges[selected], sizeof(s_challenges[selected]));
        taskEXIT_CRITICAL(&s_lock);
    }
    web_httpd_wipe_request(r, web_httpd_unread_body(r));
    secure_wipe(&view, sizeof(view));
    secure_wipe(token, sizeof(token));
    secure_wipe(csrf, sizeof(csrf));
    secure_wipe(set_cookie, sizeof(set_cookie));
    secure_wipe(body, sizeof(body));
    secure_wipe(&credentials, sizeof(credentials));
    secure_wipe(&candidate, sizeof(candidate));
    secure_wipe(digest, sizeof(digest));
    return result;
}