/*
 * SPDX-FileCopyrightText: 2015-2021 Espressif Systems (Shanghai) CO LTD
 *
 * SPDX-License-Identifier: Apache-2.0
 */

#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <stdarg.h>

#include "sys/socket.h"
#include "esp_rom_md5.h"
#include "esp_tls_crypto.h"
#include "mbedtls/sha256.h"

#include "esp_log.h"
#include "esp_check.h"

#include "http_utils.h"
#include "http_auth.h"

#define MD5_MAX_LEN (33)
#define SHA256_LEN (32)
#define SHA256_HEX_LEN (65)
#define HTTP_AUTH_BUF_LEN (1024)

static const char *TAG = "HTTP_AUTH";

/**
 * @brief      This function hash a formatted string with MD5 and format the result as ascii characters
 *
 * @param      md         The buffer will hold the ascii result
 * @param[in]  fmt        The format
 *
 * @return     Length of the result
 */
static int md5_printf(char *md, const char *fmt, ...)
{
    unsigned char *buf;
    unsigned char digest[MD5_MAX_LEN];
    int len, i;
    md5_context_t md5_ctx;
    va_list ap;
    va_start(ap, fmt);
    len = vasprintf((char **)&buf, fmt, ap);
    if (buf == NULL) {
        va_end(ap);
        return ESP_FAIL;
    }

    esp_rom_md5_init(&md5_ctx);
    esp_rom_md5_update(&md5_ctx, buf, len);
    esp_rom_md5_final(digest, &md5_ctx);

    for (i = 0; i < 16; ++i) {
        sprintf(&md[i * 2], "%02x", (unsigned int)digest[i]);
    }
    va_end(ap);

    free(buf);
    return MD5_MAX_LEN;
}

/**
 * @brief      This function hash a formatted string with SHA256 and format the result as ascii characters
 *
 * @param      sha          The buffer will hold the ascii result
 * @param[in]  fmt          The format
 *
 * @return     Length of the result
 */
static int sha256_sprintf(char *sha, const char *fmt, ...)
{

    unsigned char *buf;
    unsigned char digest[SHA256_LEN];
    int len, i;
    va_list ap;
    va_start(ap, fmt);
    len = vasprintf((char **)&buf, fmt, ap);
    if (buf == NULL) {
        va_end(ap);
        return ESP_FAIL;
    }

    int ret = 0;
    mbedtls_sha256_context sha256;
    mbedtls_sha256_init(&sha256);
    if (mbedtls_sha256_starts(&sha256, 0) != 0) {
        goto exit;
    }
    if (mbedtls_sha256_update(&sha256, buf, len) != 0) {
        goto exit;
    }
    if (mbedtls_sha256_finish(&sha256, digest) != 0) {
        goto exit;
    }

    for (i = 0; i < 32; ++i) {
        sprintf(&sha[i * 2], "%02x", (unsigned int)digest[i]);
    }
    sha[SHA256_HEX_LEN - 1] = '\0';
    ret = SHA256_HEX_LEN;

exit:
    free(buf);
    mbedtls_sha256_free(&sha256);
    va_end(ap);
    return ret;
}

