/*
 * Copyright (C) 2021 Amlogic, Inc. All rights reserved.
 *
 * All information contained herein is Amlogic confidential.
 *
 * This software is provided to you pursuant to Software License
 * Agreement (SLA) with Amlogic Inc ("Amlogic"). This software may be
 * used only in accordance with the terms of this agreement.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification is strictly prohibited without prior written permission
 * from Amlogic.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <openssl/evp.h>
#include <string.h>
#include <getopt.h>

#define DEBUG       (0)

#define BLOCK_SIZE  (16)
#define LBN_SIZE   (512)
#define LBN_SHIFT  (9)
#define IV_FACTOR_BASE_ADDR (0x800)

uint8_t key_bin[16] = { 0x00 };

uint32_t key_bin_len = sizeof(key_bin);

#if DEBUG
static void dump_bufs(const char *name, const uint8_t *buf, int32_t len)
{
    char line_buf[128] = {0x0};
    char *line_buf_ptr = line_buf;
    uint32_t i = 0;
    const uint8_t *ptr = buf;
    const int32_t sz = len;

    printf("\n========= Start dump %s(%d)\n", name, sz);
    while (len - 16 >= 0) {
        printf("%02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x\n",
                ptr[0], ptr[i+1], ptr[i+2], ptr[i+3], ptr[4], ptr[i+5], ptr[i+6], ptr[i+7],
                ptr[i+8], ptr[i+9], ptr[i+10], ptr[i+11], ptr[i+12], ptr[i+13], ptr[i+14], ptr[i+15]);
        ptr += 16;
        len -= 16;
    }

    if (len * 3 <= sizeof(line_buf)) {
        while (len) {
            snprintf(line_buf_ptr, sizeof(line_buf), "%02x ", *ptr);
            line_buf_ptr += 3; /*"xx "*/
            len--;
        }
        if (strlen(line_buf))
            printf("%s\n", line_buf);
        printf("========= End dump %s(%d)\n", name, sz);
    } else {
        printf("remaining length exceeds linebuf length\n");
    }
}
#endif

/* Get logical block number for encryption.
 * Size of each lbn is 512 bytes. */
static inline uint64_t get_lbn(uint64_t offset)
{
    return (offset >> LBN_SHIFT) + IV_FACTOR_BASE_ADDR;
}

uint32_t calc_md5(uint8_t *in, uint32_t in_sz, uint8_t *md5_hash, uint32_t *md5_hash_len)
{
    const EVP_MD *md = EVP_md5();

    if (!md || !in || !in_sz || !md5_hash || !md5_hash_len) {
        printf("%s:%d: Invalid input\n", __func__, __LINE__);
        return -1;
    }

    if (!(*md5_hash_len)) {
        printf("%s:%d: No output space\n", __func__, __LINE__);
        return -1;
    }

    if (EVP_Digest(in, in_sz, md5_hash, md5_hash_len, md, NULL) != 1) {
        printf("%s:%d: EVP_Digest failed\n", __func__, __LINE__);
        return -1;
    }

    return 0;
}

int32_t calc_aes_128_cbc(uint8_t *in, uint32_t in_sz,
        uint8_t *key, uint32_t key_sz,
        uint8_t *iv, uint32_t iv_sz,
        uint8_t *out, int32_t *out_sz)
{
    EVP_CIPHER_CTX *ctx = NULL;
    int32_t len;
    int32_t ret = 0;

    if (!in || !key || !iv || !out || !out_sz) {
        printf("%s:%d: Invalid input\n", __func__, __LINE__);
        ret = -1;
        goto out;
    }

    if (!(*out_sz)) {
        printf("%s:%d: No output space\n", __func__, __LINE__);
        return -1;
    }

    ctx = EVP_CIPHER_CTX_new();
    if (!ctx) {
        printf("%s:%d: failed to allocate ctx\n", __func__, __LINE__);
        return -1;
    }

    if (EVP_EncryptInit_ex(ctx, EVP_aes_128_cbc(), NULL, key, iv) != 1) {
        printf("%s:%d: EVP_EncryptInit_ex failed\n", __func__, __LINE__);
        ret = -1;
        goto out;
    }

    /* set no-padding */
    EVP_CIPHER_CTX_set_padding(ctx, 0);

    if (EVP_EncryptUpdate(ctx, out, out_sz, in, in_sz) != 1) {
        printf("%s:%d: EVP_EncryptUpdate failed\n", __func__, __LINE__);
        ret = -1;
        goto out;
    }

    if (EVP_EncryptFinal_ex(ctx, out + *out_sz, &len) != 1) {
        printf("%s:%d: EVP_EncryptFinal_ex failed\n", __func__, __LINE__);
        ret = -1;
    }
out:
    if (ctx)
        EVP_CIPHER_CTX_free(ctx);
    return ret;
}

