#define _DEFAULT_SOURCE
#define _GNU_SOURCE
#include <string.h>
#include <stdlib.h>
#include <glib.h>
#include <gio/gio.h>
#include <gcrypt.h>
#include <unistd.h>
#include <glib/gi18n.h>

#include "gquarks.h"
#include "common.h"
#include "file-size.h"

static GSList *get_otps_from_encrypted_backup (const gchar       *path,
                                               const gchar       *password,
                                               gint32             max_file_size,
                                               GFile             *in_file,
                                               GFileInputStream  *in_stream,
                                               GError           **err);

static GSList *get_otps_from_plain_backup     (const gchar       *path,
                                               GError           **err);

static GSList *parse_authpro_json_data        (const gchar       *data,
                                               GError           **err);


GSList *
get_authpro_data (const gchar  *path,
                  const gchar  *password,
                  gint32        max_file_size,
                  gsize         db_size,
                  GError      **err)
{
    int safe_fd = path_open_safe_regular_file (path, err);
    if (safe_fd < 0) {
        return NULL;
    }

    goffset input_size = get_file_size (path);
    if (!is_secmem_available ((db_size + input_size)  * SECMEM_REQUIRED_MULTIPLIER, err)) {
        g_autofree gchar *msg = g_strdup_printf (_(
            "Your system's secure memory limit is not enough to securely import the data.\n"
            "You need to increase your system's memlock limit by following the instructions on our "
            "<a href=\"https://github.com/paolostivanin/OTPClient/wiki/Secure-Memory-Limitations\">secure memory wiki page</a>.\n"
            "This requires administrator privileges and is a system-wide setting that OTPClient cannot change automatically."
        ));
        g_clear_error (err);
        g_set_error (err, secmem_alloc_error_gquark (), NO_SECMEM_AVAIL_ERRCODE, "%s", msg);
        close (safe_fd);
        return NULL;
    }

    g_autofree gchar *fd_path = g_strdup_printf ("/proc/self/fd/%d", safe_fd);
    GFile *in_file = g_file_new_for_path (fd_path);
    GFileInputStream *in_stream = g_file_read (in_file, NULL, err);
    if (in_stream == NULL) {
        g_object_unref (in_file);
        close (safe_fd);
        return NULL;
    }

    GSList *otps;
    if (password != NULL) {
        otps = get_otps_from_encrypted_backup (fd_path, password, max_file_size, in_file, in_stream, err);
    } else {
        g_object_unref (in_stream);
        g_object_unref (in_file);
        otps = get_otps_from_plain_backup (fd_path, err);
    }
    close (safe_fd);
    return otps;
}


