#include "osfp_common.h"

#include "osfp.h"
#include "osfp_fingerprint.h"
#include "osfp_score_db.h"
#include "osfp_log.h"

#define OSFP_PERCENTILE 100

#define OSFP_SCORE_DB_FIELD_UINT_VALUE_MAX 65536

struct osfp_score_db_array_data {
    struct osfp_os_class_score *array_head;
    unsigned int    array_len;
};

struct osfp_score_db_hash_element {
    char *key;
    unsigned int keylen;
    struct osfp_os_class_score *score;
    UT_hash_handle hh;
};

struct osfp_score_db_hash_data {
    struct osfp_score_db_hash_element  *hash_head;
};

int osfp_score_db_array_add(void *data, struct osfp_os_class_score *os_class_score, void *value, unsigned int len)
{
    int ret = -1, i;
    unsigned int index;
    struct osfp_score_db_array_data *array_data = (struct osfp_score_db_array_data *)data;

    if (array_data == NULL || os_class_score == NULL || value == NULL || len != sizeof(unsigned int)) {
        goto exit;
    }

    if (array_data->array_head == NULL || array_data->array_len == 0) {
        goto exit;
    }

    index = *(unsigned int *)value;

    if (index >= array_data->array_len) {
        goto exit;
    }

    for (i = 0; i < OSFP_OS_CLASS_MAX; i++) {
        array_data->array_head[index].scores[i] += os_class_score->scores[i];
    }

    return 0;
exit:
    return ret;
}

struct osfp_os_class_score *osfp_score_db_array_match(void *data, void *value, unsigned int len)
{
    unsigned int index;
    struct osfp_score_db_array_data *array_data = (struct osfp_score_db_array_data *)data;

    if (array_data == NULL || value == NULL || len != sizeof(unsigned int)) {
        return NULL;
    }

    if (array_data->array_head == NULL || array_data->array_len == 0) {
        return NULL;
    }

    index = *(unsigned int *)value;

    if (index >= array_data->array_len) {
        return NULL;
    }

    return &((array_data->array_head)[index]);
}

void *osfp_score_db_array_create(void)
{
    struct osfp_score_db_array_data *array_data = calloc(1, sizeof(struct osfp_score_db_array_data));
    if (array_data == NULL) {
        return NULL;
    }

    array_data->array_head = calloc(OSFP_SCORE_DB_FIELD_UINT_VALUE_MAX, sizeof(struct osfp_os_class_score));
    if (array_data->array_head == NULL) {
        free(array_data);
        return NULL;
    }

    array_data->array_len = OSFP_SCORE_DB_FIELD_UINT_VALUE_MAX;

    return (void *)array_data;
}

void osfp_score_db_array_destroy(void *data) {
    struct osfp_score_db_array_data *array_data = (struct osfp_score_db_array_data *)data;

    if (array_data) {
        if (array_data->array_head) {
            free(array_data->array_head);
        }
        free(array_data);
    }
}

int osfp_score_db_hash_add(void *data, struct osfp_os_class_score *os_class_score, void *value, unsigned int len)
{
    int ret = -1, i;
    struct osfp_score_db_hash_data *hash_data = (struct osfp_score_db_hash_data *)data;
    struct osfp_score_db_hash_element *element = NULL;

    if (hash_data == NULL || os_class_score == NULL || value == NULL || len == 0) {
        goto exit;
    }

    HASH_FIND(hh, hash_data->hash_head, value, len, element);
    if (element == NULL) {
        element = (struct osfp_score_db_hash_element *)calloc(1, sizeof(struct osfp_score_db_hash_element));
        if (element == NULL) {
            goto exit;
        }
        element->key = strdup(value);
        element->keylen = len;
        element->score = (struct osfp_os_class_score *)calloc(1, sizeof(struct osfp_os_class_score));
        if (element->score == NULL) {
            free(element);
            element = NULL;
            goto exit;
        }
        HASH_ADD_KEYPTR(hh, hash_data->hash_head, element->key, element->keylen, element);
    }

    for (i = 0; i < OSFP_OS_CLASS_MAX; i++) {
        element->score->scores[i] += os_class_score->scores[i];
    }

    return 0;
exit:
    return ret;
}

