#!/usr/bin/env python

from __future__ import print_function
import argparse
import os, os.path as osp
from tempfile import mkdtemp
from subprocess import check_output, STDOUT, CalledProcessError
from uuid import uuid4
from hashlib import md5
from itertools import chain
import sys
import tarfile
import re
import json


#the script we run inside the container to probe used files
guest_script=r"""
#!/bin/sh

parse_strace() {
    grep -oP "\(\".*?\"," | cut -c 3- | rev | cut -c 3- | rev | sort -u
}

# first strace the command to get all the files it uses
# then strace a readlink of all the files to resolve the entire symlink tree for all of them
/.strace -f -e trace=file -o /.strace.out "$@"
cat /.strace.out | parse_strace | xargs /.strace -f -e trace=lstat readlink -f 2>&1 | parse_strace > /.usedfiles
"""

def sh(*args,**kwargs):
    return check_output(stderr=STDOUT,*args,**kwargs).strip()


parser = argparse.ArgumentParser(prog='docker-stfd')
parser.add_argument('-t','--tag',help='what to tag the slimmed image (default: <image>-slim)')
parser.add_argument('--entrypoint',help='overwrite the default ENTRYPOINT of the image')
parser.add_argument('-v', action='append', dest='volumes',default=[],help='bind mount a volume')
parser.add_argument('image',help='the image name')
parser.add_argument('cmd',nargs='*',help='the command to run (which should access all necessary files)')
if '--' in sys.argv:
    isplit=sys.argv.index('--')
    argv,extrav=sys.argv[1:isplit],sys.argv[isplit+1:]
else:
    argv,extrav=sys.argv[1:],[]
args,unknown = parser.parse_known_args(argv)
args.cmd += (unknown+extrav)
if ':' not in args.image: args.image += ':latest'
if args.tag is None: args.tag=args.image+'-slim'


try:

    tmpdir = mkdtemp()
    cname = uuid4().hex

    #get reverse engineered Dockerfile
    print("Reverse engineering Dockerfile...")

    def fixline(l):
        """workaround for https://github.com/CenturyLinkLabs/dockerfile-from-image/issues/13"""
        l=re.sub(r'ENTRYPOINT &{(.*)}',r'ENTRYPOINT \1',l)
        if any([l.startswith(x) for x in ['CMD','ENTRYPOINT']]):
            l = re.sub(r'\[(.*)\]',lambda m: '[%s]'%','.join(['"%s"'%x for x in m.group(0).split('"')[1::2]]),l)
        return l

    dockerfile = sh([
        "docker","run","--rm",
        "-v","/var/run/docker.sock:/var/run/docker.sock",
        "centurylink/dockerfile-from-image",args.image
    ]).splitlines()
    dockerfile = map(fixline,dockerfile)
    baseimage = next(l for l in dockerfile if l.startswith('FROM')).split()[1]


    #get used files
    print("Inspecting container file usage...")
    if args.entrypoint:
        entrypoint = args.entrypoint.split()
    else:
        entrypoint = map(str,json.loads(sh(["docker","inspect","-f","{{json .Config.Entrypoint}}",args.image])))
        
    stfdfile=osp.join(tmpdir,'stfd')
    with open(stfdfile,'w') as f: f.write(guest_script)
    os.chmod(stfdfile,0o755)
    
    usedfiles=osp.join(tmpdir,'usedfiles')
    sh(['touch',usedfiles])
    
    cmd = ["docker","run","--rm","--privileged",
           "-v",sh(["which","strace"])+":/.strace",
           "-v",stfdfile+":/stfd",
           "-v",usedfiles+":/.usedfiles"]
    for v in args.volumes: 
        cmd += ['-v',v]
    cmd += ["--entrypoint","sh",args.image,"-c","/stfd "+' '.join('"%s"'%x for x in entrypoint+args.cmd)]
    print(sh(cmd))
    usedfiles = open(usedfiles).read().splitlines()
    
    #get image filesystem
    print("Getting image filesystem...")
    sh(["docker","run","--name",cname,"--entrypoint","sh",args.image,"-c","true"])
    sh(["docker","export","-o",osp.join(tmpdir,"slim.tar"),cname])
    sh(["docker","rm",cname])

    #get base filesystem
    print("Getting base filesystem...")
    sh(["docker","run","--name",cname,"--entrypoint","sh",baseimage,"-c","true"])
    sh(["docker","export","-o",osp.join(tmpdir,"base.tar"),cname])
    sh(["docker","rm",cname])


    #create tar with files in slim not already in base
    print("Creating new slimmed filesystem...")
    os.makedirs(osp.join(tmpdir,"build"))
    with tarfile.open(osp.join(tmpdir,"base.tar")) as base:
        with tarfile.open(osp.join(tmpdir,"slim.tar")) as slim:
            with tarfile.open(osp.join(tmpdir,"build","rootfs.tar"),'w') as rootfs:
                for m in slim.getmembers():
                    #if file doesn't exist, symlink was changed, or contents was changed
                    if (('/'+m.name in usedfiles)
                        and 
                        ((m.name not in base.getnames())
                         or (m.issym() and m.linkpath != base.getmember(m.name).linkpath)
                         or (m.isfile() and md5(slim.extractfile(m).read()).digest() != md5(base.extractfile(m.name).read()).digest())
                       )):
                        rootfs.addfile(m,slim.extractfile(m) if m.isfile() else None)


    #build new container
    print("Building new container...")
    dockerfile = [l for l in dockerfile if not any([l.startswith(x) for x in ['COPY','ADD','RUN']])]
    dockerfile += ['ADD rootfs.tar /']
    with open(osp.join(tmpdir,"build","Dockerfile"),'w') as df:  df.write('\n'.join(dockerfile))
    sh(["docker","build","-t",args.tag,osp.join(tmpdir,"build")])


    #print stats
    def get_size(image):
        return sh(["docker","images","--format","{{.Size}}",image])
    print("Uncompressed size shrunk from %s to %s"%(get_size(args.image),get_size(args.tag)))
    print("To test, run:")
    cmd = ["docker","run","--rm"]
    if args.entrypoint: cmd += ["--entrypoint",args.entrypoint]
    for v in args.volumes: cmd += ["-v",v]
    cmd += [args.tag]
    for x in args.cmd: cmd += ['"%s"'%x]
    print(" ".join(cmd))

except KeyboardInterrupt:
    print("Cleaning up temporary files...")
except CalledProcessError as e:
    print("Error calling: "+str(e.cmd))
    print(e.output.strip())
finally:
    #cleanup
    try:
        sh(["rm","-rf",tmpdir])
    except:
        pass
    try:
        sh(["docker","inspect",cname]) #will fail if container doesn't exist
        sh(["docker","rm","-f",cname])
    except:
        pass