gchar *
export_authpro (const gchar *export_path,
                const gchar *password,
                json_t      *json_db_data)
{
    GError *err = NULL;
    json_t *root = NULL;
    GFile *out_gfile = NULL;
    GFileOutputStream *out_stream = NULL;
    gchar *json_data = NULL;
    guchar *salt = NULL, *iv = NULL;
    guchar *derived_key = NULL;
    gchar *enc_buf = NULL;
    gcry_cipher_hd_t hd = NULL;

    gsize db_size = json_dumpb (json_db_data, NULL, 0, 0);
    if (!is_secmem_available (db_size * SECMEM_REQUIRED_MULTIPLIER, &err)) {
        g_autofree gchar *msg = g_strdup_printf (_(
            "Your system's secure memory limit is not enough to securely export the database.\n"
            "You need to increase your system's memlock limit by following the instructions on our "
            "<a href=\"https://github.com/paolostivanin/OTPClient/wiki/Secure-Memory-Limitations\">secure memory wiki page</a>.\n"
            "This requires administrator privileges and is a system-wide setting that OTPClient cannot change automatically."
        ));
        g_clear_error (&err);
        g_set_error (&err, secmem_alloc_error_gquark (), NO_SECMEM_AVAIL_ERRCODE, "%s", msg);
        return g_strdup (err->message);
    }

    // set_new for fresh literals throughout: ownership flows up to root, so
    // the cleanup at `end:` only decrefs root and the children go with it.
    root = json_object ();
    json_t *auth_array = json_array ();
    json_object_set_new (root, "Authenticators", auth_array);

    /* Build Categories and AuthenticatorCategories from groups */
    json_t *categories_array = json_array ();
    json_t *auth_cat_array = json_array ();
    GHashTable *group_id_map = g_hash_table_new_full (g_str_hash, g_str_equal, g_free, g_free);
    {
        gsize idx;
        json_t *tmp_obj;
        guint cat_counter = 0;
        json_array_foreach (json_db_data, idx, tmp_obj) {
            const gchar *group = json_string_value (json_object_get (tmp_obj, "group"));
            if (group != NULL && group[0] != '\0' && !g_hash_table_contains (group_id_map, group)) {
                g_autofree gchar *cat_id = g_strdup_printf ("cat-%u", cat_counter++);
                g_hash_table_insert (group_id_map, g_strdup (group), g_strdup (cat_id));
                json_t *cat_obj = json_object ();
                json_object_set_new (cat_obj, "Id", json_string (cat_id));
                json_object_set_new (cat_obj, "Name", json_string (group));
                json_array_append_new (categories_array, cat_obj);
            }
        }
        json_array_foreach (json_db_data, idx, tmp_obj) {
            const gchar *group = json_string_value (json_object_get (tmp_obj, "group"));
            if (group != NULL && group[0] != '\0') {
                const gchar *cat_id = g_hash_table_lookup (group_id_map, group);
                const gchar *secret = json_string_value (json_object_get (tmp_obj, "secret"));
                if (cat_id != NULL && secret != NULL) {
                    json_t *ac_obj = json_object ();
                    json_object_set_new (ac_obj, "CategoryId", json_string (cat_id));
                    json_object_set_new (ac_obj, "AuthenticatorSecret", json_string (secret));
                    json_array_append_new (auth_cat_array, ac_obj);
                }
            }
        }
    }
    g_hash_table_destroy (group_id_map);
    json_object_set_new (root, "Categories", categories_array);
    json_object_set_new (root, "AuthenticatorCategories", auth_cat_array);
    json_object_set_new (root, "CustomIcons", json_array ());

    json_t *db_obj;
    gsize index;
    json_array_foreach (json_db_data, index, db_obj) {
        gboolean is_steam = FALSE;
        json_t *export_obj = json_object ();
        const gchar *issuer = json_string_value (json_object_get (db_obj, "issuer"));
        if (issuer != NULL) {
            if (g_ascii_strcasecmp (issuer, "steam") == 0) {
                json_object_set_new (export_obj, "Issuer", json_string ("Steam"));
                is_steam = TRUE;
            } else {
                json_object_set (export_obj, "Issuer", json_object_get (db_obj, "issuer"));
            }
        }
        const gchar *label = json_string_value (json_object_get (db_obj, "label"));
        if (label != NULL) {
            json_object_set (export_obj, "Username", json_object_get (db_obj, "label"));
        }
        json_object_set (export_obj, "Secret", json_object_get (db_obj, "secret"));
        json_object_set (export_obj, "Digits", json_object_get (db_obj, "digits"));
        json_object_set_new (export_obj, "Ranking", json_integer (0));
        json_object_set_new (export_obj, "Icon", json_null ());
        json_object_set_new (export_obj, "Pin", json_null ());
        const gchar *algo = json_string_value (json_object_get (db_obj, "algo"));
        if (algo != NULL && g_ascii_strcasecmp (algo, "SHA1") == 0) {
            json_object_set_new (export_obj, "Algorithm", json_integer (0));
        } else if (algo != NULL && g_ascii_strcasecmp (algo, "SHA256") == 0) {
            json_object_set_new (export_obj, "Algorithm", json_integer (1));
        } else if (algo != NULL && g_ascii_strcasecmp (algo, "SHA512") == 0) {
            json_object_set_new (export_obj, "Algorithm", json_integer (2));
        }
        const gchar *type = json_string_value (json_object_get (db_obj, "type"));
        if (type != NULL && g_ascii_strcasecmp (type, "TOTP") == 0) {
            json_object_set (export_obj, "Period", json_object_get (db_obj, "period"));
            json_object_set_new (export_obj, "Counter", json_integer (0));
            json_object_set_new (export_obj, "Type", is_steam ? json_integer (4) : json_integer (2));
        } else {
            json_object_set (export_obj, "Counter", json_object_get (db_obj, "counter"));
            json_object_set_new (export_obj, "Period", json_integer (0));
            json_object_set_new (export_obj, "Type", json_integer (1));
        }
        json_array_append_new (auth_array, export_obj);
    }

    json_data = json_dumps (root, JSON_COMPACT);
    if (json_data == NULL) {
        g_set_error (&err, generic_error_gquark (), GENERIC_ERRCODE, "Couldn't dump json data");
        goto end;
    }
    gsize json_data_size = strlen (json_data);

    out_gfile = g_file_new_for_path (export_path);
    out_stream = g_file_replace (out_gfile, NULL, FALSE, G_FILE_CREATE_REPLACE_DESTINATION | G_FILE_CREATE_PRIVATE, NULL, &err);
    if (out_stream == NULL) {
        // C1: g_file_replace already populated err on failure; nothing to add.
        // Previously the code blindly used out_stream in the password branch
        // and crashed on disk-full / R-O / no-perms.
        goto end;
    }

    if (password != NULL) {
        // encrypt the content and write the encrypted file to disk
        const gchar *header = "AUTHENTICATORPRO";
        salt = g_malloc0 (AUTHPRO_SALT_TAG);
        gcry_create_nonce (salt, AUTHPRO_SALT_TAG);
        iv = g_malloc0 (AUTHPRO_IV);
        gcry_create_nonce (iv, AUTHPRO_IV);
        derived_key = get_authpro_derived_key (password, salt);
        if (derived_key == NULL) {
            g_set_error (&err, key_deriv_gquark (), KEY_DERIVATION_ERRCODE, "Error while deriving the key.");
            goto end;
        }
        hd = open_cipher_and_set_data (derived_key, iv, AUTHPRO_IV);
        if (hd == NULL) {
            g_set_error (&err, generic_error_gquark (), GENERIC_ERRCODE, "Error while opening the cipher.");
            goto end;
        }
        enc_buf = gcry_calloc_secure (json_data_size, 1);
        if (enc_buf == NULL) {
            g_set_error (&err, secmem_alloc_error_gquark (), SECMEM_ALLOC_ERRCODE,
                         "Couldn't allocate secure memory for the encrypted payload.");
            goto end;
        }
        gpg_error_t gpg_err = gcry_cipher_encrypt (hd, enc_buf, json_data_size, json_data, json_data_size);
        if (gpg_err != GPG_ERR_NO_ERROR) {
            g_set_error (&err, generic_error_gquark (), GENERIC_ERRCODE,
                         "Failed to encrypt data: %s/%s",
                         gcry_strsource (gpg_err), gcry_strerror (gpg_err));
            goto end;
        }
        guchar tag[AUTHPRO_SALT_TAG];
        gcry_cipher_gettag (hd, tag, AUTHPRO_SALT_TAG);

        if (!output_stream_write_all_exact (G_OUTPUT_STREAM(out_stream), header, 16, &err) ||
            !output_stream_write_all_exact (G_OUTPUT_STREAM(out_stream), salt, AUTHPRO_SALT_TAG, &err) ||
            !output_stream_write_all_exact (G_OUTPUT_STREAM(out_stream), iv, AUTHPRO_IV, &err) ||
            !output_stream_write_all_exact (G_OUTPUT_STREAM(out_stream), enc_buf, json_data_size, &err) ||
            !output_stream_write_all_exact (G_OUTPUT_STREAM(out_stream), tag, AUTHPRO_SALT_TAG, &err)) {
            // err is already populated by g_output_stream_write; don't re-set.
            explicit_bzero (tag, AUTHPRO_SALT_TAG);
            goto end;
        }
        explicit_bzero (tag, AUTHPRO_SALT_TAG);
    } else {
        // write the plain json to disk
        if (!output_stream_write_all_exact (G_OUTPUT_STREAM(out_stream), json_data, json_data_size, &err)) {
            // err is set by g_output_stream_write.
            goto end;
        }
    }

end:
    if (hd != NULL) gcry_cipher_close (hd);
    if (derived_key != NULL) {
        explicit_bzero (derived_key, 32);
        gcry_free (derived_key);
    }
    if (enc_buf != NULL) gcry_free (enc_buf);
    g_free (iv);
    g_free (salt);
    if (json_data != NULL) gcry_free (json_data);
    if (root != NULL) json_decref (root);
    if (out_stream != NULL) g_object_unref (out_stream);
    if (out_gfile != NULL) g_object_unref (out_gfile);

    return (err != NULL ? g_strdup (err->message) : NULL);
}


