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

#include "firm_check.h"
#include "burn_intf.h"
#include "lcrypt/sha.h"
#include "lcrypt/lcrypt.h"
#include "lcrypt_pgp/lcrypt_pgp_helpers.h"
#include "radix64.h"

#include <libxml/parser.h>
#include <libxml/tree.h>

#define LFRIM_VERSION_SIZE      32
#define LFIRM_DEV_NAME_SIZE     32
#define LFIRM_DEV_CHIPNAME_SIZE 32
#define LFIRM_DEV_REVISION      32


struct st_lxfw_info {
    xmlDoc *doc; /* the resulting document tree */
    xmlNode *root;
    char* msg;
    const char *devname;
    size_t data_size;
    size_t msg_size;
};


/* Получаем текст из XML-узла */
static xmlChar* f_get_node_text(xmlNode* node) {
    xmlChar* ret = NULL;
    for (node=node->children; node; node=node->next) {
        if ((node->type == XML_TEXT_NODE) && !xmlIsBlankNode(node)) {
            ret = node->content;
        }
    }
    return ret;
}

void lxfw_clean(t_lxfw_info_hnd hinfo) {
    if (hinfo!=NULL) {
        if (hinfo->doc) {
            xmlFreeDoc(hinfo->doc);
        }
        free(hinfo);
    }
}

void lxfw_cleanup(void) {
    /*
     * Cleanup function for the XML library.
     */
    xmlCleanupParser();
    /*
     * this is to debug memory for regression tests
     */
    xmlMemoryDump();
}


