/*
 * Copyright (c) 2015, Google Inc.
 *
 * SPDX-License-Identifier:	GPL-2.0+
 */

#include <common.h>

#if defined(CC_SUPPORT_AVB)||defined(CONFIG_ANDROID_SIGNATURE)

#include <image.h>
#include <android_image.h>
#include <u-boot/sha256.h>
#include <u-boot/rsa-mod-exp.h>
#include <u-boot/rsa.h>
#include <u-boot/rsa-checksum.h>
#include "gcpu_if.h"
#include "x_bim.h"

#define DEVICE_UNLOCK   0
#define DEVICE_LOCK     1
#define BOOT_GREEN      0
#define BOOT_YELLOW     1
#define BOOT_ORANGE     2
/**
 * tee_binding()
 */
static void tee_binding(unsigned int os_version, unsigned int state, uint8_t *pubkey, uint8_t *vbmeta_digest)
{
    uint32_t *tee = (uint32_t *)(SRAM_START + 0x6000);

    tee[0] = 0xA355A254;
    tee[1] = os_version;
    tee[2] = state == BOOT_ORANGE ? DEVICE_UNLOCK : DEVICE_LOCK;
    tee[3] = state;
    if (pubkey != NULL)
        memcpy(tee + 4, pubkey, RSA2048_BYTES);
    if (vbmeta_digest != NULL)
    {
        /* offset = 4 + public_key_len + 1(drmdebug flag) */
        memcpy(tee + 4 + RSA2048_BYTES/sizeof(uint32_t) + 1, vbmeta_digest, SHA256_SUM_LEN);
    }
}

#ifdef CC_SUPPORT_AVB
#include "../lib/libavb_ab_uboot/libavb_ab_uboot.h"

void check_android_boot_state(AvbSlotVerifyData *avb_data, AvbABFlowResult avb_ab_flow_result, unsigned int os_version)
{
    AvbABOps* ab_ops;
    bool device_is_unlocked = false;

    ab_ops = avb_get_ab_ops();
    ab_ops->ops->read_is_device_unlocked(ab_ops->ops, &device_is_unlocked);

    if(avb_ab_flow_result == AVB_AB_FLOW_RESULT_OK
        || avb_ab_flow_result == AVB_AB_FLOW_RESULT_OK_WITH_VERIFICATION_ERROR)
    {
        //green state.
        if(avb_ab_flow_result == AVB_AB_FLOW_RESULT_OK && device_is_unlocked == false)
        {
            env_add_boot_state(BOOT_GREEN);
            tee_binding(os_version, BOOT_GREEN, avb_data->vbmeta_images->public_key_data,
                        avb_data->vbmeta_images->vbmeta_digest_data);
            printf("Boot state is: Green !\n\n");
            return;
        }

        #ifdef  CC_ANDROID_SECURE_BOOT_CLASS_B
        if(device_is_unlocked)
        {// orange state.
            PanelLogo(1);
            env_add_boot_state(BOOT_ORANGE);
            tee_binding(os_version, BOOT_ORANGE, avb_data->vbmeta_images->public_key_data,
                        avb_data->vbmeta_images->vbmeta_digest_data);
            printf("\nBoot state is: Orange !\n\n");
            return;
        }
        else if(avb_ab_flow_result == AVB_AB_FLOW_RESULT_OK_WITH_VERIFICATION_ERROR
            && avb_data->vbmeta_images->verify_result == AVB_SLOT_VERIFY_RESULT_ERROR_PUBLIC_KEY_REJECTED)
        {// yellow state.
            PanelLogo(2);
            env_add_boot_state(BOOT_YELLOW);
            tee_binding(os_version, BOOT_YELLOW, avb_data->vbmeta_images->public_key_data,
                        avb_data->vbmeta_images->vbmeta_digest_data);
            printf("\nBoot state is: Yellow !\n\n");
            return;
        }
        #endif
        // In normal boot, it won't reach here.
        printf("\n%s:%d: Android Verify Boot 2.0 fail, unexpected verify result: %d\n\n",__FILE__,__LINE__,(avb_data->vbmeta_images->verify_result));
    }
    else
    {
        printf("\n%s:%d: Android Verify Boot 2.0 fail, AvbABFlowResult: %d\n\n",__FILE__,__LINE__,avb_ab_flow_result);
    }
    //halted for red state.
    PanelLogo(0);
    printf("\nBoot state is: Red, System halted !\n\n");
    while (1);
}
#endif //CC_SUPPORT_AVB