static GSList *
get_otps_from_encrypted_backup (const gchar       *path,
                                const gchar       *password,
                                gint32             max_file_size,
                                GFile             *in_file,
                                GFileInputStream  *in_stream,
                                GError           **err)
{
    guchar header[16];
    gsize bytes_done = 0;
    if (!g_input_stream_read_all (G_INPUT_STREAM (in_stream), header, sizeof (header),
                                  &bytes_done, NULL, err) ||
        bytes_done != sizeof (header)) {
        if (err != NULL && *err == NULL)
            g_set_error (err, generic_error_gquark (), GENERIC_ERRCODE,
                         "Short read while reading Authenticator Pro header.");
        g_object_unref (in_stream);
        g_object_unref (in_file);
        return NULL;
    }

    gchar *decrypted_json = get_data_from_encrypted_backup (path, password, max_file_size, AUTHPRO, in_file, in_stream, err);
    if (decrypted_json == NULL) {
        return NULL;
    }

    GSList *otps = parse_authpro_json_data (decrypted_json, err);
    gcry_free (decrypted_json);

    return otps;
}


static GSList *
get_otps_from_plain_backup (const gchar  *path,
                            GError      **err)
{
    json_error_t j_err;
    json_t *json = json_load_file (path, JSON_DISABLE_EOF_CHECK | JSON_ALLOW_NUL, &j_err);
    if (!json) {
        g_printerr ("Error loading json: %s\n", j_err.text);
        g_set_error (err, generic_error_gquark (), GENERIC_ERRCODE,
                     "Error while loading the Authenticator Pro backup: %s", j_err.text);
        return NULL;
    }

    gchar *dumped_json = json_dumps (json, 0);
    json_decref (json);

    GSList *otps = parse_authpro_json_data (dumped_json, err);
    gcry_free (dumped_json);

    return otps;
}