struct osfp_os_class_score *osfp_score_db_hash_match(void *data, void *value, unsigned int len)
{
    struct osfp_score_db_hash_data *hash_data = (struct osfp_score_db_hash_data *)data;
    struct osfp_score_db_hash_element *element = NULL;

    if (data == NULL || value == NULL || len == 0) {
        return NULL;
    }

    if (hash_data->hash_head == NULL) {
        return NULL;
    }

    HASH_FIND(hh, hash_data->hash_head, value, len, element);
    if (element == NULL) {
        return NULL;
    }

    return element->score;
}

void *osfp_score_db_hash_create(void)
{
    return (void*)calloc(1, sizeof(struct osfp_score_db_hash_data));
}

void osfp_score_db_hash_destroy(void *data) {
    struct osfp_score_db_hash_data *hash_data = (struct osfp_score_db_hash_data *)data;
    struct osfp_score_db_hash_element *element = NULL;
    struct osfp_score_db_hash_element *tmp = NULL;

    if (hash_data) {
        if (hash_data->hash_head) {
            HASH_ITER(hh, hash_data->hash_head, element, tmp) {
                HASH_DELETE(hh, hash_data->hash_head, element);
                if (element) {
                    if (element->key) {
                        free(element->key);
                    }
                    if (element->score) {
                        free(element->score);
                    }
                    free(element);
                }
            }
        }
        free(hash_data);
    }
}

struct osfp_os_class_score *osfp_score_db_filed_match(struct osfp_field_score_db *db, void *value, unsigned int len)
{
    if (db == NULL || value == NULL || len == 0) {
        return NULL;
    }

    return  db->match(db->data, value, len);
}

char *osfp_score_db_read_file(char *fp_file)
{
    int ret = -1;
    char *file_buffer = NULL;
    unsigned int file_len = 0;
    FILE *fp = NULL;
    struct stat st;

    if (0 > stat(fp_file, &st)) {
        printf("stat() on '%s' failed.\n", fp_file);
        goto exit;
    }

    if (st.st_size == 0) { 
        printf("Empty file: %s.\n", fp_file);
        goto exit;
    }

    file_len = (unsigned int)st.st_size;
    file_buffer = malloc(file_len + 1);
    if (file_buffer == NULL) {
        printf("Not enough memory for file buffer. file name: %s\n",fp_file);
        goto exit;
    }

    fp = fopen(fp_file, "r");
    if (fp == NULL) {
        printf("Cannot open '%s' for reading.\n", fp_file);
        goto exit;
    }

    ret = fread(file_buffer, 1, file_len, fp);
    if (ret < 0) {
        free(file_buffer);
        fclose(fp);
        goto exit;
    }

    fclose(fp);

    file_buffer[file_len] = 0;

    return file_buffer;
exit:
    return NULL;
}

int osfp_score_db_load_field(struct osfp_field_score_db *db, cJSON *field, enum osfp_os_class_id os_class)
{
    int ret = -1;
    struct osfp_os_class_score os_class_score = {0};

    void *value_ptr;
    unsigned int value_len;

    switch (field->type) {
        case cJSON_Number:
            value_ptr = (void *)&field->valueint;
            value_len = sizeof(field->valueint);
            break;
        case cJSON_String:
            value_ptr = (void *)field->valuestring;
            value_len = strlen(field->valuestring) + 1;
            break;
        case cJSON_NULL:
            ret = 0;
            goto exit;
        default:
            goto exit;
    }

    os_class_score.scores[os_class] = 1;

    ret = db->add(db->data, &os_class_score, value_ptr, value_len);
    if (ret != 0) {
        goto exit;
    }

    db->entry_count++;

    return 0;
exit:
    return ret;
}