#ifdef CONFIG_ANDROID_SIGNATURE

extern INT32 RSADecryption65537(UINT32 *pu1Signature, UINT32 *pu4PublicKey, UINT32 *pu4CheckSum);
extern void env_add_boot_state(unsigned int state);

#define BOOT_SIGNATURE_VERSION      1

#define DER_TYPE_INTEGER            0x2
#define DER_TYPE_BITSTRING          0x3
#define DER_TYPE_OCTETSTRING        0x4
#define DER_TYPE_OBJECTIDENTIFIER   0x6
#define DER_TYPE_PRINTABLESTRING    0x13
#define DER_TYPE_SEQUENCE           0x30
#define DER_TYPE_PRIVATE            0xA0

static const uint8_t oid_sha256_rsa[] = {
    0x2A, 0x86, 0x48, 0x86, 0xF7, 0x0D, 0x01, 0x01, 0x0B
};

static const uint8_t andr_pubkey[] = {CONFIG_ANDROID_SIGNATURE_PUBKEY};

struct andr_boot_sig {
    const uint8_t *auth_attrs;
    uint32_t auth_attrs_length;

    const uint8_t *signature;
    uint32_t signature_length;

    uint32_t image_length;
    uint8_t digest[SHA256_SUM_LEN];

    const uint8_t *certificate;
    const uint8_t *public_key;
};

#define RSA2048_WORDS (RSA2048_BYTES / sizeof(uint32_t))

struct mincrypt_rsa_publickey {
    int len;                    /* Length of n[] in number of uint32_t */
    uint32_t n0inv;             /* -1 / n[0] mod 2^32 */
    uint32_t n[RSA2048_WORDS];  /* modulus as little endian array */
    uint32_t rr[RSA2048_WORDS]; /* R^2 as little endian array */
    int exponent;               /* 3 or 65537 */
};

/**
 * get_der_integer() - decodes a big endian DER integer
 * @p:		Pointer to the beginning of the encoded content
 * @length:	Length of the value in octets
 * @value:	Receives the value
 *
 * Return: Zero on success, otherwise on failure.
 */
static int get_der_integer(const uint8_t *p, uint32_t length, uint32_t *value)
{
    switch (length) {
    case 1:
        *value = *p;
        break;
    case 2:
        *value = be16_to_cpu(*(uint16_t *)p);
        break;
    case 3:
        *value = be32_to_cpu((*(uint32_t *)p & 0xFFFFFF))>>8;
        break;
    case 4:
        *value = be32_to_cpu(*(uint32_t *)p);
        break;
    default:
        printf("Unsupported integer length (%u octets)\n", length);
        return -1;
    }
    return 0;
}

/**
 * get_der_tlv() - decodes simple DER type-length-value tuples
 * @p:		Pointer to the beginning of the encoded content
 * @expected:	If non-zero, the expected type identifier
 * @type:	Receives the type identifier
 * @length:	Receives the length of the value
 * @value:	Receives a pointer to the value
 * @next:	Receives a pointer to the start of the next tuple
 *
 * Return: Zero on success, otherwise on failure.
 */
static int get_der_tlv(const uint8_t *p, uint8_t expected, uint8_t *type,
                       uint32_t *length, const uint8_t **value, const uint8_t **next)
{
    ulong start = 2;
    *type = p[0];

    if ((expected && *type != expected) ||
        (expected == DER_TYPE_PRIVATE && *type & DER_TYPE_PRIVATE == 0))
    {
        printf("Unexpected type %02x (expected %02x)\n", *type, expected);
        return -1;
    }

    if (p[1] & 0x80) { /* long form */
        uint8_t octets = p[1] ^ 0x80;

        if (get_der_integer(&p[2], octets, length) < 0) {
            printf("Unsupported length\n");
            return -1;
        }
        start += octets;
    } else {
        *length = p[1]; /* short form */
    }
    *value = &p[start];

    if (*type == DER_TYPE_SEQUENCE)
        *next = *value;
    else
        *next = &(*value)[*length];

    return 0;
}

/**
 * get_image_digest() - calculates a SHA256 hash of the boot image and
 *			authenticated attributes
 * @hdr:	Pointer to the boot image
 * @s:		Pointer to the signature structure
 */