static GSList *
parse_authpro_json_data (const gchar *data,
                         GError     **err)
{
    json_error_t jerr;
    json_t *root = json_loads (data, JSON_DISABLE_EOF_CHECK, &jerr);
    if (root == NULL) {
        g_set_error (err, generic_error_gquark (), GENERIC_ERRCODE, "%s", jerr.text);
        return NULL;
    }

    json_t *array = json_object_get (root, "Authenticators");
    if (array == NULL) {
        // jerr is stale here (json_loads succeeded above); reusing jerr.text
        // produced a misleading or empty error message.
        g_set_error (err, generic_error_gquark (), GENERIC_ERRCODE,
                     "Malformed Authenticator Pro backup: missing 'Authenticators' array.");
        json_decref (root);
        return NULL;
    }

    /* Build category ID -> name map */
    GHashTable *cat_map = g_hash_table_new_full (g_str_hash, g_str_equal, g_free, g_free);
    json_t *categories = json_object_get (root, "Categories");
    if (categories != NULL && json_is_array (categories)) {
        for (guint ci = 0; ci < json_array_size (categories); ci++) {
            json_t *cat = json_array_get (categories, ci);
            const gchar *id = json_string_value (json_object_get (cat, "Id"));
            const gchar *name = json_string_value (json_object_get (cat, "Name"));
            if (id != NULL && name != NULL)
                g_hash_table_insert (cat_map, g_strdup (id), g_strdup (name));
        }
    }

    /* Build secret -> category name map (first category only) */
    GHashTable *secret_cat_map = g_hash_table_new_full (g_str_hash, g_str_equal, g_free, g_free);
    json_t *auth_categories = json_object_get (root, "AuthenticatorCategories");
    if (auth_categories != NULL && json_is_array (auth_categories)) {
        for (guint ci = 0; ci < json_array_size (auth_categories); ci++) {
            json_t *ac = json_array_get (auth_categories, ci);
            const gchar *secret = json_string_value (json_object_get (ac, "AuthenticatorSecret"));
            const gchar *cat_id = json_string_value (json_object_get (ac, "CategoryId"));
            if (secret != NULL && cat_id != NULL) {
                const gchar *cat_name = g_hash_table_lookup (cat_map, cat_id);
                if (cat_name != NULL && !g_hash_table_contains (secret_cat_map, secret))
                    g_hash_table_insert (secret_cat_map, g_strdup (secret), g_strdup (cat_name));
            }
        }
    }

    GSList *otps = NULL;
    for (guint i = 0; i < json_array_size (array); i++) {
        json_t *obj = json_array_get (array, i);

        const gchar *secret_str = json_string_value (json_object_get (obj, "Secret"));
        if (secret_str == NULL) {
            g_printerr ("Skipping AuthPro entry with missing secret\n");
            continue;
        }

        otp_t *otp = g_new0 (otp_t, 1);
        otp->issuer = g_strdup (json_string_value (json_object_get (obj, "Issuer")));
        otp->account_name = g_strdup (json_string_value (json_object_get (obj, "Username")));
        otp->secret = secure_strdup (secret_str);
        otp->digits = (guint32)json_integer_value (json_object_get(obj, "Digits"));
        /* period and counter share a union, so only the member matching the
         * token type may be written; it is set per-type in the switch below.
         * Writing both here clobbered the HOTP counter with the Period. */

        gboolean skip = FALSE;
        guint32 algo = (guint32)json_integer_value (json_object_get(obj, "Algorithm"));
        switch (algo) {
            case 0:
                otp->algo = g_strdup ("SHA1");
                break;
            case 1:
                otp->algo = g_strdup ("SHA256");
                break;
            case 2:
                otp->algo = g_strdup ("SHA512");
                break;
            default:
                g_printerr ("Skipping token due to unsupported algo: %d\n", algo);
                skip = TRUE;
                break;
        }

        guint32 type = (guint32)json_integer_value (json_object_get(obj, "Type"));
        switch (type) {
            case 1:
                otp->type = g_strdup ("HOTP");
                otp->counter = json_integer_value (json_object_get (obj, "Counter"));
                break;
            case 2:
                otp->type = g_strdup ("TOTP");
                otp->period = (guint32)json_integer_value (json_object_get (obj, "Period"));
                break;
            case 4:
                otp->type = g_strdup ("TOTP");
                otp->period = (guint32)json_integer_value (json_object_get (obj, "Period"));
                g_free (otp->issuer);
                otp->issuer = g_strdup ("Steam");
                break;
            default:
                g_printerr ("Skipping token due to unsupported type: %d (3=Mobile-OTP, 5=Yandex)\n", type);
                skip = TRUE;
                break;
        }

        if (!skip) {
            const gchar *cat_name = g_hash_table_lookup (secret_cat_map, secret_str);
            otp->group = (cat_name != NULL) ? g_strdup (cat_name) : NULL;
            otps = g_slist_append (otps, otp);
        } else {
            gcry_free (otp->secret);
            g_free (otp->issuer);
            g_free (otp->account_name);
            g_free (otp->algo);
            g_free (otp->type);
            g_free (otp);
        }
    }

    g_hash_table_destroy (cat_map);
    g_hash_table_destroy (secret_cat_map);

    json_decref (root);

    return otps;
}