int osfp_score_db_load_entry(struct osfp_score_db *score_db, cJSON *entry)
{
    int ret = -1, i;
    cJSON *field = NULL;
    struct osfp_field_score_db *db;
    enum osfp_os_class_id os_class;

    if (score_db == NULL || entry == NULL) {
        goto exit;
    }

    field = cJSON_GetObjectItem(entry, osfp_fingerprint_get_field_name(OSFP_FIELD_OS));
    if (field == NULL || field->valuestring == NULL) {
        goto exit;
    }

    os_class = osfp_os_class_name_to_id(field->valuestring);
    if (os_class >= OSFP_OS_CLASS_MAX) {
        goto exit;
    }

    for (i = 0; i < OSFP_FIELD_OS; i++) {
        db = &score_db->field_score_dbs[i];
        if (db == NULL) {
            goto exit;
        }

        if (!db->enabled) {
            continue;
        }

        field = cJSON_GetObjectItem(entry, osfp_fingerprint_get_field_name(i));
        if (field == NULL) {
            printf("json entry missing field: %s\n%s\n",
                        osfp_fingerprint_get_field_name(i), cJSON_Print(entry));
            continue;
        }

        ret = osfp_score_db_load_field(db, field, os_class);
        if (ret != 0) {
            printf("json entry field load failed. field: %s\n%s\n",
                        osfp_fingerprint_get_field_name(i), cJSON_Print(entry));
            continue;
        }
    }

    score_db->entry_count++;
    score_db->os_class_entry_count[os_class]++;

    return 0;
exit:
    return ret;
}

int osfp_score_db_load(struct osfp_score_db *score_db, char *fp_file)
{
    int ret = OSFP_EINVAL, i;
    char *file_buffer;
    cJSON *root = NULL;
    cJSON *entry;
    struct osfp_field_score_db *field_score_db;

    if (score_db == NULL) {
        goto exit;
    }

    file_buffer = osfp_score_db_read_file(fp_file);
    if (file_buffer == NULL) {
        ret = OSFP_ERR_SCORE_DB_READ_FILE;
        goto exit;
    }

    root = cJSON_Parse(file_buffer);
    if (root == NULL) {
        ret = OSFP_ERR_SCORE_DB_PARSE_FILE;
        goto exit;
    }

    score_db->entry_count = cJSON_GetArraySize(root);

    for (i = 0; i < score_db->entry_count; i++) {
        entry = cJSON_GetArrayItem(root, i);
        if (entry) {
            ret = osfp_score_db_load_entry(score_db, entry);
            if (ret != 0) {
                osfp_log_debug("json entry load failed.\n%s\n", cJSON_Print(entry));
                continue;
            }
        } 
    }

    for (i = 0; i < OSFP_FIELD_MAX; i++) {
        field_score_db = &score_db->field_score_dbs[i];
        if (field_score_db->enabled) {
            score_db->perfect_score += osfp_fingerprint_get_field_importance(i);
        }
    }

    ret = OSFP_NOERR;
exit:
    if (root) {
        cJSON_Delete(root);
    }
    if (file_buffer) {
        free((char*)file_buffer);
    }
    return ret;
}

int osfp_score_db_score(struct osfp_score_db *score_db, unsigned int flags, struct osfp_fingerprint *fp, struct osfp_os_class_score *result_score)
{
    int ret = OSFP_EINVAL, i, j;

    unsigned int tmp_score;
    unsigned int perfect_score;
    unsigned int entry_count;
    unsigned int importance;

    struct osfp_os_class_score *os_class_score_matched;
    enum osfp_os_class_id os_class_id;

    struct osfp_fingerprint_field *field;
    struct osfp_field_score_db *field_score_db;

    if (score_db == NULL || fp == NULL || result_score == NULL) {
        goto exit;
    }

    memset(result_score, 0, sizeof(struct osfp_os_class_score));

    perfect_score = score_db->perfect_score;
    if (perfect_score == 0) {
        goto exit;
    }

    for (i = 0; i < OSFP_FIELD_MAX; i++) {
        field_score_db = &score_db->field_score_dbs[i];
        if (!field_score_db->enabled) {
            continue;
        }

        field = &fp->fields[i];
        if (!field->enabled) {
            continue;
        }

        os_class_score_matched = osfp_score_db_filed_match(field_score_db, field->value, field->value_len);
        if (os_class_score_matched == NULL) {
            continue;
        }

        importance = osfp_fingerprint_get_field_importance(i);

        for (j = 0; j < OSFP_OS_CLASS_MAX; j++) {
            entry_count = score_db->os_class_entry_count[j];
            tmp_score = os_class_score_matched->scores[j];
            if (entry_count == 0 || tmp_score == 0) {
                continue;
            }
            if (0 == flags || flags & OSFP_BIT_U32(j)) {
                printf("%s %s: ((%d * %u / %u) * %u ) / %u\n", osfp_fingerprint_get_field_name(i), osfp_os_class_id_to_name(j), OSFP_PERCENTILE, importance, perfect_score, os_class_score_matched->scores[j], entry_count);
                result_score->scores[j] += ((OSFP_PERCENTILE * importance / perfect_score) * os_class_score_matched->scores[j]) / entry_count;
            }
        }

        if (i == OSFP_FIELD_TCP_OPTIONS) {
            i++;
        }
    }

    return OSFP_NOERR;
exit:
    return ret;
}