char *http_auth_digest(const char *username, const char *password, esp_http_auth_data_t *auth_data)
{
    char *ha1, *ha2 = NULL;
    char *digest = NULL;
    char *auth_str = NULL;
    char *temp_auth_str = NULL;
    esp_err_t ret = ESP_OK;

    if (username == NULL ||
            password == NULL ||
            auth_data->nonce == NULL ||
            auth_data->uri == NULL ||
            auth_data->realm == NULL) {
        return NULL;
    }

    int digest_size = MD5_MAX_LEN;
    int (*digest_func)(char *digest, const char *fmt, ...) = md5_printf;
    if (!memcmp(auth_data->algorithm, "SHA256", strlen("SHA256")) ||
            !memcmp(auth_data->algorithm, "SHA-256", strlen("SHA-256"))) {
        digest_size = SHA256_HEX_LEN;
        digest_func = sha256_sprintf;
    }

    ha1 = calloc(1, digest_size);
    ESP_GOTO_ON_FALSE(ha1, ESP_FAIL, _digest_exit, TAG, "Memory exhausted");

    ha2 = calloc(1, digest_size);
    ESP_GOTO_ON_FALSE(ha2, ESP_FAIL, _digest_exit, TAG, "Memory exhausted");

    digest = calloc(1, digest_size);
    ESP_GOTO_ON_FALSE(digest, ESP_FAIL, _digest_exit, TAG, "Memory exhausted");

    if (digest_func(ha1, "%s:%s:%s", username, auth_data->realm, password) <= 0) {
        goto _digest_exit;
    }

    ESP_LOGD(TAG, "%s %s %s %s", "Digest", username, auth_data->realm, password);
    if ((strcasecmp(auth_data->algorithm, "md5-sess") == 0) ||
            (strcasecmp(auth_data->algorithm, "SHA256") == 0) ||
            (strcasecmp(auth_data->algorithm, "md5-sess") == 0)) {
        if (digest_func(ha1, "%s:%s:%016llx", ha1, auth_data->nonce, auth_data->cnonce) <= 0) {
            goto _digest_exit;
        }
    }
    if (digest_func(ha2, "%s:%s", auth_data->method, auth_data->uri) <= 0) {
        goto _digest_exit;
    }

    //support qop = auth
    if (auth_data->qop && strcasecmp(auth_data->qop, "auth-int") == 0) {
        if (digest_func(ha2, "%s:%s", ha2, "entity") <= 0) {
            goto _digest_exit;
        }
    }

    if (auth_data->qop) {
        // response=digest_func(HA1:nonce:nonceCount:cnonce:qop:HA2)
        if (digest_func(digest, "%s:%s:%08x:%016llx:%s:%s", ha1, auth_data->nonce, auth_data->nc, auth_data->cnonce, auth_data->qop, ha2) <= 0) {
            goto _digest_exit;
        }
    } else {
        // response=digest_func(HA1:nonce:HA2)
        if (digest_func(digest, "%s:%s:%s", ha1, auth_data->nonce, ha2) <= 0) {
            goto _digest_exit;
        }
    }
    int rc = asprintf(&auth_str, "Digest username=\"%s\", realm=\"%s\", nonce=\"%s\", uri=\"%s\", algorithm=%s, "
                      "response=\"%s\", qop=%s, nc=%08x, cnonce=\"%016"PRIx64"\"",
                      username, auth_data->realm, auth_data->nonce, auth_data->uri, auth_data->algorithm, digest, auth_data->qop, auth_data->nc, auth_data->cnonce);
    if (rc < 0) {
        ESP_LOGE(TAG, "asprintf() returned: %d", rc);
        ret = ESP_FAIL;
        goto _digest_exit;
    }
    if (auth_data->opaque) {
        rc = asprintf(&temp_auth_str, "%s, opaque=\"%s\"", auth_str, auth_data->opaque);
        // Free the previous memory allocated for `auth_str`
        free(auth_str);
        if (rc < 0) {
            ESP_LOGE(TAG, "asprintf() returned: %d", rc);
            ret = ESP_FAIL;
            goto _digest_exit;
        }
        auth_str = temp_auth_str;
    }
_digest_exit:
    free(ha1);
    free(ha2);
    free(digest);
    return (ret == ESP_OK) ? auth_str : NULL;
}

char *http_auth_basic(const char *username, const char *password)
{
    size_t out;
    char *user_info = NULL;
    char *digest = NULL;
    esp_err_t ret = ESP_OK;
    size_t n = 0;
    if (asprintf(&user_info, "%s:%s", username, password) < 0) {
        return NULL;
    }
    ESP_RETURN_ON_FALSE(user_info, NULL, TAG, "Memory exhausted");
    esp_crypto_base64_encode(NULL, 0, &n, (const unsigned char *)user_info, strlen(user_info));
    digest = calloc(1, 6 + n + 1);
    ESP_GOTO_ON_FALSE(digest, ESP_FAIL, _basic_exit, TAG, "Memory exhausted");
    strcpy(digest, "Basic ");
    esp_crypto_base64_encode((unsigned char *)digest + 6, n, &out, (const unsigned char *)user_info, strlen(user_info));
_basic_exit:
    free(user_info);
    return (ret == ESP_OK) ? digest : NULL;
}