#include "lcrypt_pgp_helpers.h"
#include <string.h>

#ifndef LCRYPT_HASH_BLOCK_SIZE
    #define LCRYPT_HASH_BLOCK_SIZE (64*1024)
#endif

static uint32_t get_uint32_be(const uint8_t *data) {
    return ((uint32_t)data[0] << 24) | ((uint32_t)data[1] << 16) | ((uint32_t)data[2] << 8) | data[3];
}

static uint16_t get_uint16_be(const uint8_t *data) {
    return ((uint16_t)data[0] << 8) | data[1];
}

#ifdef LCRYPT_CRYPT_ENGINE_GCRYPT
    t_lcrypt_err get_mpi(const uint8_t *msg, size_t size, t_lcrypt_pgp_mpi *mpi, size_t *parsed) {
        gcry_error_t err = gcry_mpi_scan(mpi, GCRYMPI_FMT_PGP, msg, size, parsed);
        return err == GPG_ERR_NO_ERROR ? 0 : LCRYPT_ERR_PGP_INVALID_FORMAT;
    }

    void free_mpi(t_lcrypt_pgp_mpi *mpi) {
        gcry_mpi_release(*mpi);
    }
#elif defined  LCRYPT_MPI_STATIC
    #define MPI_MSG_BYTES(mpi) ((mpi)->len + 2)

    static t_lcrypt_err get_mpi(const uint8_t *data, size_t size, t_lcrypt_pgp_mpi *mpi, size_t *parsed) {
        t_lcrypt_err err = 0;
        if (size < 2) {
            err = LCRYPT_ERR_PGP_INSUF_DATA;
        } else {
            uint16_t bits = get_uint16_be(data);
            mpi->len = (bits + 7)/8;
            if (mpi->len > LCRYPT_MPI_MAX_SIZE) {
                err = LCRYPT_ERR_PGP_MPI_UNSUP_SIZE;
            } else if (size < (size_t)MPI_MSG_BYTES(mpi)) {
                err = LCRYPT_ERR_PGP_INSUF_DATA;
            } else {
                memcpy(mpi->bytes, &data[2], mpi->len);
            }
            *parsed = MPI_MSG_BYTES(mpi);
        }
        return err;
    }

    void free_mpi(t_lcrypt_pgp_mpi *mpi) {

    }
#endif



t_lcrypt_err pgp_parse_packet_hdr(const uint8_t *msg, size_t size, t_lcrypt_pgp_hdr *hdr) {
    t_lcrypt_err err = size >= 2 ? 0 : LCRYPT_ERR_PGP_INSUF_DATA;
    uint8_t tag;

    if (!err) {
        tag = msg[0];
        if (!(tag & 0x80)) {
            err = LCRYPT_ERR_PGP_INVALID_FORMAT;
        }
    }

    if (!err) {
        uint8_t hdr_size;
        uint32_t pkt_size;
        uint8_t type = (tag & 0x3C) >> 2;

        if (!(tag & 0x40)) {
            switch (msg[0] & 0x3) {
                case 0:
                    hdr_size = 2;
                    if (size >= hdr_size) {
                        pkt_size = msg[1];
                    } else {
                        err = LCRYPT_ERR_PGP_INSUF_DATA;
                    }
                    break;
                case 1:
                    hdr_size = 3;
                    if (size >= hdr_size) {
                        pkt_size = get_uint16_be(&msg[1]);
                    } else {
                        err = LCRYPT_ERR_PGP_INSUF_DATA;
                    }
                    break;
                case 2:
                    hdr_size = 5;
                    if (size >= hdr_size) {
                       pkt_size = get_uint32_be(&msg[1]);
                    } else {
                        err = LCRYPT_ERR_PGP_INSUF_DATA;
                    }
                    break;
                default:
                    err = LCRYPT_ERR_PGP_UNSUP_FEATURE;
                    break;
            }
        } else {

            if (msg[1] < 192) {
                hdr_size = 2;
                pkt_size = msg[1];
            } else if (msg[1] < 224) {
                hdr_size = 3;
                if (size >= hdr_size) {
                    pkt_size = ((msg[1] - 192) << 8) + msg[2] + 192;
                } else {
                    err = LCRYPT_ERR_PGP_INSUF_DATA;
                }
            } else if (msg[1] == 255) {
                hdr_size = 6;
                if (size >= hdr_size) {
                    pkt_size = get_uint32_be(&msg[2]);
                } else {
                    err = LCRYPT_ERR_PGP_INSUF_DATA;
                }
            } else {
                err = LCRYPT_ERR_PGP_UNSUP_FEATURE;
            }
        }

        if (!err && (size < (hdr_size + pkt_size))) {
            err = LCRYPT_ERR_PGP_INSUF_DATA;
        }

        if (!err) {
            hdr->hdr_size = hdr_size;
            hdr->type = type;
            hdr->pkt_size = pkt_size;
        }
    }


    return err;
}