void osfp_score_db_debug_print(struct osfp_score_db *score_db)
{
    int i;
    printf("score_db:\n");
    printf("entry_count: %u\n", score_db->entry_count);
    printf("perfect_score: %u\n", score_db->perfect_score);

    for (i = 0; i < OSFP_OS_CLASS_MAX; i++) {
        const char *name = osfp_os_class_id_to_name(i);
        printf("os class %s ", name);

        printf("entry_count: %u\n", score_db->os_class_entry_count[i]);
        printf("os class %s entry_count: %u\n", osfp_os_class_id_to_name(i), score_db->os_class_entry_count[i]);
    }

    for (i = 0; i < OSFP_FIELD_MAX; i++) {
        printf("field %s enabled: %u\n", osfp_fingerprint_get_field_name(i), score_db->field_score_dbs[i].enabled);
        printf("field %s type: %u\n", osfp_fingerprint_get_field_name(i), score_db->field_score_dbs[i].type);
        printf("field %s entry_count: %u\n", osfp_fingerprint_get_field_name(i), score_db->field_score_dbs[i].entry_count);
        printf("field %s enabled: %u\n", osfp_fingerprint_get_field_name(i), score_db->field_score_dbs[i].data);
    }
}

struct osfp_score_db *osfp_score_db_create(void)
{
    int i;
    struct osfp_score_db *score_db;
    struct osfp_field_score_db *db;

    score_db = calloc(1, sizeof(struct osfp_score_db));
    if (score_db == NULL) {
        goto exit;
    }

    for (i = 0; i < OSFP_FIELD_MAX; i++) {
        db = &score_db->field_score_dbs[i];

        db->enabled = osfp_fingerprint_get_field_enabled(i);
        if (!db->enabled) {
            osfp_log_warning("field disabled: %s", "");
            continue;
        }

        db->type = osfp_fingerprint_get_field_type(i);
        switch (db->type) {
            case OSFP_FIELD_TYPE_UINT:
                db->create = osfp_score_db_array_create;
                db->destroy = osfp_score_db_array_destroy;
                db->add = osfp_score_db_array_add;
                db->match = osfp_score_db_array_match;
                break;
            case OSFP_FIELD_TYPE_STRING:
                db->create = osfp_score_db_hash_create;
                db->destroy = osfp_score_db_hash_destroy;
                db->add = osfp_score_db_hash_add;
                db->match = osfp_score_db_hash_match;
                break;
            default:
                osfp_log_debug("unsupported type: %u", db->type);
                goto exit;
        }

        db->data = db->create();
        if (db->data == NULL) {
            osfp_log_debug("create failed. field: %s", osfp_fingerprint_get_field_name(i));
            goto exit;
        }
    }

    return score_db;
exit:
    if (score_db) {
        osfp_score_db_destroy(score_db);
    }
    return NULL;
}

void osfp_score_db_destroy(struct osfp_score_db *score_db)
{
    int i;
    struct osfp_field_score_db *db;

    if (score_db) {
        for (i = 0; i < OSFP_FIELD_MAX; i++) {
            db = &score_db->field_score_dbs[i];
            db->destroy(db->data);
            db->data = NULL;
        }
        free(score_db);
    }
}