/* -*- c-basic-offset: 8 -*-
   rdesktop: A Remote Desktop Protocol client.
   ASN.1 utility functions
   Copyright 2012-2017 Henrik Andersson <hean01@cendio.se> for Cendio AB
   Copyright 2017 Alexander Zakharov <uglym8@gmail.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include <gnutls/gnutls.h>
#include <libtasn1.h>
#include <stdlib.h>

#include "rdesktop.h"
#include "asn.h"

// Generated by asn1Parser
#include "pkix_asn1_tab.c"

static asn1_node *asn_defs = NULL;

#define MAX_ERROR_DESCRIPTION_SIZE 1024
static char errstr[MAX_ERROR_DESCRIPTION_SIZE];

/* Parse an ASN.1 BER header */
static int ber_parse_header(STREAM s, int tagval, uint32 *length)
{
    int tag, len;

    if (tagval > 0xff) {
        if (!s_check_rem(s, 2)) {
            return ASN_PARSE_ERROR;
        }
        in_uint16_be(s, tag);
    } else {
        if (!s_check_rem(s, 1)) {
            return ASN_PARSE_ERROR;
        }
        in_uint8(s, tag);
    }

    if (tag != tagval) {
        logger(Core, Error, "ber_parse_header(), expected tag %d, got %d", tagval, tag);
        return ASN_TAG_ERROR;
    }

    if (!s_check_rem(s, 1)) {
        return ASN_PARSE_ERROR;
    }
    in_uint8(s, len);

    if (len & 0x80) {
        len &= ~0x80;
        if (!s_check_rem(s, len)) {
            return ASN_PARSE_ERROR;
        }
        *length = 0;
        while (len--) {
            next_be(s, *length);
        }
    } else {
        *length = len;
    }

    return ASN_SUCCESS;
}

static void ber_out_sequence(STREAM out, STREAM content)
{
    size_t length;
    length = (content ? s_length(content) : 0);
    ber_out_header(out, BER_TAG_SEQUENCE | BER_TAG_CONSTRUCTED, length);
    if (content) {
        out_stream(out, content);
    }
}

/* Output an ASN.1 BER header */
static void ber_out_header(STREAM s, int tagval, int length)
{
    if (tagval > 0xff) {
        out_uint16_be(s, tagval);
    } else {
        out_uint8(s, tagval);
    }

    if (length >= 0x80) {
        out_uint8(s, 0x82);
        out_uint16_be(s, length);
    } else {
        out_uint8(s, length);
    }
}

/* Output an ASN.1 BER integer */
static void ber_out_integer(STREAM s, int value)
{
    ber_out_header(s, BER_TAG_INTEGER, 2);
    out_uint16_be(s, value);
}

static int ber_in_header(STREAM s, int *tagval, int *decoded_len)
{
    in_uint8(s, *tagval);
    in_uint8(s, *decoded_len);

    if (*decoded_len < 0x80) {
        return ASN_SUCCESS;
    } else if (*decoded_len == 0x81) {
        in_uint8(s, *decoded_len);
        return ASN_SUCCESS;
    } else if (*decoded_len == 0x82) {
        in_uint16_be(s, *decoded_len);
        return ASN_SUCCESS;
    }

    return ASN_PARSE_ERROR;
}

static int init_asn1_lib(void)
{
    int asn1_rv;

    if (asn_defs) {
        return ASN_SUCCESS;
    }

    asn_defs = malloc(sizeof(*asn_defs));

    if (!asn_defs) {
        logger(Core, Error, "%s:%s:%d Failed to allocate memory for ASN.1 parser\n",
               __FILE__, __func__, __LINE__);
        return ASN_MEM_ERROR;
    }

    *asn_defs = NULL;

    if (ASN1_SUCCESS != (asn1_rv = asn1_array2tree(pkix_asn1_tab, asn_defs, errstr))) {
        logger(Core, Error, "%s:%s:%d Failed to init ASN.1 parser. Error = 0x%x (%s)\n",
               __FILE__, __func__, __LINE__, asn1_rv, asn1_strerror(asn1_rv));

        free(asn_defs);
        asn_defs = NULL;
        return asn1_rv;
    }

    return ASN_SUCCESS;
}