t_lcrypt_err pgp_public_key_parse(const uint8_t *msg, size_t size, t_lcrypt_pgp_pubkey *key) {
    t_lcrypt_err err = 0;
    uint8_t ver;
    if (size < 1) {
        err = LCRYPT_ERR_PGP_INSUF_DATA;
    } else {
        ver = msg[0];
        if (ver == 3) {
            key->creation_time = get_uint32_be(&msg[1]);

        } else if (ver == 4) {
            if (size < 6) {
                err = LCRYPT_ERR_PGP_INSUF_DATA;
            } else {
                unsigned pos = 6;
                size_t parsed;
                key->creation_time = get_uint32_be(&msg[1]);
                key->algo = msg[5];
                switch (key->algo) {
                    case PGP_PUBKEY_ALGO_RSA:
                    case PGP_PUBKEY_ALGO_RSA_ENCRYPT_ONLY:
                    case PGP_PUBKEY_ALGO_RSA_SIGN_ONLY:

                        err = get_mpi(&msg[pos], size-pos, &key->rsa.n, &parsed);
                        if (!err) {
                            pos += parsed;
                            err = get_mpi(&msg[pos], size-pos, &key->rsa.e, &parsed);
                            if (!err) {
                                pos += parsed;
                            }
                        }
                        break;
                    default:
                        err = LCRYPT_ERR_PGP_UNSUP_ALGO;
                        break;
                }
            }
        } else {
            err = LCRYPT_ERR_PGP_INVALID_FORMAT;
        }
    }

    return err;
}

void pgp_public_key_release(t_lcrypt_pgp_pubkey *key) {
    switch (key->algo) {
        case PGP_PUBKEY_ALGO_RSA:
        case PGP_PUBKEY_ALGO_RSA_ENCRYPT_ONLY:
        case PGP_PUBKEY_ALGO_RSA_SIGN_ONLY:
            free_mpi(&key->rsa.n);
            free_mpi(&key->rsa.e);
            break;
        default:
            break;
    }
}

static t_lcrypt_err pgp_sign_algo_parse(const uint8_t *msg, size_t size,  t_lcrypt_pgp_sign *sign) {
    t_lcrypt_err err = 0;
    size_t parsed;
    switch (sign->pubkey_algo) {
        case PGP_PUBKEY_ALGO_RSA:
        case PGP_PUBKEY_ALGO_RSA_ENCRYPT_ONLY:
        case PGP_PUBKEY_ALGO_RSA_SIGN_ONLY:
            err = get_mpi(msg, size, &sign->rsa.val, &parsed);
            break;
        default:
            err = LCRYPT_ERR_PGP_UNSUP_ALGO;
            break;
    }
    return err;
}