static int32_t get_enc_iv(uint64_t offset, uint8_t *md5, uint32_t md5_sz, uint8_t *out, int32_t *out_sz)
{
    uint64_t lbn = get_lbn(offset);
    uint8_t iv[BLOCK_SIZE] = {};
    uint8_t iv_iv[BLOCK_SIZE] = {0xff};
    int32_t ret = 0;

    if (!out || !out_sz || !*out_sz) {
        printf("%s:%d: Invalid input.\n", __func__, __LINE__);
        return -1;
    }

    iv[0] = lbn & 0xFF;
    iv[1] = (lbn >> 8) & 0xFF;
    iv[2] = (lbn >> 16) & 0xFF;
    iv[3] = (lbn >> 24) & 0xFF;

    ret = calc_aes_128_cbc(iv, sizeof(iv),
            md5, md5_sz,          /* key */
            iv_iv, sizeof(iv_iv), /* iv  */
            out, out_sz);
    return ret;
}

static uint32_t get_file_sz(FILE *fp) {
    uint32_t sz = 0;
    uint32_t cur = ftell(fp);;

    fseek(fp, 0, SEEK_END);
    sz = ftell(fp);
    fseek(fp, cur, SEEK_SET);
    return sz;
}

void print_usage(void)
{
    printf("Usage: ./enc_part [options] (%s)\n", VERSION);
    printf("Options:\n");
    printf("  -i, --in\t\t\tInput file to be encrypted.\n");
    printf("  -o, --out\t\t\tEncrypted output file.\n");
    printf("  -k, --key\t\t\tEncryption key in hexdump format, e.g. aabb1122\n");
}

static int hex2bin(const char *hex, void *bin, size_t binlen)
{
    int i, c, n1, n2, hexlen, k;

    if (!hex || !bin || !binlen) {
        printf("%s:%d: invalid input.\n", __func__, __LINE__);
        return 0;
    }

    hexlen = strnlen(hex, 64);
    if (hexlen > binlen * 2) {
        printf("%s:%d: input overflow(%d, %zu)\n", __func__, __LINE__, hexlen, binlen);
        return 0;
    }

    k = 0;
    n1 = -1;
    n2 = -1;
    for (i = 0; i < hexlen; i++) {
        n2 = n1;
        c = hex[i];
        if (c >= '0' && c <= '9') {
            n1 = c - '0';
        } else if (c >= 'a' && c <= 'f') {
            n1 = c - 'a' + 10;
        } else if (c >= 'A' && c <= 'F') {
            n1 = c - 'A' + 10;
        } else if (c == ' ') {
            n1 = -1;
            continue;
        } else {
            return -1;
        }

        if (n1 >= 0 && n2 >= 0) {
            ((uint8_t *)bin)[k] = (n2 << 4) | n1;
            n1 = -1;
            k++;
        }
    }
    return k;
}

static void dump_params(const char *in, const char *out,
						const char *key) {
	printf("****************enc_part(%s)*********************\n", VERSION);
	printf("in file :\t%s\n", in);
	printf("out file:\t%s\n", out);
	printf("key:\t\t%s\n", key);
	printf("******************************************\n");
}