static void get_image_digest(const struct andr_img_hdr *hdr,
                             struct andr_boot_sig *s)
{
#ifdef CONFIG_SHA256
    sha256_context ctx;

    sha256_starts(&ctx);
    sha256_update(&ctx, (const uint8_t *)hdr, s->image_length);
    sha256_update(&ctx, s->auth_attrs, s->auth_attrs_length);
    sha256_finish(&ctx, s->digest);
#else
    GCPU_SHA_HDL_T ctx;

    GCPU_Init(0);
    GCPU_SHA256_Init(&ctx, FALSE);
    GCPU_SHA256_Update(&ctx, (const uint8_t *)hdr, s->image_length, FALSE);
    GCPU_SHA256_Update(&ctx, s->auth_attrs, s->auth_attrs_length, TRUE);
    GCPU_SHA256_Final(&ctx, s->digest);
#endif // CONFIG_SHA256
}

/**
 * parse_boot_signature() - decodes and validates boot signature format
 * @hdr:	Pointer to the boot image
 * @s:		Pointer to the signature structure to fill
 *
 * Return: Zero on success, otherwise on failure.
 */
static int parse_boot_signature(const struct andr_img_hdr *hdr,
                                struct andr_boot_sig *s)
{
    const uint8_t *p = (const uint8_t *)android_image_get_end(hdr);
    const uint8_t *value;
    uint8_t type;
    uint32_t image_length;
    uint32_t length;
    uint32_t tmp;

    memset(s, 0, sizeof(struct andr_boot_sig));
    image_length = (uint32_t)(p - (const uint8_t *)hdr);

    if (get_der_tlv(p, DER_TYPE_SEQUENCE, &type, &length, &value, &p) < 0)
        return -1;

    if (length == 0) {
        printf("Empty signature block?\n");
        return -1;
    }

    /* format version */
    if (get_der_tlv(p, DER_TYPE_INTEGER, &type, &length, &value, &p) < 0)
        return -1;

    if (get_der_integer(value, length, &tmp) < 0)
        return -1;

    if (tmp != BOOT_SIGNATURE_VERSION) {
        printf("Unsupported boot signature format %u\n", tmp);
        return -1;
    }

    /* certificate */
    if (get_der_tlv(p, DER_TYPE_SEQUENCE, &type, &length, &value, &p) < 0)
        return -1;
    s->certificate = p;
    p = &value[length]; /* ignore */

    /* algorithm identifier */
    if (get_der_tlv(p, DER_TYPE_SEQUENCE, &type, &length, &value, &p) < 0)
        return -1;

    if (get_der_tlv(p, DER_TYPE_OBJECTIDENTIFIER, &type, &length, &value, &p) < 0)
        return -1;

    if (length != ARRAY_SIZE(oid_sha256_rsa) ||
        memcmp(value, oid_sha256_rsa, length)) {
        printf("Unsupported signature algorithm\n");
        return -1;
    }
    /* authenticated attributes */
    s->auth_attrs = p;

    if (get_der_tlv(p, DER_TYPE_SEQUENCE, &type, &length, &value, &p) < 0)
        return -1;
    s->auth_attrs_length = (value - s->auth_attrs) + length;

    if (get_der_tlv(p, DER_TYPE_PRINTABLESTRING, &type, &length, &value, &p) < 0)
        return -1;

    if (get_der_tlv(p, DER_TYPE_INTEGER, &type, &length, &value, &p) < 0)
        return -1;

    if (get_der_integer(value, length, &tmp) < 0)
        return -1;

    if (tmp != image_length) {
        printf("Unexpected boot image length %u (expected %u)\n", tmp, image_length);
        return -1;
    }
    s->image_length = tmp;

    /* signature */
    if (get_der_tlv(p, DER_TYPE_OCTETSTRING, &type, &length, &value, &p) < 0)
        return -1;

    if (length != RSA2048_BYTES) {
        printf("Unsupported signature length %u (expected %u)\n", length, RSA2048_BYTES);
        return -1;
    }
    s->signature = value;
    s->signature_length = length;

    get_image_digest(hdr, s);

    return 0;
}

/**
 * parse_certificate() - decodes public key from certificate
 * @s:		Pointer to boot signature
 *
 * Return: Zero on success, otherwise on failure.
 */
