/*
 * SPDX-License-Identifier: LicenseRef-Unisoc-General-1.0
 *
 * Copyright 2016-2024 Unisoc (Shanghai) Technologies Co., Ltd
 *
 * Licensed under the Unisoc General Software License, version 1.0 (the License);
 * You may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * https://www.unisoc.com/en_us/license/UNISOC_GENERAL_LICENSE_V1.0-EN_US
 *
 * Software distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OF ANY KIND, either express or implied.
 * See the Unisoc General Software License, version 1.0 for more details.
 */

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <time.h>
#include <iostream>
#include <fstream>
#include <string>
#include <fcntl.h>
#include <unistd.h>

#include "log.h"
#include "anti-rollback.h"

#define SYS_HEADER_SIZE 0x200
#define MAX_CERT_SIZE    2048

const char* v_file="/mnt/etc/anti-rollback/version";
const char* s_file="/mnt/etc/anti-rollback/signature";
char boot_dev[15]={0};

uint8_t signature[4096]={0};
char cert[MAX_CERT_SIZE] = {0};

#define MAX_HASH_BITS_LEN (256)
#define MAX_HASH_BYTES_LEN (MAX_HASH_BITS_LEN >> 3)
#define IV_BYTE_LEN 16
#define MAGIC_SIZE 8

#define MAX_RSA4096_KEY_BITS_LEN (4096)
#define MAX_RSA_KEY_BITS_LEN (MAX_RSA4096_KEY_BITS_LEN)
#define MAX_RSA_KEY_BYTES_LEN (MAX_RSA_KEY_BITS_LEN>>3)

typedef struct sys_img_header {
    uint32_t magic_num;        // "BTHD"=="0x42544844"=="boothead"
    uint32_t version;         // 1
    uint8_t payload_hash[MAX_HASH_BYTES_LEN]; // sha256 hash value
    uint64_t imgaddr;         // image loaded address
    uint32_t imgsize;         // image size
    uint8_t iv_data[IV_BYTE_LEN]; // the parameter of AES crypto
    uint32_t is_packed;        // packed image flag 0:false 1:true
    uint32_t firmware_size;    // runtime firmware size
    uint32_t img_real_size;      //image real size before sign
    uint32_t size_after_sign;      //image all size after sign
    uint8_t reserved[424];     // 424 + 17*4 + 16 + 4= 512
    uint32_t slot;             // spl double slot: slota or slotb
} sys_img_header;

typedef struct sprdsignedimageheader {
    /* Magic number */
    uint8_t magic[MAGIC_SIZE];
    /* Version of this header format */
    uint32_t header_version_major;
    /* Version of this header format */
    uint32_t header_version_minor;

    /*image body, plain or cipher text */
    uint64_t payload_size;
    uint64_t payload_offset;

    /*offset from itself start */
    /*content certification size,if 0,ignore */
    uint64_t cert_size;
    uint64_t cert_offset;

    /*(opt)private content size,if 0,ignore */
    uint64_t priv_size;
    uint64_t priv_offset;

    /*(opt)debug/rma certification primary size,if 0,ignore */
    uint64_t cert_dbg_prim_size;
    uint64_t cert_dbg_prim_offset;

    /*(opt)debug/rma certification second size,if 0,ignore */
    uint64_t cert_dbg_developer_size;
    uint64_t cert_dbg_developer_offset;

} sprdsignedimageheader;

bool is_little_endian(void)
{
    uint16_t data = 1;
    return *(uint8_t*)(&data);
}

void write32(uint8_t* buffer, uint32_t value)
{
    if (is_little_endian()) {
        buffer[0] = value & 0xFF;
        buffer[1] = (value >> 8) & 0xFF;
        buffer[2] = (value >> 16) & 0xFF;
        buffer[3] = (value >> 24) & 0xFF;
    } else {
        buffer[3] = value & 0xFF;
        buffer[2] = (value >> 8) & 0xFF;
        buffer[1] = (value >> 16) & 0xFF;
        buffer[0] = (value >> 24) & 0xFF;
    }
}

void hexdump(const unsigned char *s, int l)
{
    int n=0;

    for( ; n < l ; ++n)
    {
        if((n%16) == 0)
        printf("\n%04x",n);
        printf(" %02x",s[n]);
    }
    printf("\n");
}