int main(int argc, char** argv)
{
    int32_t opt, long_index;
    FILE *in = NULL, *out = NULL;
    const char* in_file = NULL;
    uint32_t in_sz = 0;
    const char* out_file = NULL;
    const char* key = NULL;

    int32_t remaining, read_sz, ret, retVal = 0;
    uint8_t *buf = NULL, *o_buf = NULL;
    uint8_t page_iv[BLOCK_SIZE] = {};
    int32_t page_iv_sz = BLOCK_SIZE;
    uint64_t offset = 0;
    int32_t wr_sz;
    uint32_t padding = 0;
    uint32_t temp_sz = 0;
    uint32_t remain_of_blk = 0;

    uint8_t md5[16];
    uint32_t md5_sz = sizeof(md5);

    if (argc < 2) {
        print_usage();
        return -1;
    }

    struct option long_options[] = {
        { "in", required_argument, 0, 'i'},
        { "out", required_argument, 0, 'o'},
        { "key", required_argument, 0, 'k'},
    };

    while ((opt = getopt_long(argc, argv,"i:o:k:",
                    long_options, &long_index )) != -1) {
        switch (opt) {
            case 'i' :
                in_file = optarg;
                break;
            case 'o' :
                out_file = optarg;
                break;
            case 'k' :
                key = optarg;
                break;
            default:
                print_usage();
                return -1;
        }
    }

    if (!in_file || !out_file) {
        print_usage();
        return -1;
    }

    in = fopen(in_file, "rb");
    if (!in) {
        printf("failed to open file: %s\n", in_file);
        goto out;
    }
    in_sz = get_file_sz(in);

    /* Try to open file in place (in_file == out_file) */
    out = fopen(out_file, "r+b");
    if (!out) {
        /* Try to open file as a new one */
        out = fopen(out_file, "w+b");
        if (!out) {
            printf("failed to open file: %s\n", out_file);
            goto out;
        }
    }

    buf = (uint8_t*)malloc(LBN_SIZE);
    if (!buf) {
        printf("failed to malloc buffer(%u)\n", LBN_SIZE);
        goto out;
    }
    o_buf = (uint8_t*)malloc(LBN_SIZE);
    if (!o_buf) {
        printf("failed to malloc buffer(%u)\n", LBN_SIZE);
        goto out;
    }

    remaining = in_sz;
    key_bin_len = hex2bin(key, key_bin, key_bin_len);

    /* partition key hash */
    calc_md5(key_bin, key_bin_len, md5, &md5_sz);

    dump_params(in_file, out_file, key);

    while (remaining > 0) {
        /* Read in a lbn if it is possible */
        read_sz = (remaining >= LBN_SIZE)? LBN_SIZE: remaining;
        ret = fread(buf, sizeof(uint8_t), read_sz, in);
        if (ret != read_sz) {
            printf("%s:%d: Error to read in file\n", __func__, __LINE__);
            retVal = -1;
            goto out;
        }
        remain_of_blk = read_sz & (BLOCK_SIZE - 1);
        if (remain_of_blk) {
            /* Since LBN_SIZE is multiple of BLOCK_SIZE, entering
             * this case means it is the last read */
            padding = BLOCK_SIZE - remain_of_blk;
            temp_sz = read_sz + padding;
        }
        /* update iv */
        page_iv_sz = BLOCK_SIZE;
        ret = get_enc_iv(offset, md5, md5_sz, page_iv, &page_iv_sz);
        if (ret < 0) {
            printf("%s:%d: Error to get enc iv\n", __func__, __LINE__);
            retVal = -1;
            goto out;
        }
        if (remain_of_blk) {
            wr_sz = temp_sz;
            ret = calc_aes_128_cbc(buf, temp_sz,   /* input */
                    key_bin, key_bin_len,          /* key */
                    page_iv, page_iv_sz,           /* iv */
                    o_buf, &wr_sz);                /* output */
        } else {
            wr_sz = read_sz;
            ret = calc_aes_128_cbc(buf, read_sz, /* input*/
                    key_bin, key_bin_len,        /* key */
                    page_iv, page_iv_sz,         /* iv */
                    o_buf, &wr_sz);              /* output */
        }
        if (ret < 0) {
            printf("%s:%d: Error to calc aes 128 cbc\n", __func__, __LINE__);
            retVal = -1;
            goto out;
        }

        ret = fwrite(o_buf, sizeof(uint8_t), read_sz, out);
        if (ret != read_sz) {
            printf("Error to write output file\n");
            retVal = -1;
            goto out;
        }

        offset += read_sz;
        remaining -= read_sz;
    }
out:
    if (in) fclose(in);
    if (out) fclose(out);
    if (buf) free(buf);
    if (o_buf) free(o_buf);

    return retVal;
}