t_lcrypt_err pgp_sign_parse(const uint8_t *msg, size_t size,  t_lcrypt_pgp_sign *sign) {
    t_lcrypt_err err = 0;
    if (size < 1) {
        err = LCRYPT_ERR_PGP_INSUF_DATA;
    } else {
        sign->version = msg[0];
        if (sign->version == 3) {
            if (size < 18) {
                err = LCRYPT_ERR_PGP_INSUF_DATA;
            } else {
                sign->hashed_data_len = msg[1];
                if (sign->hashed_data_len != 5) {
                    err = LCRYPT_ERR_PGP_INVALID_FORMAT;
                } else {
                    sign->hashed_data = malloc(sign->hashed_data_len);
                    if (sign->hashed_data == NULL) {
                        err = LCRYPT_ERR_PGP_MEMORY_ALLOC;
                    }
                }
            }

            if (!err) {
                memcpy(sign->hashed_data, &msg[2], 5);
                memcpy(sign->key_id, &msg[7], 8);
                sign->pubkey_algo = msg[15];
                sign->hash_algo = msg[16];
                sign->first_word = get_uint16_be(&msg[17]);
                err = pgp_sign_algo_parse(&msg[19], size-19, sign);
            }
        } else if (sign->version == 4) {
            unsigned pos;
            if (size < 5) {
                err = LCRYPT_ERR_PGP_INSUF_DATA;
            } else {
                sign->sign_type = msg[1];
                sign->pubkey_algo = msg[2];
                sign->hash_algo = msg[3];
                sign->hashed_data_len = 6 + get_uint16_be(&msg[4]);
                if (size < (size_t)(sign->hashed_data_len + 2)) {
                    err = LCRYPT_ERR_PGP_INSUF_DATA;
                } else {
                    if (sign->hashed_data_len) {
                        sign->hashed_data = malloc(sign->hashed_data_len);
                        if (sign->hashed_data == NULL) {
                            err = LCRYPT_ERR_PGP_MEMORY_ALLOC;
                        } else {
                            memcpy(sign->hashed_data, msg, sign->hashed_data_len);
                        }
                    } else {
                        sign->hashed_data = NULL;
                    }

                    if (!err) {
                        pos = sign->hashed_data_len;
                    }
                }
            }

            if (!err) {
                uint16_t unheshed_subpacket_len = get_uint16_be(&msg[pos]);
                if (size < (pos + 2 + unheshed_subpacket_len + 2)) {
                    err = LCRYPT_ERR_PGP_INSUF_DATA;
                } else {
                    pos += (2 + unheshed_subpacket_len);
                }
            }

            if (!err) {
                sign->first_word = get_uint16_be(&msg[pos]);
                err = pgp_sign_algo_parse(&msg[pos+2], size-(pos+2), sign);
            }
        }
    }
    return err;
}

void pgp_sign_release(t_lcrypt_pgp_sign *sign) {
    switch (sign->pubkey_algo) {
        case PGP_PUBKEY_ALGO_RSA:
        case PGP_PUBKEY_ALGO_RSA_ENCRYPT_ONLY:
        case PGP_PUBKEY_ALGO_RSA_SIGN_ONLY:
            free_mpi(&sign->rsa.val);
            free(sign->hashed_data);
            break;
        default:
            break;
    }
}