static int parse_certificate(struct andr_boot_sig *s)
{
    const uint8_t *p = s->certificate;
    const uint8_t *value;
    uint8_t type;
    uint32_t length;
    uint32_t i;

    /* Certificate */
    if (get_der_tlv(p, DER_TYPE_SEQUENCE, &type, &length, &value, &p) < 0)
        return -1;

    /* Version Number */
    if (get_der_tlv(p, DER_TYPE_PRIVATE, &type, &length, &value, &p) < 0)
        return -1;

    /* Serial Number */
    if (get_der_tlv(p, DER_TYPE_INTEGER, &type, &length, &value, &p) < 0)
        return -1;

    /* Signature Algorithm ID */
    if (get_der_tlv(p, DER_TYPE_SEQUENCE, &type, &length, &value, &p) < 0)
        return -1;
    p = &value[length]; /* ignore */

    /* Issuer Name */
    if (get_der_tlv(p, DER_TYPE_SEQUENCE, &type, &length, &value, &p) < 0)
        return -1;
    p = &value[length]; /* ignore */

    /* Validity period */
    if (get_der_tlv(p, DER_TYPE_SEQUENCE, &type, &length, &value, &p) < 0)
        return -1;
    p = &value[length]; /* ignore */

    /* Subject name */
    if (get_der_tlv(p, DER_TYPE_SEQUENCE, &type, &length, &value, &p) < 0)
        return -1;
    p = &value[length]; /* ignore */

    /* Subject Public Key Info */
    if (get_der_tlv(p, DER_TYPE_SEQUENCE, &type, &length, &value, &p) < 0)
        return -1;

    /* Public Key Algorithm */
    if (get_der_tlv(p, DER_TYPE_SEQUENCE, &type, &length, &value, &p) < 0)
        return -1;
    p = &value[length]; /* ignore */

    /* Subject Public Key */
    if (get_der_tlv(p, DER_TYPE_BITSTRING, &type, &length, &value, &p) < 0)
        return -1;
    p = &value[1];

    /* RSA Public Key */
    if (get_der_tlv(p, DER_TYPE_SEQUENCE, &type, &length, &value, &p) < 0)
        return -1;

    /* Modulus */
    if (get_der_tlv(p, DER_TYPE_INTEGER, &type, &length, &value, &p) < 0)
        return -1;
    s->public_key = &value[1];

    /* Exponent */
    /*if (get_der_tlv(p, DER_TYPE_INTEGER, &type, &length, &value, &p) < 0)
        return -1;*/
    return 0;
}

/**
 * reverse() - reverses an array of bytes from s to d
 * @d:		Pointer to the destination buffer
 * @s:		Pointer to the source buffer
 * @length:	Buffer length
 */
static void reverse(void *d, const void *s, uint32_t length)
{
    uint8_t *dst = (uint8_t *)d;
    const uint8_t *src = (const uint8_t *)s;
    uint32_t i;

    for (i = 0; i < length; ++i)
        dst[i] = src[length - i - 1];
}

/**
 * get_public_key() - converts Android mincrypt public key to struct key_prop
 * @prop:	Pointer to struct key_prop to fill
 * @n:		Pointer to RSA2048_BYTES byte buffer to use for n
 * @rr:		Pointer to RSA2048_BYTES byte buffer to use for r^2
 *
 * Return: Zero on success, otherwise on failure.
 */
static int get_public_key(struct key_prop *prop, uint8_t *n, uint8_t *rr)
{
    const struct mincrypt_rsa_publickey *key =
        (struct mincrypt_rsa_publickey *)andr_pubkey;
    uint32_t len = le32_to_cpu(key->len) * sizeof(uint32_t);

    if (len != RSA2048_BYTES) {
        printf("Unsupported public key length %u\n", len);
        return -1;
    }
    memset(prop, 0, sizeof(struct key_prop));

    reverse(n, key->n, len);
    reverse(rr, key->rr, len);

    prop->modulus = n;
    prop->n0inv = le32_to_cpu(key->n0inv);
    prop->num_bits = len * 8;
    prop->rr = rr;

    return 0;
}

/**
 * rsa_verify_user_key() - verify user key
 * @s:		Pointer to boot signature
 *
 * Return: Zero on success, otherwise on failure.
 */
