/*
 * Copyright (C) 2011-2020 Intel Corporation. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above copyright
 *     notice, this list of conditions and the following disclaimer in
 *     the documentation and/or other materials provided with the
 *     distribution.
 *   * Neither the name of Intel Corporation nor the names of its
 *     contributors may be used to endorse or promote products derived
 *     from this software without specific prior written permission.
 *
 * 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 <string.h>
 #include <assert.h>
 #include <iostream>
 #include <unistd.h>
 #include <pwd.h>
 #include <fstream>
 #define MAX_PATH FILENAME_MAX
 #include <sgx_uswitchless.h>
 #include "sgx_urts.h"
 #include "App.h"
 #include "Enclave_u.h"
 #include <sgx_eid.h>
 #include <sgx_error.h>
 
 void *CurrentMemoryContext;
 void *pg_detoast_datum;
 void *MemoryContextAlloc;
 void *MemoryContextAllocZero;
 
 // extern "C"
 // {
 //     void errfinish();
 //     int pg_sprintf(char *__restrict__ __s, const char *__restrict__ __format, ...){};
 //     void errmsg();
 //     void errstart();
 //     void errcode();
 //     void ocall_ereport(const char*) {}
 //     void ocall_usleep(size_t sz){  
 //         sleep(sz);
 //     }
 //     void* palloc(int sz) {
 //         return malloc(sz);
 //     }
 //     void* palloc0(int sz) {
 //         return malloc(sz);
 //     }
 // }
 extern "C"
 {
     // 只保留声明，删除实现
     void errfinish();
     int pg_sprintf(char *__restrict__ __s, const char *__restrict__ __format, ...);
     void errmsg();
     void errstart();
     void errcode();
     void ocall_ereport(const char*);
     void ocall_usleep(size_t sz);
     void* palloc(int sz);
     void* palloc0(int sz);
 }
 
 // EnclaveTest.cpp
 
 extern "C" {
     void* pg_detoast_datum_packed(void* datum) {
         // 简单占位实现，返回输入
         return datum;
     }
 
     void ExceptionalCondition(const char* conditionName, const char* errorType, const char* fileName, int lineNumber) {
         // 简单占位实现，打印错误信息
         fprintf(stderr, "ExceptionalCondition: %s, %s, %s, %d\n", conditionName, errorType, fileName, lineNumber);
     }
 }
 
 // 将实现移到这里，并重命名避免冲突
 void ocall_usleep_impl(size_t sz) {  
     sleep(sz);
 }
 
 void* test_palloc(int sz) {
     return malloc(sz);
 }
 
 void* test_palloc0(int sz) {
     return malloc(sz);
 }
 
 void errfinish(){};
 void errmsg(){};
 void errstart(){};
 void errcode(){};
 
 using namespace std;
 /* Global EID shared by multiple threads */
 sgx_enclave_id_t global_eid = 0;
 
 typedef struct _sgx_errlist_t
 {
     sgx_status_t err;
     const char *msg;
     const char *sug; /* Suggestion */
 } sgx_errlist_t;
 
 /* Error code returned by sgx_create_enclave */
 static sgx_errlist_t sgx_errlist[] = {
     {SGX_ERROR_UNEXPECTED,
      "Unexpected error occurred.",
      NULL},
     {SGX_ERROR_INVALID_PARAMETER,
      "Invalid parameter.",
      NULL},
     {SGX_ERROR_OUT_OF_MEMORY,
      "Out of memory.",
      NULL},
     {SGX_ERROR_ENCLAVE_LOST,
      "Power transition occurred.",
      "Please refer to the sample \"PowerTransition\" for details."},
     {SGX_ERROR_INVALID_ENCLAVE,
      "Invalid enclave image.",
      NULL},
     {SGX_ERROR_INVALID_ENCLAVE_ID,
      "Invalid enclave identification.",
      NULL},
     {SGX_ERROR_INVALID_SIGNATURE,
      "Invalid enclave signature.",
      NULL},
     {SGX_ERROR_OUT_OF_EPC,
      "Out of EPC memory.",
      NULL},
     {SGX_ERROR_NO_DEVICE,
      "Invalid SGX device.",
      "Please make sure SGX module is enabled in the BIOS, and install SGX driver afterwards."},
     {SGX_ERROR_MEMORY_MAP_CONFLICT,
      "Memory map conflicted.",
      NULL},
     {SGX_ERROR_INVALID_METADATA,
      "Invalid enclave metadata.",
      NULL},
     {SGX_ERROR_DEVICE_BUSY,
      "SGX device was busy.",
      NULL},
     {SGX_ERROR_INVALID_VERSION,
      "Enclave version was invalid.",
      NULL},
     {SGX_ERROR_INVALID_ATTRIBUTE,
      "Enclave was not authorized.",
      NULL},
     {SGX_ERROR_ENCLAVE_FILE_ACCESS,
      "Can't open enclave file.",
      NULL},
 };
 
 /* Check error conditions for loading enclave */
 void print_error_message(sgx_status_t ret)
 {
     size_t idx = 0;
     size_t ttl = sizeof sgx_errlist / sizeof sgx_errlist[0];
 
     for (idx = 0; idx < ttl; idx++)
     {
         if (ret == sgx_errlist[idx].err)
         {
             if (NULL != sgx_errlist[idx].sug)
                 printf("Info: %s\n", sgx_errlist[idx].sug);
             printf("Error: %s\n", sgx_errlist[idx].msg);
             break;
         }
     }
 
     if (idx == ttl)
         printf("Error code is 0x%X. Please refer to the \"Intel SGX SDK Developer Reference\" for more details.\n", ret);
 }
 
 const char* find_enclave_file() {
     // 按优先级尝试不同路径
     const char* paths[] = {
         ENCLAVE_FILENAME,                  // 默认路径
         "../lib/enclave.signed.so",        // 相对于测试程序的路径
         "./enclave.signed.so",             // 当前目录
         "../enclave.signed.so",            // 上级目录
         "../../lib/enclave.signed.so",     // 上上级lib目录
         "../../../lib/enclave.signed.so",  // 再上一级lib目录
         "./lib/enclave.signed.so",         // 当前目录下的lib
         "/etc/encryptsql/enclave.signed.so" // 系统安装路径
     };
     
     for (const char* path : paths) {
         struct stat st;
         if (stat(path, &st) == 0 && st.st_size > 0) {
             printf("找到有效的enclave文件: %s\n", path);
             return path;
         }
     }
     
     // 如果找不到，打印搜索路径
     printf("在以下路径中未找到有效的enclave文件:\n");
     for (const char* path : paths) {
         struct stat st;
         if (stat(path, &st) == 0) {
             printf("- %s (大小: %ld bytes)\n", path, st.st_size);
         } else {
             printf("- %s (不存在)\n", path);
         }
     }
     
     return NULL;  // 如果都找不到，返回NULL
 }
 
 /* Initialize the enclave:
  *   Call sgx_create_enclave to initialize an enclave instance
  */
 int initialize_enclave(void)
 {
     sgx_status_t ret = SGX_ERROR_UNEXPECTED;
     
     // 获取正确的enclave文件路径
     const char* enclave_file = find_enclave_file();
     if (!enclave_file) {
         printf("错误: 未找到有效的enclave文件\n");
         return -1;
     }
     
     // 打印完整路径
     char full_path[PATH_MAX];
     if (realpath(enclave_file, full_path) != NULL) {
         printf("Enclave完整路径: %s\n", full_path);
     } else {
         printf("无法解析Enclave路径: %s\n", enclave_file);
         perror("realpath error");
         return -1;
     }
     
     // 检查文件权限
     struct stat st;
     if (stat(enclave_file, &st) == 0) {
         printf("文件权限: %o\n", st.st_mode & 0777);
         printf("所有者: %d\n", st.st_uid);
         printf("组: %d\n", st.st_gid);
         
         // 检查文件大小
         if (st.st_size == 0) {
             printf("错误: Enclave文件是空的\n");
             return -1;
         }
         printf("Enclave文件大小: %ld bytes\n", st.st_size);
         
         // 检查权限是否过于开放
         if ((st.st_mode & 0777) == 0777) {
             printf("警告: Enclave文件权限过于开放，建议修改为644\n");
         }
     } else {
         printf("无法获取文件信息: %s\n", strerror(errno));
         return -1;
     }
     
     // 检查文件是否存在且可读
     if (access(enclave_file, R_OK) != 0) {
         printf("错误: 无法访问Enclave文件: %s\n", strerror(errno));
         return -1;
     }
     
     // 检查SGX设备是否可用
    //  if (access("/dev/sgx", F_OK) == -1 && access("/dev/isgx", F_OK) == -1) {
    //      printf("错误: 未找到SGX设备。请确认SGX驱动已正确安装。\n");
    //      return -1;
    //  }


     //模拟状态下
    bool sgx_available = false;
    const char *sgx_device_paths[] = {"/dev/sgx", "/dev/isgx", "/dev/sgx_enclave"};
    for (int i = 0; i < sizeof(sgx_device_paths)/sizeof(sgx_device_paths[0]); i++) {
        if (access(sgx_device_paths[i], F_OK) == 0) {
            sgx_available = true;
            printf("检测到SGX物理设备: %s\n", sgx_device_paths[i]);
            break;
        }
    }

    if (!sgx_available) {
        printf("警告: 未找到SGX物理设备，将尝试使用模拟模式\n");
    }

     printf("正在创建enclave: %s\n", enclave_file);
     
     // 尝试创建enclave
     ret = sgx_create_enclave(enclave_file,
                             0,  // 非调试模式
                             NULL,
                             NULL,
                             &global_eid,
                             NULL);
     
     if (ret != SGX_SUCCESS) {
         print_error_message(ret);
         if (ret == SGX_ERROR_OUT_OF_MEMORY) {
             printf("内存不足。请检查:\n");
             printf("1. 系统可用内存\n");
             printf("2. SGX EPC大小设置\n");
             printf("3. 其他占用内存的进程\n");
             
             // 打印当前内存使用情况
             FILE* fp = popen("free -h", "r");
             if (fp) {
                 char buffer[256];
                 while (fgets(buffer, sizeof(buffer), fp)) {
                     printf("%s", buffer);
                 }
                 pclose(fp);
             }
             
             // 打印SGX相关信息
             printf("\nSGX 设备信息:\n");
             system("ls -l /dev/sgx*");
             printf("\nSGX 驱动状态:\n");
             system("systemctl status aesmd | head -n 3");
             
             // 查EPC大小
             FILE* fp2 = popen("dmesg | grep -i sgx", "r");
             if (fp2) {
                 char buffer[256];
                 printf("\nSGX 内核日志:\n");
                 while (fgets(buffer, sizeof(buffer), fp2)) {
                     printf("%s", buffer);
                 }
                 pclose(fp2);
             }
         }
         return -1;
     }
 
     printf("Enclave创建成功，ID: %lu\n", global_eid);
     return 0;
 }
 
 /* OCall functions */
 void ocall_print_string(const char *str)
 {
     /* Proxy/Bridge will check the length and null-terminate 
      * the input string to prevent buffer overflow. 
      */
     printf("%s", str);
 }
 
 char plaPath[] = "/etc/encryptsql/masterKey.plain";
 char encPath[] = "/etc/encryptsql/masterKey.encrypted";
 constexpr int keySize = 32;
 
 void encryptKey() {
 
     // ecall_genMasterKey(global_eid, path);
     // ecall_openMasterKey(global_eid, path);
     using namespace std;
     ifstream fin;
     fin.open(plaPath, ios::binary | ios::in);
     if(!fin.good()) {
         printf("can not open key file:%s\n", plaPath);
         exit(0);
     }
     char key[keySize];
 
     fin.get(key, keySize + 1);
     ecall_encryptKeyWithEnclave(global_eid, key, keySize, encPath);
 }
 
 void getMasterKeyFromEnclave()
 {
     char masterKey[keySize+1];
     auto ret = ecall_openMasterKey(global_eid, encPath, masterKey, keySize);
     if(ret < 0) {
         print_error_message(ret);
     }
     masterKey[keySize] = '\0';
     cout << "MasterKey from Enclave is : " << masterKey << endl;
 }
 
 /* Application entry */
 int SGX_CDECL main(int argc, char *argv[])
 {
     // 检查环境变量
     const char* sgx_path = getenv("SGX_ENCLAVE_PATH");
     if (sgx_path) {
         printf("SGX_ENCLAVE_PATH=%s\n", sgx_path);
     }
     
     (void)(argc);
     (void)(argv);
 
     printf("正在初始化enclave...\n");
     
     // 在初始化前清理一些资源
     system("sync"); // 同步文件系统缓存
     
     int ret = initialize_enclave();
     if (ret < 0) {
         printf("初始化Enclave失败，错误代码: %d\n", ret);
         return -1;
     }
     
     printf("Enclave初始化成功。\n");
     
     try {
         encryptKey();
         getMasterKeyFromEnclave();
     } catch (const std::exception& e) {
         printf("操作执行失败: %s\n", e.what());
         sgx_destroy_enclave(global_eid);
         return -1;
     }
 
     sgx_destroy_enclave(global_eid);
     return 0;
 }