#ifndef LCRYPT_CRYPT_ENGINE_NONE
t_lcrypt_err pgp_sign_check(const t_lcrypt_pgp_sign *sign, const t_lcrypt_pgp_pubkey *key,
                                   const uint8_t *data, size_t data_size, t_lcrypt_progress_cb cb, void *cb_data) {
    t_lcrypt_err err = 0;
    int gcry_hash_algo;
    const char *hash_algo_str;
    const unsigned char *hash = NULL;
    switch (sign->hash_algo) {
        case PGP_HASH_ALGO_SHA256:
            gcry_hash_algo = GCRY_MD_SHA256;
            hash_algo_str = "sha256";
            break;
        default:
            err = LCRYPT_ERR_PGP_UNSUP_ALGO;
            break;
    }

    if (!err) {
        gcry_md_hd_t hmd;
        gcry_error_t gerr = gcry_md_open(&hmd, gcry_hash_algo, 0);
        if (gerr != GPG_ERR_NO_ERROR) {
            err = LCRYPT_ERR_PGP_INTERNAL;
        }


        if (!err) {
            size_t done_size = 0;
            while (done_size != data_size) {
                size_t block_size = data_size - done_size;
                if (block_size > LCRYPT_HASH_BLOCK_SIZE)
                    block_size = LCRYPT_HASH_BLOCK_SIZE;
                gcry_md_write(hmd, &data[done_size], block_size);
                done_size += block_size;
                if (cb != NULL) {
                    cb(cb_data, done_size, data_size + 1024);
                }
            }
            gcry_md_write(hmd, sign->hashed_data, sign->hashed_data_len);
            if (sign->version == 4) {
                gcry_md_putc(hmd, sign->version);
                gcry_md_putc(hmd, 0xFF);
                gcry_md_putc(hmd, 0);
                gcry_md_putc(hmd, 0);
                gcry_md_putc(hmd, (sign->hashed_data_len >> 8));
                gcry_md_putc(hmd, (sign->hashed_data_len & 0xFF));
            }

            hash = gcry_md_read(hmd, 0);
            if ((hash[0] != ((sign->first_word >> 8) & 0xFF)) ||
                    (hash[1] != (sign->first_word & 0xFF))) {
                err = LCRYPT_ERR_PGP_BAD_SIGN;
            }
        }

        if (!err) {
            gcry_sexp_t key_exp = 0, data_exp = 0, sign_exp = 0;
            gcry_mpi_t  data_mpi = 0;
            size_t err_offs;
            gcry_error_t gerr;
            switch (key->algo) {
                case PGP_PUBKEY_ALGO_RSA:
                case PGP_PUBKEY_ALGO_RSA_ENCRYPT_ONLY:
                case PGP_PUBKEY_ALGO_RSA_SIGN_ONLY:
                    gerr = gcry_sexp_build(&key_exp, &err_offs,
                                           "(public-key (rsa (n %M) (e %M)))",
                                            key->rsa.n, key->rsa.e);
                    if (gerr == GPG_ERR_NO_ERROR) {
                        gerr = gcry_sexp_build(&sign_exp, NULL,
                                   "(sig-val (rsa (s %M)))", sign->rsa.val);
                    }
                    break;
                default:
                    err = LCRYPT_ERR_PGP_UNSUP_ALGO;
                    break;
            }

            if (!err) {
                if (gerr == GPG_ERR_NO_ERROR)  {
                    gerr = gcry_mpi_scan(&data_mpi, GCRYMPI_FMT_USG, hash, 256/8, NULL);
                }

                if (gerr == GPG_ERR_NO_ERROR) {
                    gerr = gcry_sexp_build(&data_exp, NULL,
                                       "(data (flags pkcs1) (hash %s %M))",
                                       hash_algo_str, data_mpi);
                }

                if (gerr == GPG_ERR_NO_ERROR) {
                    gerr = gcry_pk_verify(sign_exp, data_exp, key_exp);
                    if (gerr != GPG_ERR_NO_ERROR) {
                        err = LCRYPT_ERR_PGP_BAD_SIGN;
                    }
                }
            }

            gcry_mpi_release(data_mpi);
            gcry_sexp_release(key_exp);
            gcry_sexp_release(data_exp);
            gcry_sexp_release(sign_exp);
        }

        gcry_md_close(hmd);
    }
    return err;
}
#endif

t_lcrypt_err pgp_data_parse(const uint8_t *msg, size_t size, t_lcrypt_pgp_data *data) {
    t_lcrypt_err err  = 0;
    size_t pos = 0;
    if (size < 2) {
        err = LCRYPT_ERR_PGP_INVALID_FORMAT;
    } else {
        data->type = msg[0];
        data->filename_len = msg[1];
        pos += 2;
    }

    if (!err) {
        if (size < (pos + data->filename_len)) {
            err = LCRYPT_ERR_PGP_INVALID_FORMAT;
        } else {
            data->filename = (const char*)&msg[pos];
            pos += data->filename_len;
        }
    }

    if (!err) {
        if (size < (pos + 4)) {
            err = LCRYPT_ERR_PGP_INVALID_FORMAT;
        } else {
            data->creation_date = get_uint32_be(&msg[pos]);
            pos += 4;
        }
    }

    if (!err) {
        data->body = &msg[pos];
        data->body_len = size - pos;
    }
    return err;
}