t_lxfw_errs lxfw_init(const char *filename, t_lxfw_info_hnd *hinfo) {
    t_lxfw_errs err = 0;
    t_lxfw_info_hnd pinfo = calloc(sizeof(t_lxfw_info), 1);
    xmlNode *child;

    if (pinfo==NULL) {
        err = LXFW_ERR_MEMORY_ALLOC;
    } else {
        pinfo->doc = NULL;
        pinfo->root = NULL;
    }


    if (!err) {
        LIBXML_TEST_VERSION;
        /*parse the file and get the DOM */
        pinfo->doc = xmlReadFile(filename, NULL, 0);
        if (pinfo->doc == NULL)
            err = LXFW_ERR_FIRM_PARSE;
    }


    /* получаем root-элемент и проверяем его имя */
    if (!err) {
        pinfo->root = xmlDocGetRootElement(pinfo->doc);
        if (strcmp((char*)pinfo->root->name, "firmware"))
            err = LXFW_ERR_FIRM_PARSE_ROOT_ELEM;
    }



    if (!err) {
        for (child = pinfo->root->children; child; child = child->next) {
            if (child->type == XML_ELEMENT_NODE) {
                if (!strcmp((char*)child->name, "data")) {
                    pinfo->msg = (char*)f_get_node_text(child);
                    pinfo->msg_size = strlen(pinfo->msg);
                    if ((pinfo->msg_size==0) || (pinfo->msg_size %4)) {
                        err = LXFW_ERR_FIRM_PARSE_NO_DATA;
                    } else {
                        pinfo->data_size = pinfo->msg_size*3/4;
                        if (pinfo->msg[pinfo->msg_size-1]=='=')
                            pinfo->data_size--;
                        if (pinfo->msg[pinfo->msg_size-2]=='=')
                            pinfo->data_size--;
                    }
                } else if (!strcmp((char*)child->name, "devices")) {
                    xmlNode *devnode = child->children;
                    for (; (devnode != NULL) && (pinfo->devname == NULL); devnode = devnode->next) {
                        if ((devnode->type == XML_ELEMENT_NODE) && !strcmp((char*)devnode->name, "device")) {
                            xmlAttr* attr;
                            for (attr = devnode->properties; (attr != NULL) &&
                                 (pinfo->devname == NULL); attr=attr->next) {

                                const char *val =  (char*)attr->children->content;
                                if (!strcmp((char*)attr->name, "name")) {
                                    pinfo->devname = val;
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    if (!err && (pinfo->msg == NULL))
        err = LXFW_ERR_FIRM_PARSE_NO_DATA;

    if (!err && (pinfo->devname == NULL))
        err = LXFW_ERR_FIRM_PARSE_NO_DEVS;

    if (err && pinfo) {
        lxfw_clean(pinfo);
    }

    if (!err)
        *hinfo = pinfo;


    return err;
}


const char *lxfw_get_property(t_lxfw_info_hnd hinfo, const char *name) {
    const char *retval = NULL;
    xmlAttr* attr;
    for (attr = hinfo->root->properties; (attr != NULL) && (retval == NULL); attr=attr->next) {
        if (!strcmp((char*)attr->name, name)) {
            retval = (char*)attr->children->content;
        }
    }
    return retval;
}


const char *lxfw_get_version(t_lxfw_info_hnd hinfo) {
    return lxfw_get_property(hinfo, "version");
}

t_lxfw_errs lxfw_get_size(t_lxfw_info_hnd hinfo, unsigned *size) {
    *size = hinfo->data_size;
    return 0;
}

t_lxfw_errs lxfw_get_wr_sign(t_lxfw_info_hnd hinfo, uint32_t *sign) {
    t_lxfw_errs err = 0;
    const char *sign_str = lxfw_get_property(hinfo, "wr_sign");
    if (sign_str != NULL) {
        if (sscanf(sign_str, "%i", sign)!=1) {
            err = LXFW_ERR_INVALID_PROPERTY_VALUE;
        }
    } else {
        err = LXFW_ERR_NO_PROPERTY;
    }
    return err;
}

t_lxfw_errs lxfw_get_crc32(t_lxfw_info_hnd hinfo, uint32_t *sign) {
    t_lxfw_errs err = 0;
    const char *sign_str = lxfw_get_property(hinfo, "crc32");
    if (sign_str != NULL) {
        if (sscanf(sign_str, "%i", sign)!=1) {
            err = LXFW_ERR_INVALID_PROPERTY_VALUE;
        }
    } else {
        err = LXFW_ERR_NO_PROPERTY;
    }
    return err;
}


t_lxfw_errs firm_check_device(t_lxfw_info_hnd hinfo, const char *devname, const char *chip, const char *rev) {
    int fnd = 0;
    xmlNode *child = hinfo->root->children;

    for (; !fnd && child; child = child->next) {
        if ((child->type == XML_ELEMENT_NODE) && !strcmp((char*)child->name, "devices")) {
            xmlNode *devnode = child->children;
            for (; devnode; devnode = devnode->next) {
                if ((devnode->type == XML_ELEMENT_NODE) && !strcmp((char*)devnode->name, "device")) {
                    int ok_name=0, ok_chip=chip==NULL, ok_rev = rev==NULL;
                    xmlAttr* attr;
                    for (attr = devnode->properties; attr; attr=attr->next) {
                        char *val =  (char*)attr->children->content;
                        if (!strcmp((char*)attr->name, "name")) {
                            if (!strcmp(val, devname))
                                ok_name = 1;
                        } else if ((chip!=NULL) && !strcmp((char*)attr->name, "chip")) {
                            if (!strcmp(chip, val))
                                ok_chip = 1;
                        } else if ((rev!=NULL) && !strcmp((char*)attr->name, "revision")) {
                            if (!strcmp(rev, val))
                                ok_rev = 1;
                        }
                    }

                    if  (ok_name && ok_chip && ok_rev)
                        fnd = 1;
                }
            }
        }
    }

    return fnd ? 0 : LXFW_ERR_DEVICE_CHECK;
}



int lxfw_get_data(t_lxfw_info_hnd hinfo, unsigned char *buf, size_t pos, size_t size) {
    size_t first_pos = (pos/3)*4, first_offs = (pos*4)%3;
    size_t last_offs, cont_size, buf_pos=0;
    uint8_t tmp[3];

    if (pos >= hinfo->data_size)
        return 0;

    if ((pos+size) > hinfo->data_size)
        size = hinfo->data_size-pos;

    cont_size = size;
    if (first_offs != 0)
        cont_size -= (3-first_offs);
    last_offs = (cont_size*4)%3;
    cont_size -= last_offs;


    if (first_offs != 0) {
        Radix64_Decode(&hinfo->msg[first_pos], 4, tmp);
        memcpy(buf, &tmp[first_offs], 3-first_offs);
        buf_pos+=(3-first_offs);
        first_pos+=4;
    }

    Radix64_Decode(&hinfo->msg[first_pos], cont_size*4/3, &buf[buf_pos]);
    buf_pos += cont_size;
    first_pos+=(cont_size*4/3);

    if (last_offs != 0) {
        Radix64_Decode(&hinfo->msg[first_pos], 4, tmp);
        memcpy(&buf[buf_pos], &tmp[0], last_offs);
    }
    return size;
}





/* проверка подписи прошивки */
t_lxfw_errs lxfw_check_sign(const char *filename) {
    FILE *f=NULL, *f_sig=NULL;
    char *signame = NULL;
    t_lxfw_errs err = 0;


    signame = malloc(strlen(filename)+5);
    if (signame==NULL)
        err = LXFW_ERR_MEMORY_ALLOC;

    if (!err) {
        f = fopen(filename, "rb");
        if (f==NULL) {
            err = LXFW_ERR_FIRM_FILE_OPEN;
        }
    }

    if (!err) {
        strcpy(signame, filename);
        strcat(signame, ".sig");
        f_sig = fopen(signame, "rb");
        if (f_sig==NULL) {
            err = LXFW_ERR_SIGN_FILE_OPEN;
        }
    }

    if (!err) {
        SHA256Context sha256;
        static uint8_t f_buf[1024];
        size_t sig_size;
        t_lcrypt_pgp_hdr hdr;
        t_lcrypt_pgp_sign sign_hdr;
        int sign_free_req = 0;
        uint8_t hash[SHA256HashSize];
        SHA256Reset(&sha256);

        do {

            size_t rd_size = fread(f_buf, 1, sizeof(f_buf), f);
            if (ferror(f)) {
                err = LXFW_ERR_FILE_READ;
            } else if (rd_size != 0) {
                SHA256Input(&sha256, f_buf, rd_size);
            }
        } while (!err && !feof(f));


        if (!err) {
            sig_size = fread(f_buf, 1, sizeof(f_buf), f_sig);
            if (ferror(f_sig)) {
                err = LXFW_ERR_FILE_READ;
            } else if (!feof(f_sig)) {
                err = LXFW_ERR_FILE_READ;
            }
        }

        if (!err) {
            t_lcrypt_err crypt_err = pgp_parse_packet_hdr(f_buf, sig_size, &hdr);
            if (!crypt_err) {
                if (hdr.type == PGP_PACKET_SIGNATURE) {
                    crypt_err = pgp_sign_parse(&f_buf[hdr.hdr_size], hdr.pkt_size, &sign_hdr);
                    sign_free_req = 1;
                } else {
                    crypt_err = LCRYPT_ERR_PGP_INVALID_FORMAT;
                }
            }

            if (crypt_err) {
                err = LXFW_ERR_INVALID_SIGNATURE;
            }
        }

        if (!err) {
            int i;
            static long_t sign[2*SIGN_SIZE + 10], dec_sign[2*SIGN_SIZE + 10], n_buf[2*RSA_N_SIZE+10];
            static long_t buf1[2*SIGN_SIZE + 10], buf2[2*SIGN_SIZE + 10];

            SHA256Input(&sha256, sign_hdr.hashed_data, sign_hdr.hashed_data_len);
            if (sign_hdr.version == 4) {
                uint8_t bytes[] = {sign_hdr.version, 0xFF, 0, 0, (sign_hdr.hashed_data_len >> 8), (sign_hdr.hashed_data_len & 0xFF)};
                SHA256Input(&sha256, bytes, sizeof(bytes));
            }
            SHA256Result(&sha256, hash);

            //проверка корректности хеша по 2 байтам из пакета
            if ((hash[0] != ((sign_hdr.first_word >> 8) & 0xFF)) ||
                    (hash[1] != (sign_hdr.first_word & 0xFF))) {
                err = LXFW_ERR_INVALID_SIGNATURE;
            }

            if (!err) {
                memset(sign, 0, sizeof(sign));
                for (i = 0; i < sign_hdr.rsa.val.len; i++) {
                    sign[i] = sign_hdr.rsa.val.bytes[sign_hdr.rsa.val.len-i-1];
                }
                //проверяем подпись по rsa
                memcpy(n_buf, rsa_n, RSA_N_SIZE);

                vlong_fast_spow_mod(sign, SIGN_SIZE/sizeof(long_t), rsa_e, n_buf, RSA_N_SIZE/sizeof(long_t),
                                    dec_sign, buf1, buf2);

                //проверка результата - сравниваем с нужным форматом
                for (i=0; (i < SHA256HashSize)&&!err; i++) {
                    if (dec_sign[SHA256HashSize - 1 - i]!=hash[i]) {
                        err = LXFW_ERR_INVALID_SIGNATURE;
                    }
                }

                for (i=SHA256HashSize; (i< (SHA256HashSize + SHA256_ASN_SIZE))&&!err; i++) {
                    if (dec_sign[SHA256_ASN_SIZE+SHA256HashSize+SHA256HashSize-1-i]!=SHA256_ASN[i- SHA256HashSize]) {
                        err = LXFW_ERR_INVALID_SIGNATURE;
                    }
                }

                if ((dec_sign[SHA256_ASN_SIZE+SHA256HashSize] != 0x0) ||
                        (dec_sign[SIGN_SIZE-1] != 0x0) ||
                        (dec_sign[SIGN_SIZE-2] != 0x1)) {
                    err = LXFW_ERR_INVALID_SIGNATURE;
                }

                for (i=SHA256HashSize + SHA256_ASN_SIZE + 1; (i< (SIGN_SIZE -2))&&!err; i++) {
                    if (dec_sign[i]!=0xFF) {
                        err = LXFW_ERR_INVALID_SIGNATURE;
                    }
                }
            }
        }

        if (sign_free_req) {
            pgp_sign_release(&sign_hdr);
        }
    }

    if (f)
        fclose(f);
    if (f_sig)
        fclose(f_sig);
    free(signame);

    return err;
}


typedef struct {
    int32_t err;
    const char *str;
} t_err_table;

static const char* f_unknow_err = "Неизвестная ошибка";

static const t_err_table f_err_tbl[] = {
    {LXFW_ERR_INVALID_PARAM,        "Ошибка входных параметров"},
    {LXFW_ERR_FIRM_FILE_OPEN,       "Не удалось открыть файл с прошивкой"},
    {LXFW_ERR_SIGN_FILE_OPEN,       "Не удалось открыть файл с подписью"},
    {LXFW_ERR_FILE_READ,            "Ошибка чтения файла"},
    {LXFW_ERR_INVALID_BOOT_FLAG,    "Неверное значение флага загрузки в модуле"},
    {LXFW_ERR_DEVICE_CHECK,         "Прошивка не соответствует указанному устройству"},
    {LXFW_ERR_WRITE_FILE,           "Ошибка записи в файл"},
    {LXFW_ERR_DEVICE_OPEN,          "Ошибка открытия устройства"},
    {LXFW_ERR_INVALID_INTERFACE,    "Выбран неизвестный интерфейс"},
    {LXFW_ERR_MEMORY_ALLOC,         "Ошибка выделения памяти"},
    {LXFW_ERR_INVALID_SIGNATURE,    "Неверная подпись прошивки"},
    {LXFW_ERR_FIRM_PARSE,           "Ошибка формата файла прошивки"},
    {LXFW_ERR_FIRM_PARSE_ROOT_ELEM, "Файл не является прошивкой (неверный тип элемента)"},
    {LXFW_ERR_FIRM_PARSE_NO_DATA,   "В файле не найдено содержимое прошивки"},
    {LXFW_ERR_FIRM_PARSE_NO_DEVS,   "В прошивке не указано ни одного устройства, для которой она предназначена"},
    {LXFW_ERR_FIRM_OLD_VERSION,     "Версия прошивки более старая, чем записана в устройтсве"},
    {LXFW_ERR_UNSUPPORTED_DEVICE,   "Неподдерживаемое имя устройства"},
    {LXFW_ERR_NO_PROPERTY,          "Указанное свойство не найдено"},
    {LXFW_ERR_INVALID_PROPERTY_VALUE,"Неверное значение совойства"}
};


const char *lxfw_get_err_str(t_lxfw_errs err_code) {
    size_t i;
    const char* str = f_unknow_err;

    for (i=0; (i < sizeof(f_err_tbl)/sizeof(f_err_tbl[0])) && (str==f_unknow_err); i++) {
        if (f_err_tbl[i].err == err_code)
            str = f_err_tbl[i].str;
    }
    return str;
}


const char *lxfw_get_devname(t_lxfw_info_hnd hinfo) {
    return hinfo->devname;
}