/* Encode RSA public key into DER PKCS#1 */
static int write_pkcs1_der_pubkey(const gnutls_datum_t *m, const gnutls_datum_t *e, uint8_t *out, int *out_len)
{
    int asn1_rv;
    asn1_node asn_cert;

    if (!asn_defs) {
        if (init_asn1_lib() != ASN_SUCCESS) {
            return ASN_INIT_ERROR;
        }
    }

    if (ASN1_SUCCESS != (asn1_rv = asn1_create_element(*asn_defs, "PKIX1Implicit88.RSAPublicKey", &asn_cert))) {
        logger(Core, Error, "%s:%s:%d Failed to create ASN.1 parser element for RSAPublicKey. Error = 0x%x (%s)\n",
               __FILE__, __func__, __LINE__, asn1_rv, asn1_strerror(asn1_rv));
        return ASN_CREATE_ERROR;
    }

    if (ASN1_SUCCESS != (asn1_rv = asn1_write_value(asn_cert, "modulus", m->data, m->size))) {
        logger(Core, Error, "%s:%s:%d Failed to write modulus. Error = 0x%x (%s)\n",
               __FILE__, __func__, __LINE__, asn1_rv, asn1_strerror(asn1_rv));

        return ASN_WRITE_ERROR;
    }

    if (ASN1_SUCCESS != (asn1_rv = asn1_write_value(asn_cert, "publicExponent", e->data, e->size))) {
        logger(Core, Error, "%s:%s:%d Failed to write publicExponent. Error = 0x%x (%s)\n",
               __FILE__, __func__, __LINE__, asn1_rv, asn1_strerror(asn1_rv));
        return ASN_WRITE_ERROR;
    }

    if (ASN1_SUCCESS != (asn1_rv = asn1_der_coding(asn_cert, "", out, out_len, errstr))) {
        logger(Core, Error, "%s:%s:%d Failed to encode PKIX1Implicit88.RSAPublicKey. Error = 0x%x (%s)\n",
               __FILE__, __func__, __LINE__, asn1_rv, asn1_strerror(asn1_rv));

        return ASN_ENCODE_ERROR;
    }

    return ASN_SUCCESS;
}

