package org.bdware.sc.engine;

import com.google.gson.reflect.TypeToken;
import org.apache.commons.io.FileUtils;
import org.bdware.sc.bean.ContractRequest;
import org.bdware.sc.node.ContractNode;
import org.bdware.sc.node.FunctionNode;
import org.bdware.sc.util.JsonUtil;
import wrp.jdk.nashorn.api.scripting.ScriptObjectMirror;
import wrp.jdk.nashorn.internal.objects.Global;

import javax.script.Invocable;
import javax.script.ScriptException;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.lang.reflect.Type;
import java.nio.charset.StandardCharsets;
import java.util.*;

public class ConfidentialContractUtil {

    public static final String CONFIDENTIAL_TEMPLATE_PATH = System.getenv("GRAPHENE_DIR")
            + File.separator + "Examples" + File.separator + "nodejs-secret";
    public static final String CONFIDENTIAL_SCRIPT_PATH = System.getenv("GRAPHENE_DIR")
            + File.separator + "App";
    public static final String[] COMMAND = {"bash", "executeContract.sh"};
    private static final Type MapType = TypeToken.getParameterized(HashMap.class, String.class, String.class).getType();

    public static String executeConfidentialContract(ContractRequest input) throws IOException, InterruptedException {
        File runDir = new File(CONFIDENTIAL_SCRIPT_PATH + File.separator + input.getRequestID());
        ProcessBuilder pb = new ProcessBuilder(COMMAND);
        pb.directory(runDir);
        Process p = pb.start();
        p.waitFor();
        File resultFile = new File(CONFIDENTIAL_SCRIPT_PATH + File.separator + input.getRequestID() +
                File.separator + "result.json");
        return FileUtils.readFileToString(resultFile, StandardCharsets.UTF_8);
    }

    public static void generateConfidentialContract(ContractNode cn, ScriptObjectMirror globalVars, Global global) {
        List<FunctionNode> functionNodes = cn.getFunctions();
        for (FunctionNode fn : functionNodes) {
            // assuming only one confidential function for now
            if (fn.isConfidential()) {
                StringBuilder jsStr = new StringBuilder();
                // find all dependent functions
                Set<String> dependentFunctions = findAllDependentFunctions(cn, fn);
                // add self and all dependent function declaration
                for (FunctionNode fNode : functionNodes) {
                    if (dependentFunctions.contains(fNode.functionName)) {
                        jsStr.append(fNode.plainText()).append("\n");
                    }
                }
                // load necessary Node.js libraries
                jsStr.append("var fs = require('fs');\n" + "var crypto = require('crypto');\n" + "var sm2 = require('sm-crypto').sm2;\n");
                // load Global variables and arguments from files
                jsStr.append("let rawGlobal = fs.readFileSync('global.json').toString();\n" + "let Global = JSON.parse(rawGlobal);\n");
                jsStr.append("let rawArg = fs.readFileSync('arg.json').toString();\n" + "let jsonArg = JSON.parse(rawArg);\n" + "let requester = jsonArg.requester;\n" + "let arg = jsonArg.arg;\n");
                jsStr.append("let srcStr = fs.readFileSync('contract.js').toString();\n");
                // verify signatures and decrypt all confidential variables Important!!!!!
                jsStr.append("for (var k in Global) {\n" + "  if (Global.hasOwnProperty(k)) {\n" + "    if (k.startsWith('conf_')) {\n" + "      let sig = Global[k].signature;\n" + "      let pubKey = Global[k].owner;\n" + "      let verifyResult = sm2.doVerifySignature(srcStr, sig, pubKey);\n" + "      if (verifyResult) {\n" + "        let newKey = k.substring(5);\n" + "        let decKey = Buffer.from(process.env['KEY_'+pubKey.substring(0,10).toUpperCase()], 'hex');\n" + "        let decIv = Buffer.from(Global[k].iv, 'hex');\n" + "        let cipherText = Buffer.from(Global[k].cipherText, 'hex');\n" + "        let decipher = crypto.createDecipheriv('aes-256-cbc', decKey, decIv);\n" + "        let decrypted = decipher.update(cipherText);\n" + "        decrypted = Buffer.concat([decrypted, decipher.final()]);\n" + "        let plaintext = decrypted.toString();\n" + "        Global[newKey] = plaintext;\n" + "      }\n" + "    }\n" + "  }\n" + "}\n");
                // call function
                jsStr.append("var ret = ").append(fn.functionName).append("(arg, requester, null);\n");
                // TODO: encrypt all confidential variables so state can be updated in confidential function @shujunyi
                // encrypt return value and write to a file
                jsStr.append("var retStr = JSON.stringify(ret);\n" + "var key = Buffer.from(process.env['KEY_'+requester.substring(0,10).toUpperCase()], 'hex');\n" + "var iv = crypto.randomBytes(16);\n" + "let cipher = crypto.createCipheriv('aes-256-cbc', key, iv); \n" + "let encRet =  cipher.update(retStr);\n" + "encRet = Buffer.concat([encRet, cipher.final()]);\n" + "let result = {iv: iv.toString('hex'), encryptedData: encRet.toString('hex')};\n" + "let resultStr = JSON.stringify(result);\n" + "fs.writeFileSync('result.json', resultStr);\n");
                // put script into Global so owner can send it and collect signatures
                Object som = ScriptObjectMirror.wrap(jsStr.toString(), global);
                globalVars.put("src_" + fn.functionName, som);
                break;
            }
        }
    }