static int get_verison_and_signature(uint32_t* size) {
    std::ifstream file(v_file);
    std::ifstream file1(s_file, std::ios::binary);

    if (file.is_open() && file1.is_open()) {
        std::string line;
        std::getline(file, line);
        std::cout << "int version:" << atoi(line.c_str()) << std::endl;

        write32(signature + *size, line.length());
        *size += 4;
        memcpy(signature + *size, line.c_str(), line.length());
        *size += line.length();

        std::getline(file, line);
        std::cout << "algorithm:" << line << std::endl;

        //write32(signature + size, line.length());
        //size += 4;
        //memcpy(signature + size, line.c_str(), line.length());
        //size += line.length();

        std::getline(file, line);
        //std::cout << "hash:" << line << std::endl;

        write32(signature + *size, line.length());
        *size += 4;
        memcpy(signature + *size, line.c_str(), line.length());
        //hexdump(signature + *size, line.length());
        *size += line.length();

        write32(signature + *size, MAX_CERT_SIZE);
        *size += 4;
        memcpy(signature + *size, cert, MAX_CERT_SIZE);
        //std::cout << "pub key:" <<std::endl;
        //hexdump(signature + size, MAX_CERT_SIZE);
        *size += MAX_CERT_SIZE;

        write32(signature + *size, MAX_RSA_KEY_BYTES_LEN);
        *size += 4;
        file1.seekg(0);
        file1.read((char*)(signature + *size), MAX_RSA_KEY_BYTES_LEN);
        //std::cout << "signature:" <<std::endl;
        //hexdump(signature + *size, MAX_RSA_KEY_BYTES_LEN);
        *size += MAX_RSA_KEY_BYTES_LEN;

        std::cout << "total size:" << *size << std::endl;
        file.close();
        file1.close();
    } else {
        std::cout << "open file " << v_file <<" failed" << std::endl;
    }

    return 0;
}

bool load_pubkey_from_boot()
{
    sys_img_header img_header;
    sprdsignedimageheader sprd_img_header;
    uint64_t img_size = 0;
    uint64_t cert_size = 0;
    uint64_t cert_offset = 0;
    bool ret = false;

    int fd = open64(boot_dev,  O_RDONLY | O_CLOEXEC);
    if (fd < 0) {
        LOGE("open blk dev %s , err=%d", boot_dev, errno);
        goto failed1;
    }

    // 1.read imageheader
    if (pread64(fd, (char *)&img_header, sizeof(sys_img_header), 0) != sizeof(sys_img_header)) {
        LOGE("read imageheader error %d\n", errno);
        goto failed2;
    }
    img_size = img_header.imgsize;

    // 2.read sprdsignedimageheader
    if (pread64(fd, (char *)&sprd_img_header, sizeof(sprdsignedimageheader),
        img_size + SYS_HEADER_SIZE) != sizeof(sprdsignedimageheader)) {
        LOGE("read sprdsignedimageheader error %d\n", errno);
        goto failed2;
    }

    cert_size = sprd_img_header.cert_size;
    cert_offset = sprd_img_header.cert_offset;
    LOGE("[load_pubkey]cert_size %lu, cert_offset 0x%lX\n", cert_size, cert_offset);

    // 3.read cert
    if (pread64(fd, cert, MAX_CERT_SIZE, cert_offset) != MAX_CERT_SIZE) {
        LOGE("read cert error %d\n", errno);
        goto failed2;
    }
    ret = true;

failed2:
    close(fd);
failed1:
    return ret;
}

int main(int argc, char** argv)
{
    int ret = -1;
    uint8_t result = -1;
    uint32_t result_size = 0;
    char boot_system[15] = {0};

    snprintf(boot_system, sizeof(boot_system)-1, "%s", argv[1]);
    if (!strcmp(boot_system, "system_b"))
        snprintf(boot_dev, sizeof(boot_dev)-1, "%s", "/dev/boot_b");
    else if (!strcmp(boot_system, "system_a"))
        snprintf(boot_dev, sizeof(boot_dev)-1, "%s", "/dev/boot_a");
    else
        snprintf(boot_dev, sizeof(boot_dev)-1, "%s", "/dev/boot");
    LOGE("boot dev: %s\n", boot_dev);
    if (!load_pubkey_from_boot())
        goto failed;

    get_verison_and_signature(&result_size);

    ret = update_rootfs_verison(signature, &result_size, &result);
    LOGE("ret:%d, result:%d", ret, result);
    if (ret >= 0)
        ret = result;


failed:
    return ret;
}