static int libtasn_read_cert_pk(uint8_t *data, size_t len, char *oid, size_t *oid_size, gnutls_datum_t *m, gnutls_datum_t *e)
{
    int asn1_rv;
    asn1_node asn_cert;

    /* Parse DER encoded x.509 certificate */
    if (!asn_defs) {
        if (init_asn1_lib() != ASN_SUCCESS) {
            return ASN_INIT_ERROR;
        }
    }

    if (ASN1_SUCCESS != (asn1_rv = asn1_create_element(*asn_defs, "PKIX1Implicit88.Certificate", &asn_cert))) {
        logger(Core, Error, "%s:%s:%d Failed to create ASN.1 parser element. Error = 0x%x (%s)\n",
               __FILE__, __func__, __LINE__, asn1_rv, asn1_strerror(asn1_rv));
        return ASN_CREATE_ERROR;
    }

    if (ASN1_SUCCESS != (asn1_rv = asn1_der_decoding(&asn_cert, data, len, errstr))) {
        logger(Core, Error, "%s:%s:%d Failed to decode certificate object. Error = 0x%x (%s)\n",
               __FILE__, __func__, __LINE__, asn1_rv, asn1_strerror(asn1_rv));
        return ASN_DECODE_ERROR;
    }

    if (oid && oid_size) {
        if (ASN1_SUCCESS != (asn1_rv = asn1_read_value(asn_cert, "tbsCertificate.subjectPublicKeyInfo.algorithm.algorithm",
                                                       oid, (int *)oid_size)))
        {
            logger(Core, Error, "%s:%s:%d Failed to get cert's public key algorithm. Error = 0x%x (%s)\n",
                   __FILE__, __func__, __LINE__, asn1_rv, asn1_strerror(asn1_rv));
            return ASN_READ_ERROR;
        }
    }

    if (m && e) {
        int buflen;
        uint8_t buf[16384];

        asn1_node asn_key;
        int nblen;
        uint8_t newbuf[16384];

        buflen = sizeof(buf) - 1;
        if (ASN1_SUCCESS != (asn1_rv = asn1_read_value(asn_cert, "tbsCertificate.subjectPublicKeyInfo.subjectPublicKey", buf, &buflen))) {
            logger(Core, Error, "%s:%s:%d Failed to get cert's public key. Error = 0x%x (%s)\n",
                   __FILE__, __func__, __LINE__, asn1_rv, asn1_strerror(asn1_rv));
            return ASN_READ_ERROR;
        }

        if (ASN1_SUCCESS != (asn1_rv = asn1_create_element(*asn_defs, "PKIX1Implicit88.RSAPublicKey", &asn_key))) {
            logger(Core, Error, "%s:%s:%d Failed to create ASN.1 parser element. Error = 0x%x (%s)\n",
                   __FILE__, __func__, __LINE__, asn1_rv, asn1_strerror(asn1_rv));
            return ASN_CREATE_ERROR;
        }

        // As it' a BIT STRING the len constitutes the number of BITS, not BYTES
        if (ASN1_SUCCESS != (asn1_rv = asn1_der_decoding(&asn_key, buf, buflen / 8, errstr))) {
            logger(Core, Error, "%s:%s:%d Failed to decode public key  object. Error = 0x%x (%s)\n",
                   __FILE__, __func__, __LINE__, asn1_rv, asn1_strerror(asn1_rv));
            return ASN_DECODE_ERROR;
        }

        /* Get RSA public key's modulus and exponent */
        nblen = sizeof(newbuf);

        if (ASN1_SUCCESS != (asn1_rv = asn1_read_value(asn_key, "modulus", newbuf, &nblen))) {
            logger(Core, Error, "%s:%s:%d Failed to get RSA public key's modulus. Error = 0x%x (%s)\n",
                   __FILE__, __func__, __LINE__, asn1_rv, asn1_strerror(asn1_rv));
            return ASN_READ_ERROR;
        }

        m->size = nblen;

        m->data = malloc(m->size);
        if (!m->data) {
            logger(Core, Error, "%s:%s:%d Failed to allocate memory for modulus.\n", __FILE__, __func__, __LINE__);
            return ASN_MEM_ERROR;
        }

        memcpy((void *)m->data, newbuf, m->size);

        nblen = sizeof(newbuf);

        if (ASN1_SUCCESS != (asn1_rv = asn1_read_value(asn_key, "publicExponent", newbuf, &nblen))) {
            logger(Core, Error, "%s:%s:%d Failed to get RSA public key's exponent. Error = 0x%x (%s)\n",
                   __FILE__, __func__, __LINE__, asn1_rv, asn1_strerror(asn1_rv));
            free(m->data);
            return ASN_READ_ERROR;
        }

        e->size = nblen;

        e->data = malloc(e->size);
        if (!e->data) {
            logger(Core, Error, "%s:%s:%d Failed to allocate memory for exponent.\n", __FILE__, __func__, __LINE__);
            free(m->data);
            return ASN_MEM_ERROR;
        }

        memcpy((void *)e->data, newbuf, e->size);
    }

    return ASN_SUCCESS;
}

int libtasn_read_cert_pk_oid(uint8_t *data, size_t len, char *oid, size_t *oid_size)
{
    return libtasn_read_cert_pk(data, len, oid, oid_size, NULL, NULL);
}

int libtasn_read_cert_pk_parameters(uint8_t *data, size_t len, gnutls_datum_t *m, gnutls_datum_t *e)
{
    return libtasn_read_cert_pk(data, len, NULL, NULL, m, e);
}