#!/usr/bin/env python2
import sys,os
import subprocess
import re

def gen_dockerfile( image_name, app_name, bin_name, proj_dir):
  if not os.path.exists(proj_dir + '/Tools/build'):
    os.makedirs(proj_dir + '/Tools/build')
  df =open(proj_dir + '/Tools/build/Dockerfile' + '.' + app_name, 'w')
  df.write('# This file is auto-generated, any edits will be overwritten\n')

  df.write('\n')
  # Choose the base image from the user input
  df.write('FROM '+ image_name + '\n')
  df.write('\n')

  # SWITCH to ROOT
  df.write('# SWITCH to root\n')
  df.write('USER root\n\n')

  # DOWNLOAD dependencies
  df.write('# Download dependencies\n')
  df.write('RUN apt-get update && \\\n')
  df.write('    apt-get install -y openssl libjemalloc-dev python python-pip python-dev\n')
  df.write('RUN pip install protobuf && \\\n')
  df.write('    pip install pycrypto\n')

  df.write('# Temporal fixes for Dependencies Issue #1: libcrypto.so.1.0.0 and libssl.so.1.0.0 have different locations\n')
  if not os.path.isfile('/lib/x86_64-linux-gnu/libcrypto.so.1.0.0'):
    df.write('RUN ln -s /usr/lib/x86_64-linux-gnu/libcrypto.so.1.0.0 /lib/x86_64-linux-gnu/libcrypto.so.1.0.0 \\\n')
  if not os.path.isfile('/lib/x86_64-linux-gnu/libssl.so.1.0.0'):
    df.write('RUN ln -s /usr/lib/x86_64-linux-gnu/libssl.so.1.0.0 /lib/x86_64-linux-gnu/libssl.so.1.0.0\n')

  # SETUP Directory Structure
  df.write('# Setup Directory Structure\n')
  df.write('RUN mkdir -p ' + proj_dir + '/LibOS/shim/test/apps/' + app_name + ' \\\n')
  df.write(' && mkdir -p ' + proj_dir + '/Pal/src/host/Linux-SGX/signer \\\n')
  df.write(' && mkdir -p ' + proj_dir + '/Runtime \\\n')
  df.write(' && mkdir /gbin\n')

  # COPY system files
  df.write('# Copy system files\n')
  df.write('COPY Runtime/* ' + proj_dir + '/Runtime/\n')
  df.write('COPY Pal/src/Makefile.Host ' + proj_dir + '/Pal/src/Makefile.Host\n')
  df.write('COPY Pal/src/host/Linux-SGX/signer/* ' + proj_dir + '/Pal/src/host/Linux-SGX/signer/\n')

  # COPY tools for building app instance
  df.write('# Copy tools for building app instance\n')
  df.write('COPY Tools/build/tools/* /gbin/\n')
  df.write('COPY Tools/gen_manifest /gbin/\n')

  # Generating manifest file for target app
  df.write('# Generating manifest for target app\n')
  df.write('RUN /gbin/gen_manifest ' + app_name + ' ' + bin_name + ' ' + proj_dir + '\n')

  # Sign Enclave
  df.write('# Signing Enclave\n')
  df.write('RUN cd ' + proj_dir + '/LibOS/shim/test/apps/' + app_name + ' && \\\n'
           '    '+ proj_dir + '/Pal/src/host/Linux-SGX/signer/pal-sgx-sign -libpal ' + proj_dir +
           '/Pal/src/host/Linux-SGX/../../../../Runtime/libpal-Linux-SGX.so -key ' + proj_dir +
           '/Pal/src/host/Linux-SGX/signer/enclave-key.pem -output ' + app_name + '.manifest.sgx ' +
           '-manifest ' + app_name + '.manifest\n')
  # Remove signing key
  df.write('# Removing key after signing\n')

  # Overwrite Entry Point
  df.write('ENTRYPOINT  ["/bin/bash", "/gbin/app_exec"]\n')
  df.close()

def make_executable(path):
  mode = os.stat(path).st_mode
  mode |= (mode & 0o444) >> 2    # copy R bits to X
  os.chmod(path, mode)

def gen_app_executor(app_name, bin_cmd, proj_dir):
  if not os.path.exists(proj_dir + '/Tools/build/tools'):
    os.makedirs(proj_dir + '/Tools/build/tools')

  e_path = proj_dir + '/Tools/build/tools/app_exec'
  ef = open(e_path, 'w')
  make_executable(e_path)

  ef.write('#!/usr/bin/env bash\n\n')
  ef.write('cd ' + proj_dir + '/LibOS/shim/test/apps/' + app_name +'\n')
  ef.write('# Generate EINITOKEN\n')
  ef.write(proj_dir + '/Pal/src/host/Linux-SGX/signer/pal-sgx-get-token -output '
           + app_name + '.token -sig ' + app_name + '.sig\n')
  ef.write('# Run the application\n')
  ef.write('SGX=1 ./' + app_name + '.manifest.sgx ' + bin_cmd + '\n')

  ef.close()

if __name__ == '__main__':
  if len(sys.argv) < 3:
    print('Usage: gsce run [Image name]')
    exit()

  image_name = sys.argv[-1]
  image_match = re.match(r'([^:]*)(:*)(.*)', image_name)
  if image_match:
    app_name = image_match.group(1)

  # application name may contain '/', remove it
  app_name = app_name.split('/')[-1]

  inspect_cmd = 'sudo docker inspect --format \'{{.Config.Cmd}}\' ' + image_name
  res = subprocess.check_output(inspect_cmd, shell=True).strip()

  # Docker image may execute '[/bin/sh -c command]', replace with just '[command]'
  if res.startswith('[/bin/sh -c '):
    res = '[' + res[len('[/bin/sh -c '):]

  match = re.match(r'\[([^\s]*)\s*(.*)\]', res)
  bin_name = match.group(1)
  bin_cmd = ''
  if match.group(2):
    bin_cmd = match.group(2)

  # Store the rest arguments as Docker run arguments
  docker_str = ' ' + ' '.join(sys.argv[2:-1])

  # print image_cmd
  proj_dir = os.path.abspath(os.getcwd() + '/../')

  # STEP 1: Generating Dockerfile
  gen_dockerfile(image_name, app_name, bin_name, proj_dir)

  # STEP 2: Generating entry point execute script
  gen_app_executor(app_name, bin_cmd, proj_dir)

  # STEP 3: Building new docker image with generated Dockerfile
  os.chdir('..')
  os.system('sudo docker build -f Tools/build/Dockerfile.' + app_name + ' -t gsc_' + app_name + ' .\n')

  # STEP 4: Run GSC with the target app
  os.system('sudo docker run -i -t' + docker_str +' --device=/dev/gsgx --device=/dev/isgx ' +
            '-v /var/run/aesmd/aesm.socket:/var/run/aesmd/aesm.socket gsc_' + app_name + '\n')