    public static void copyTemplateToDestination(ContractRequest input) {
        String dest = CONFIDENTIAL_SCRIPT_PATH + File.separator + input.getRequestID();
        File srcDir = new File(CONFIDENTIAL_TEMPLATE_PATH);
        File destDir = new File(dest);
        try {
            FileUtils.copyDirectory(srcDir, destDir);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public static void dumpScriptAndStates(
            Invocable engine,
            FunctionNode functionNode,
            ContractRequest input,
            ScriptObjectMirror globalVars)
            throws IOException, ScriptException, NoSuchMethodException {
        Map<String, Object> globalMap = (Map<String, Object>) convertIntoJavaObject(globalVars);
        String dest = CONFIDENTIAL_SCRIPT_PATH + File.separator + input.getRequestID();
        // dump script
        String jsStr = globalMap.remove("src_" + functionNode.functionName).toString();
        String scriptPath = dest + File.separator + "contract.js";
        BufferedWriter writer = new BufferedWriter(new FileWriter(scriptPath));
        writer.write(jsStr);
        writer.close();
        // dump global variables
        globalMap.remove("Resources");
        String globalStr = JsonUtil.toJson(globalMap);
        String globalPath = dest + File.separator + "global.json";
        writer = new BufferedWriter(new FileWriter(globalPath));
        writer.write(globalStr);
        writer.close();
        // dump arg and requester
        Map<String, Object> argMap = new HashMap<>();
        argMap.put("arg", input.getArg());
        argMap.put("requester", input.getRequester());
        String argStr = JsonUtil.toJson(argMap);
        String argPath = dest + File.separator + "arg.json";
        writer = new BufferedWriter(new FileWriter(argPath));
        writer.write(argStr);
        writer.close();
        // generate manifest
        String manifestStr = generateGrapheneManifestStr(engine, input);
        String manifestPath = dest + File.separator + "nodejs.manifest.template";
        writer = new BufferedWriter(new FileWriter(manifestPath));
        writer.write(manifestStr);
        writer.close();
    }

    private static Set<String> findAllDependentFunctions(ContractNode cn, FunctionNode fn) {
        Set<String> dependentFunctions = fn.getDependentFunctions();
        for (String functionName : dependentFunctions) {
            for (FunctionNode functionNode : cn.getFunctions()) {
                if (functionNode.functionName.equals(functionName)) {
                    dependentFunctions.addAll(findAllDependentFunctions(cn, functionNode));
                }
            }
        }
        dependentFunctions.add(fn.functionName);
        return dependentFunctions;
    }

    private static String generateGrapheneManifestStr(Invocable engine, ContractRequest input) throws ScriptException, NoSuchMethodException {
        String manifestStr = "# Nodejs manifest file example\n" +
                "#\n" +
                "# This manifest was prepared and tested on Ubuntu 18.04.\n" +
                "\n" +
                "loader.argv0_override = \"nodejs\"\n" +
                "\n" +
                "# LibOS layer library of Graphene. There is currently only one implementation,\n" +
                "# so it is always set to libsysdb.so.\n" +
                "loader.preload = \"file:$(GRAPHENEDIR)/Runtime/libsysdb.so\"\n" +
                "\n" +
                "# Show/hide debug log of Graphene ('inline' or 'none' respectively).\n" +
                "loader.debug_type = \"$(GRAPHENEDEBUG)\"\n" +
                "\n" +
                "# Read application arguments directly from the command line. Don't use this on production!\n" +
                "loader.insecure__use_cmdline_argv = 1\n" +
                "\n" +
                "# Specify paths to search for libraries. The usual LD_LIBRARY_PATH syntax\n" +
                "# applies. Paths must be in-Graphene visible paths, not host-OS paths (i.e.,\n" +
                "# paths must be taken from fs.mount.xxx.path, not fs.mount.xxx.uri).\n" +
                "loader.env.LD_LIBRARY_PATH = \"/lib:/usr/lib:$(ARCH_LIBDIR):/usr/$(ARCH_LIBDIR):./\"\n" +
                "\n" +
                "# Mount host-OS directory to required libraries (in 'uri') into in-Graphene\n" +
                "# visible directory /lib (in 'path').\n" +
                "fs.mount.lib.type = \"chroot\"\n" +
                "fs.mount.lib.path = \"/lib\"\n" +
                "fs.mount.lib.uri = \"file:$(GRAPHENEDIR)/Runtime\"\n" +
                "\n" +
                "fs.mount.lib2.type = \"chroot\"\n" +
                "fs.mount.lib2.path = \"$(ARCH_LIBDIR)\"\n" +
                "fs.mount.lib2.uri = \"file:$(ARCH_LIBDIR)\"\n" +
                "\n" +
                "#fs.mount.lib3.type = \"chroot\"\n" +
                "#fs.mount.lib3.path = \"/usr/$(ARCH_LIBDIR)\"\n" +
                "#fs.mount.lib3.uri = \"file:/usr/$(ARCH_LIBDIR)\"\n" +
                "\n" +
                "fs.mount.usr.type = \"chroot\"\n" +
                "fs.mount.usr.path = \"/usr\"\n" +
                "fs.mount.usr.uri = \"file:/usr\"\n" +
                "\n" +
                "# Host-level directory to NSS files required by Glibc + NSS libs\n" +
                "fs.mount.etc.type = \"chroot\"\n" +
                "fs.mount.etc.path = \"/etc\"\n" +
                "fs.mount.etc.uri = \"file:/etc\"\n" +
                "\n" +
                "# Workload needs to create temporary files\n" +
                "fs.mount.tmp.type = \"chroot\"\n" +
                "fs.mount.tmp.path = \"/tmp\"\n" +
                "fs.mount.tmp.uri = \"file:/tmp\"\n" +
                "\n" +
                "# Set enclave size to 2GB; NodeJS expects around 1.7GB of heap on startup,\n" +
                "# see e.g. https://github.com/nodejs/node/issues/13018.\n" +
                "# Recall that SGX v1 requires to specify enclave size at enclave creation time.\n" +
                "sgx.enclave_size = \"2G\"\n" +
                "\n" +
                "# Set maximum number of in-enclave threads (somewhat arbitrarily) to 8. Recall\n" +
                "# that SGX v1 requires to specify the maximum number of simultaneous threads at\n" +
                "# enclave creation time.\n" +
                "sgx.thread_num = 16\n" +
                "\n" +
                "# Specify all libraries used by Node.js and its dependencies (including all libs\n" +
                "# which can be loaded at runtime via dlopen).\n" +
                "sgx.trusted_files.ld = \"file:$(GRAPHENEDIR)/Runtime/ld-linux-x86-64.so.2\"\n" +
                "sgx.trusted_files.libc = \"file:$(GRAPHENEDIR)/Runtime/libc.so.6\"\n" +
                "sgx.trusted_files.libm = \"file:$(GRAPHENEDIR)/Runtime/libm.so.6\"\n" +
                "sgx.trusted_files.libdl = \"file:$(GRAPHENEDIR)/Runtime/libdl.so.2\"\n" +
                "sgx.trusted_files.librt = \"file:$(GRAPHENEDIR)/Runtime/librt.so.1\"\n" +
                "sgx.trusted_files.libutil = \"file:$(GRAPHENEDIR)/Runtime/libutil.so.1\"\n" +
                "sgx.trusted_files.libpthread = \"file:$(GRAPHENEDIR)/Runtime/libpthread.so.0\"\n" +
                "sgx.trusted_files.libnssdns = \"file:$(GRAPHENEDIR)/Runtime/libnss_dns.so.2\"\n" +
                "sgx.trusted_files.libresolv = \"file:$(GRAPHENEDIR)/Runtime/libresolv.so.2\"\n" +
                "\n" +
                "sgx.trusted_files.libstdc = \"file:/usr/$(ARCH_LIBDIR)/libstdc++.so.6\"\n" +
                "sgx.trusted_files.libgccs = \"file:$(ARCH_LIBDIR)/libgcc_s.so.1\"\n" +
                "sgx.trusted_files.libaptpkg = \"file:/usr/$(ARCH_LIBDIR)/libapt-pkg.so.5.0\"\n" +
                "sgx.trusted_files.liblz4 = \"file:/usr/$(ARCH_LIBDIR)/liblz4.so.1\"\n" +
                "sgx.trusted_files.libsystemd = \"file:$(ARCH_LIBDIR)/libsystemd.so.0\"\n" +
                "sgx.trusted_files.libselinux = \"file:$(ARCH_LIBDIR)/libselinux.so.1\"\n" +
                "sgx.trusted_files.libgcrypt = \"file:$(ARCH_LIBDIR)/libgcrypt.so.20\"\n" +
                "sgx.trusted_files.libpcre = \"file:$(ARCH_LIBDIR)/libpcre.so.3\"\n" +
                "sgx.trusted_files.libgpgerror = \"file:$(ARCH_LIBDIR)/libgpg-error.so.0\"\n" +
                "sgx.trusted_files.libexpat = \"file:$(ARCH_LIBDIR)/libexpat.so.1\"\n" +
                "sgx.trusted_files.libz = \"file:$(ARCH_LIBDIR)/libz.so.1\"\n" +
                "sgx.trusted_files.libz2 = \"file:$(ARCH_LIBDIR)/libbz2.so.1.0\"\n" +
                "sgx.trusted_files.liblzma = \"file:$(ARCH_LIBDIR)/liblzma.so.5\"\n" +
                "sgx.trusted_files.libmpdec = \"file:/usr/$(ARCH_LIBDIR)/libmpdec.so.2\"\n" +
                "\n" +
                "# Name Service Switch (NSS) libraries (Glibc dependencies)\n" +
                "sgx.trusted_files.libnssfiles = \"file:$(ARCH_LIBDIR)/libnss_files.so.2\"\n" +
                "sgx.trusted_files.libnsscompat = \"file:$(ARCH_LIBDIR)/libnss_compat.so.2\"\n" +
                "sgx.trusted_files.libnssnis = \"file:$(ARCH_LIBDIR)/libnss_nis.so.2\"\n" +
                "sgx.trusted_files.libnsl = \"file:$(ARCH_LIBDIR)/libnsl.so.1\"\n" +
                "sgx.trusted_files.libnssmyhostname = \"file:$(ARCH_LIBDIR)/libnss_myhostname.so.2\"\n" +
                "sgx.trusted_files.libnssmdns = \"file:$(ARCH_LIBDIR)/libnss_mdns4_minimal.so.2\"\n" +
                "\n" +
                "# Scratch space\n" +
                "sgx.allowed_files.tmp = \"file:/tmp\"\n" +
                "\n" +
                "# APT config files\n" +
                "sgx.allowed_files.aptconfd = \"file:/etc/apt/apt.conf.d\"\n" +
                "sgx.allowed_files.aptconf = \"file:/etc/apt/apt.conf\"\n" +
                "sgx.allowed_files.apport = \"file:/etc/default/apport\"\n" +
                "\n" +
                "# Name Service Switch (NSS) files (Glibc reads these files)\n" +
                "sgx.allowed_files.nsswitch = \"file:/etc/nsswitch.conf\"\n" +
                "sgx.allowed_files.group = \"file:/etc/group\"\n" +
                "sgx.allowed_files.passwd = \"file:/etc/passwd\"\n" +
                "\n" +
                "# DNS hostname resolution files (Glibc reads these files)\n" +
                "sgx.allowed_files.hostconf = \"file:/etc/host.conf\"\n" +
                "sgx.allowed_files.hosts = \"file:/etc/hosts\"\n" +
                "sgx.allowed_files.gaiconf = \"file:/etc/gai.conf\"\n" +
                "sgx.allowed_files.resolv = \"file:/etc/resolv.conf\"\n" +
                "\n" +
                "sgx.allowed_files.openssl = \"file:/etc/ssl/openssl.cnf\"\n" +
                "\n" +
                "# System's file system table\n" +
                "sgx.allowed_files.fstab = \"file:/etc/fstab\"\n" +
                "\n" +
                "$(NODEJS_TRUSTED_LIBS)\n" +
                "\n" +
                "# JavaScript (trusted)\n" +
                "sgx.allowed_files.smlib = \"file:node_modules\"\n" +
                "sgx.trusted_files.npminfo = \"file:package.json\"\n" +
                "sgx.trusted_files.contract = \"file:contract.js\"\n" +
                "sgx.trusted_files.globaljson = \"file:global.json\"\n" +
                "sgx.trusted_files.argjson = \"file:arg.json\"\n" +
                "\n" +
                "sys.insecure__allow_eventfd = 1\n" +
                "\n" +
                "sgx.remote_attestation = 1\n" +
                "\n" +
                "loader.env.LD_PRELOAD = \"libsecret_prov_attest.so\"\n" +
                "loader.env.SECRET_PROVISION_CONSTRUCTOR = \"1\"\n" +
                "loader.env.SECRET_PROVISION_SET_PF_KEY = \"1\"\n" +
                "loader.env.SECRET_PROVISION_CA_CHAIN_PATH = \"certs/test-ca-sha256.crt\"\n" +
                "loader.env.SECRET_PROVISION_SERVERS = \"localhost:4433\"\n" +
                "\n" +
                "sgx.trusted_files.libsecretprovattest = \"file:libsecret_prov_attest.so\"\n" +
                "sgx.trusted_files.cachain = \"file:certs/test-ca-sha256.crt\"\n" +
                "\n" +
                "# Specify your SPID and linkable/unlinkable attestation policy\n" +
                "sgx.ra_client_spid = \"DF3A8BA098E93F66CC64E8A215E98333\"\n" +
                "sgx.ra_client_linkable = 0\n";
        // add secret servers
        manifestStr += "loader.env.SECRET_PROVISION_CC_SERVERS = ";
        Object resultStr = engine.invokeFunction(
                "getAllSecret",
                "",
                input.getRequester(),
                input.getRequesterDOI());
        Map<String, String> resultMap = JsonUtil.fromJson(resultStr.toString(), MapType);
        Map<String, String> serverMap = JsonUtil.fromJson(resultMap.get("result"), MapType);
        List<String> entries = new ArrayList<>();
        for (Map.Entry<String, String> entry : serverMap.entrySet()) {
            String key = entry.getKey();
            String server = entry.getValue();
            String envVar = "KEY_" + key.substring(0, 10).toUpperCase();
            entries.add(envVar + "=" + server);
        }
        manifestStr += "\"" + String.join(";", entries) + "\"";
        return manifestStr;
    }

    private static Object convertIntoJavaObject(Object scriptObj) {
        if (scriptObj instanceof ScriptObjectMirror) {
            ScriptObjectMirror scriptObjectMirror = (ScriptObjectMirror) scriptObj;
            if (scriptObjectMirror.isArray()) {
                List<Object> list = new ArrayList<>();
                for (Map.Entry<String, Object> entry : scriptObjectMirror.entrySet()) {
                    list.add(convertIntoJavaObject(entry.getValue()));
                }
                return list;
            } else {
                Map<String, Object> map = new HashMap<>();
                for (Map.Entry<String, Object> entry : scriptObjectMirror.entrySet()) {
                    map.put(entry.getKey(), convertIntoJavaObject(entry.getValue()));
                }
                return map;
            }
        } else {
            return scriptObj;
        }
    }
}