static int rsa_verify_user_key(struct andr_boot_sig *s)
{
    uint8_t sig[RSA2048_BYTES];
    uint8_t chk[RSA2048_BYTES];
    uint8_t dig[RSA2048_BYTES] = {
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0x20, 0x04, 0x00, 0x05, 0x01, 0x02, 0x04, 0x03, 0x65, 0x01, 0x48, 0x86, 0x60, 0x09, 0x06, 0x0d,
        0x30, 0x31, 0x30, 0x00, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
        0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x01, 0x00};

    parse_certificate(s);
    reverse(sig, s->signature, RSA2048_BYTES);
    reverse(chk, s->public_key, RSA2048_BYTES);
    reverse(dig, s->digest, SHA256_SUM_LEN);
    RSADecryption65537(sig, chk, chk);

    if (memcmp(dig, chk, RSA2048_BYTES))
        return -EACCES;
    return 0;
}

/**
 * compare_with_oem_key() - check if user key is the same as oem key
 * @s:		Pointer to boot signature
 *
 * Return: Zero on success, otherwise on failure.
 */
static int compare_with_oem_key(struct andr_boot_sig *s)
{
    uint8_t chk[RSA2048_BYTES];
    reverse(chk, s->public_key, RSA2048_BYTES);

    if (memcmp(chk, andr_pubkey + 8, RSA2048_BYTES))
        return -EACCES;
#if 0   //test verify by user key embedded in boot image
    if (getenv("test_ukey") != NULL)
        return -1;
#endif
    return 0;
}

/**
 * verify_boot_signature() - verifies the signature
 * @hdr:	Pointer to boot image
 * @s:		Pointer to boot signature
 *
 * Return: Zero on success, otherwise on failure.
 */
static int verify_boot_signature(const struct andr_img_hdr *hdr,
                                 struct andr_boot_sig *s)
{
    int ret;
    struct key_prop prop;
    uint8_t n[RSA2048_BYTES];
    uint8_t rr[RSA2048_BYTES];

    struct checksum_algo algo = {
        "sha256",
        SHA256_SUM_LEN,
        RSA2048_BYTES,
#if IMAGE_ENABLE_SIGN
        NULL,
#endif
        NULL,
        padding_sha256_rsa2048
    };

    ret = get_public_key(&prop, n, rr);

    if (ret < 0)
        return ret;

    //get_image_digest(hdr, s);

    ret = rsa_verify_key(&prop, s->signature, s->signature_length, s->digest, &algo);

    if (ret < 0)
        printf("failed (%d)", ret);
    return ret;
}

/**
 * android_image_verify() - verifies the integrity of a signed boot image
 * @hdr:	Pointer to image header, which is at the start of the image.
 *
 * Return: Zero on success, otherwise on failure.
 */
int android_image_verify(const struct andr_img_hdr *hdr)
{
    struct andr_boot_sig sig;

    printf("signature verification ");
#ifdef  CC_ANDROID_SECURE_BOOT_CLASS_B
    {
        char *dev_state = getenv("devicestate");
        if (dev_state != NULL && strcmp(dev_state, "unlock") == 0)
        {
            printf("is bypassed\n");
            PanelLogo(1);
            env_add_boot_state(BOOT_ORANGE);
            tee_binding(hdr->os_version, BOOT_ORANGE, NULL, NULL);
            return 0; //unlock/orange state
        }
    }
#endif//CC_ANDROID_SECURE_BOOT_CLASS_B

    if (parse_boot_signature(hdr, &sig) == 0 &&
        rsa_verify_user_key(&sig) == 0)
    {
        if (compare_with_oem_key(&sig) == 0)
        {
            printf("is OK\n");
            env_add_boot_state(BOOT_GREEN);
            tee_binding(hdr->os_version, BOOT_GREEN, sig.public_key, NULL);
            return 0;   //lock/green state
        }
#ifdef  CC_ANDROID_SECURE_BOOT_CLASS_B
        else
        {
            printf("is OK by custom key\n");
            PanelLogo(2);
            env_add_boot_state(BOOT_YELLOW);
            tee_binding(hdr->os_version, BOOT_YELLOW, sig.public_key, NULL);
            return 0;   //lock/yellow state
        }
#endif//CC_ANDROID_SECURE_BOOT_CLASS_B
    }
    printf("failed\n");
    PanelLogo(0);
    printf("System is halted\n");
    while (1);          //lock/red state
}

#endif //#ifdef CONFIG_ANDROID_SIGNATURE

#endif //defined(CC_SUPPORT_AVB)||defined(CONFIG_ANDROID_SIGNATURE)

