#include <stdio.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <sys/stat.h>
#include <sys/sysmacros.h>
#include <misc.h>
#include <tlmplib.h>
#include <string>
#include <set>
#include <map>

using namespace std;

static const char *rootdir = "";	// root insert as the source
static const char *varlxc = "/var/lib/lxc";
/*
	Create a directory hierarchy
*/
static void create_dirs (const string &dir, set<string> &dirs, const string &rootfs)
{
	if (dir != "/" && dirs.count(dir)==0){
		const char *ptdir = dir.c_str();
		const char *last = strrchr(ptdir,'/');
		if (last != NULL && last > ptdir){
			string basedir(ptdir,last-ptdir);
			create_dirs(basedir,dirs,rootfs);
		}
		struct stat64 st;
		if (stat64(ptdir,&st)==-1){
			tlmp_error ("Stat directory %s failed (%s)\n",ptdir,strerror(errno));
		}else{
			printf ("mkdir -m %0o -p %s%s\n",st.st_mode&07777,rootfs.c_str(),ptdir);
			if (st.st_uid != 0 || st.st_gid != 0) printf ("chown %d:%d %s%s\n",st.st_uid,st.st_gid,rootfs.c_str(),ptdir);
			if (strcmp(ptdir,"/tmp")==0) printf ("chmod 1777 %s%s\n",rootfs.c_str(),ptdir);
		}
		dirs.insert(dir);
	}
}

static const char *ln_cmd = "ln";
static const char *cp_cmd = "cp -p";
static FILE *ffiles = NULL;

static int create_dir_from_file (const string &file, string &dir, set<string> &dirs, const string &rootfs)
{
	int ret = -1;
	const char *ptfile = file.c_str();
	const char *fname = strrchr(ptfile,'/');
	if (fname != NULL){
		dir = string(ptfile,fname-ptfile);
		create_dirs (dir,dirs,rootfs);
		ret = 0;
	}
	return ret;
}
static bool use_link = false;
static vector<string> excludedirs;	// Directories excluded
/*
	This either copies a file or make a link. The link mode should offer higher performance
	because multiple container may share the same binaries/libraries.
	But it also means that one "infected/rooted" may infect other containers by altering its
	own binaries. So this mode must be used with mechanism such as frozen files. Note that if all
	binaries in a container are "busy" all the time, then they can't be modified.
*/
static void link_file (const string &cwd, const string &rootfs, string file, set<string> &dirs)
{
	if (file[0] != '/'){
		file = string_f("%s/%s",cwd.c_str(),file.c_str());
	}
	const char *ptfile = file.c_str();
	bool excluded = false;
	for (auto &d:excludedirs){
		if (strncmp(ptfile,d.c_str(),d.size())==0
			&& ptfile[d.size()] == '/'){
			excluded = true;
			break;
		}
	}
	if (!excluded
		&& strcmp(ptfile,"/etc/passwd")!=0
		&& strcmp(ptfile,"/etc/shadow")!=0
		&& strcmp(ptfile,"/etc/group")!=0
		&& strcmp(ptfile,"/etc/hosts")!=0
		&& strcmp(ptfile,"/etc/conf.linuxconf")!=0
		&& strncmp(ptfile,"/var/lib/sss/mc/",16)!=0
		&& strncmp(file.c_str(),"/proc/",6)!=0
		&& dirs.insert(file).second){
		// Create the directory hierarchy as needed
		string dir;
		if (create_dir_from_file (file,dir,dirs,rootfs) != -1){
			const char *cmd = ln_cmd;
			if (strncmp(ptfile,"/usr/",5)!=0
				&& strncmp(ptfile,"/lib/",5)!=0
				&& strncmp(ptfile,"/bin/",5)!=0
				&& strncmp(ptfile,"/sbin/",6)!=0
				&& strncmp(ptfile,"/lib64/",7)!=0){
				cmd = cp_cmd;
			}
			struct stat64 st;
			if (stat64(ptfile,&st)==-1){
				if (errno == EACCES){
					printf ("%s %s%s %s%s\n",cp_cmd,rootdir,ptfile,rootfs.c_str(),ptfile);
				}else{
					tlmp_error ("Can't stat file %s (%s)\n",ptfile,strerror(errno));
				}
			}else if (S_ISDIR(st.st_mode)){
				create_dirs (file,dirs,rootfs);
			}else if (S_ISFIFO(st.st_mode)){
				tlmp_error ("link_file %s: fifo not supported\n",ptfile);
			}else if (S_ISLNK(st.st_mode)){
				tlmp_error ("link_file %s: symbolic link\n",ptfile);
			}else if (S_ISSOCK(st.st_mode)){
				tlmp_error ("link_file %s: unix socket\n",ptfile);
			}else if (S_ISCHR(st.st_mode)
				|| S_ISBLK(st.st_mode)){
				char letter = S_ISCHR(st.st_mode) ? 'c' : 'b';
				printf ("mknod %s%s %c %d %d\n",rootfs.c_str(),ptfile,letter,major(st.st_rdev),minor(st.st_rdev));
				if (st.st_uid != 0 || st.st_gid != 0) printf ("chown %d:%d %s%s\n",st.st_uid,st.st_gid,rootfs.c_str(),ptfile);
				printf ("chmod %0o %s%s\n",st.st_mode & 0777,rootfs.c_str(),ptfile);
			}else if (S_ISREG(st.st_mode)){
				if (ffiles != NULL) fprintf (ffiles,"%s %s/%s\n",ptfile,rootfs.c_str(),ptfile);
				if (use_link){
					// When performing a link on a file, if the file is a symlink
					// we have to follow the link to perform it on the target
					// instead of just the link.
					char buf[PATH_MAX];
					ssize_t len = readlink(ptfile,buf,sizeof(buf)-2);
					if (len > 0){
						buf[len] = '\0';
						if (buf[0] == '/'){
							printf ("%s %s%s %s%s\n",cmd,rootdir,buf,rootfs.c_str(),ptfile);
						}else{
							printf ("%s %s%s/%s %s%s\n",cmd,rootdir,dir.c_str(),buf,rootfs.c_str(),ptfile);
						}
					}else{
						printf ("%s %s%s %s%s\n",cmd,rootdir,ptfile,rootfs.c_str(),ptfile);
					}
				}else{
					printf ("%s %s%s %s%s\n",cmd,rootdir,ptfile,rootfs.c_str(),ptfile);
				}
			}
		}
	}
}

<mod>
static void link_prog (const string &cwd, const string &rootfs, const string &prog, set<string> &dirs, bool only_dll)
{
	glocal const string *cwd = &cwd;
	glocal const string *rootfs = &rootfs;
	glocal set<string> *dirs = &dirs;
	string cmd = string_f ("ldd %s",prog.c_str());
	<call walkpopen>(cmd.c_str(),5);
	<f oneline>
		vector<string> tb;
		int n = str_splitline(line,' ',tb);
		if (n == 4 && tb[1] == "=>"){
			link_file (*glocal.cwd,*glocal.rootfs,tb[2],*glocal.dirs);
		}
		return 0;
	</f>
	</call>
	if (!only_dll) link_file (cwd,rootfs,prog,dirs);
}
</mod>

static string getstrarg(
	const char *pt,
	int &result,
	const char *&ptarg)
{
	string ret;
	result = -1;
	ptarg = "";
	if (*pt == '"'){
		pt++;
		while (*pt != '"'){
			ret += *pt++;
		}
		if (*pt == '"'){
			pt++;
			ptarg = pt;
			pt = strchr(pt,')');
			if (pt != NULL){
				pt = str_skip(pt+1);
				if (*pt == '='){
					pt = str_skip(pt+1);
					result = atoi(pt);
				}
			}
		}
	}
	return ret;
}
static string getstrarg(
	const char *pt,
	bool &success,
	const char *&ptarg)
{
	int result;
	success = false;
	string file = getstrarg(pt,result,ptarg);
	if (result != -1) success = true;
	return file;
}

static void savefile_init (FILE *fout, const char *name)
{
	fprintf (fout,"#!/bin/sh\n");
	fprintf (fout,"cd %s/%s || exit 1\n",varlxc,name);
	fprintf (fout,"mkdir -p data\n");
	fchmod (fileno(fout),0755);
}

struct PIDINFO {
	string cwd;
	string lastopen;	// Start of open syscall with <unfinished ...>
	string lastclone;	// Start of clone syscall with <unfinished ...>
	string lastexecve;	// Start of execve syscall with <unfinished ...>
};

static const char *basename(PARAM_STRING path)
{
	const char *fname = strrchr(path.ptr,'/');
	if (fname == NULL){
		fname = path.ptr;
	}else{
		fname++;
	}
	return fname;
}
	

<mod>
int main (int argc, char *argv[])
{
	glocal int ret = -1;
	glocal const char *name = NULL;
	glocal vector<string>logfiles;
	glocal const char *initprog = "/usr/sbin/trli-init";
	glocal const char *program = NULL;
	glocal vector<string> extraprogs;	// Extra program to include in the container, usually to help debug
	glocal vector<string> extradirs;	// Extra directory to include in the container
	glocal const char *savefile = NULL;
	glocal const char *restorefile = NULL;
	glocal const char *filelist = NULL;
	glocal vector<string> preserve;		// List of files that must go in the save and restore file
						// even if they are not part of the strace log.
	glocal.ret = <call tlmpprogram>(argc,argv);
	<f init>
		setproginfo ("trli-lxc0",VERSION,"Setup execution parameters for LXC zero mode");
		setarg ('i',"initprog","Program used to start the other",glocal.initprog,true);
		setarg ('l',"stracelog","Log file produced by strace",glocal.logfiles,true);
		setarg ('n',"progname","Name of the program",glocal.name,true);
		setarg ('p',"program","Program path",glocal.program,true);
		setarg ('e',"extranprog","Extra programs",glocal.extraprogs,false);
		setarg ('d',"extradir","Extra directories",glocal.extradirs,false);
		setarg ('x',"exclude","Exclude files in this directory",excludedirs,false);
		setarg ('P',"preserve","Preserve this file",glocal.preserve,false);
		setgrouparg ("Files");
		setarg (' ',"savefile","Script to save statfile from container",glocal.savefile,false);
		setarg (' ',"restorefile","Script to restore files to container",glocal.restorefile,false);
		setarg (' ',"filelist","Write all file names to this file",glocal.filelist,false);
		setgrouparg ("Misc.");
		setarg (' ',"link","Perform links instead of copies",use_link,false);
		setarg (' ',"rootdir","Insert rootdir to the source path",rootdir,false);
	</f>
	<f main_noarg>
		glocal set<string> dirs;
		glocal set<string> files;
		glocal string rootfs = string_f("/var/lib/lxc/%s/rootfs",glocal.name);
		glocal map<int,PIDINFO> pids;
		glocal map<int,string> dirfd;	// Directory associated with a file handle (see openat() below)

		if (glocal.filelist != NULL){
			ffiles = fopen (glocal.filelist,"w");
			if (ffiles == NULL){
				tlmp_error ("Can't open file list %s (%s), aborting\n",glocal.filelist,strerror(errno));
				exit (-1);
			}
		}
		glocal string cwd;
		{
			char cwd[PATH_MAX];
			if (getcwd(cwd,sizeof(cwd))!=NULL){
				glocal.cwd = cwd;
			}
		}
		int ret = -1;
		if (use_link){
			ln_cmd = "ln";
		}else{
			ln_cmd = "cp -p";
		}
		// Some files we do not need in the container
		glocal.files.insert ("/etc/resolv.conf");
		//glocal.files.insert ("/etc/ld.so.cache");
		create_dirs ("/sbin",glocal.dirs,glocal.rootfs);
		create_dirs ("/tmp",glocal.dirs,glocal.rootfs);
		create_dirs ("/dev",glocal.dirs,glocal.rootfs);
		printf ("%s %s %s/sbin/init\n",ln_cmd,glocal.initprog,glocal.rootfs.c_str());
		link_prog (glocal.cwd,glocal.rootfs,glocal.initprog,glocal.dirs,true);
		for (auto x:glocal.extraprogs) link_prog (glocal.cwd,glocal.rootfs,x,glocal.dirs,false);
		for (auto x:glocal.extradirs) create_dirs (x,glocal.dirs,glocal.rootfs);
		link_file (glocal.cwd,glocal.rootfs,glocal.program,glocal.dirs);
		{
			static const char *tb[]={
				"/dev/null","/dev/zero","/dev/urandom","/lib64/ld-linux-x86-64.so.2"
			};
			for (auto x:tb) link_file (glocal.cwd,glocal.rootfs,x,glocal.dirs);
		}
		for (auto f:glocal.logfiles){
		<call loadfile>(f.c_str(),true);
		<f oneline>
			debug_printf ("line=%s\n",line);
			int pid = atoi(line);
			auto pid_it = glocal.pids.find(pid);
			string cwd = glocal.cwd;
			string lastopen;
			string lastclone;
			string lastexecve;
			if (pid_it == glocal.pids.end()){
				glocal.pids[pid].cwd = glocal.cwd;
			}else{
				cwd = pid_it->second.cwd;
				lastopen = pid_it->second.lastopen;
				lastclone = pid_it->second.lastclone;
				lastexecve = pid_it->second.lastexecve;
			}
			const char *pt = str_skipdig(line);
			pt = str_skip(pt);
			const char *resumed = strstr(pt,"<... open resumed>");
			string tmp;
			if (resumed != NULL){
				tmp = lastopen + (resumed+18);
				pt = tmp.c_str();
				//fprintf (stderr,"newpt =%s\n",pt);
			}else{
				resumed = strstr(pt,"<... clone resumed>");
				if (resumed != NULL){
					tmp = lastclone + (resumed+19);
					pt = tmp.c_str();
				}else{
					resumed = strstr(pt,"<... openat resumed>");
					if (resumed != NULL){
						tmp = lastopen + (resumed+20);
						pt = tmp.c_str();
					}else{
						resumed = strstr(pt,"<... execve resumed>");
						if (resumed != NULL){
							tmp = lastexecve + (resumed+20);
							pt = tmp.c_str();
debug_printf ("execve :%s:\n",pt);
						}
					}
				}
			}
			if (strncmp(pt,"open(",5)==0 || strncmp(pt,"openat(",7)==0){
				const char *unfin = strstr(pt,"<unfinished ...>");
				if (unfin != NULL){
					glocal.pids[pid].lastopen = string(pt,unfin-pt);
				}else{
					bool success = false;
					int result = -1;
					string file;
					if (strncmp(pt,"openat(",7)==0){
						pt += 7;
						if (isdigit(*pt)){
							int fd = atoi(pt);
							pt = str_skipdig(pt);
							if (*pt == ',' && pt[1] == ' '){
								pt += 2;
								file = getstrarg (pt,result,pt);
								cwd = glocal.dirfd[fd];
							}
						}else if (strncmp(pt,"-1, ",4)==0){
							file = getstrarg (pt+4,result,pt);
						}else if (strncmp(pt,"AT_FDCWD, ",10)!=0){
							tlmp_error ("openat, only AT_FDCWD supported for now, ending\n");
							exit (-1);
						}else{
							pt += 10;
							if (*pt != '"'){
								tlmp_error ("openat, no \" found, ending\n");
								exit (-1);
							}else{
								file = getstrarg (pt,result,pt);
							}
						}
						if (result != -1){
							success = true;
							if (file[0] != '/') file = string_f ("%s/%s",cwd.c_str(),file.c_str());
							glocal.dirfd[result] = file;
						}
					}else{
						file = getstrarg (pt+5,success,pt);
					}
					//fprintf (stderr,"file=%s pt=%s\n",file.c_str(),pt);
					if (file.size() > 0){
						if (glocal.files.count(file) == 0){
							pt = str_skip(pt);
							if (*pt == ','){
								pt = str_skip(pt+1);
								bool is_read = false;
								bool is_write = false;
								bool is_dir = false;
								while (*pt != '\0' && *pt != ')' && *pt != ',' && *pt != '<'){
									if (strncmp(pt,"O_RDONLY",8)==0){
										is_read = true;
									}else if (strncmp(pt,"O_WRONLY",8)==0){
										is_write = true;
									}else if (strncmp(pt,"O_RDWR",6)==0){
										is_write = true;
									}else if (strncmp(pt,"O_DIRECTORY",11)==0){
										is_dir = true;
									}
									while (isalpha(*pt) || *pt == '_') pt++;
									if (*pt == '|') pt++;
									pt = str_skip(pt);
								}
								while (*pt != '\0' && *pt != ')') pt++;
								debug_printf ("open %s %d '%c'\n",file.c_str(),is_write,*pt);
								if (!is_dir && *pt == ')' && success){
									glocal.files.insert(file);
									if (is_read){
										link_file (cwd,glocal.rootfs,file,glocal.dirs);
									}else if (is_write){
										if (file[0] != '/'){
											file = string_f("%s/%s",cwd.c_str(),file.c_str());
										}
										if (strncmp(file.c_str(),"/tmp/apr-tmp",12)==0){
											// do nothing with these tmp files
										}else if (strncmp(file.c_str(),"/proc/",6)!=0 && glocal.dirs.insert(file).second){
											// Create the directory hierarchy as needed
											string dir;
											debug_printf ("Create directory for %s\n",file.c_str());
											create_dir_from_file (file,dir,glocal.dirs,glocal.rootfs);
											const char *fname = strrchr(file.c_str(),'/');
											if (fname != NULL){
												fname++;
												// Check if it is a pid file. It has no value
												// and sometime confuse the process when already there
												const char *pidstr = strstr(fname,".pid");
												if (pidstr == NULL || pidstr[4] != '\0'){
													glocal.preserve.push_back(file);
												}
											}
										}
									}
								}
							}
						}
					}
				}
			}else if (strncmp(pt,"chdir(",6)==0){
				bool success;
				string dir = getstrarg (pt+6,success,pt);
				if (success){
					debug_printf ("chdir %d %s\n",pid,dir.c_str());
					glocal.pids[pid].cwd= dir;
				}
			}else if (strncmp(pt,"execve(",7)==0){
				const char *unfin = strstr(pt,"<unfinished ...>");
				if (unfin != NULL){
					glocal.pids[pid].lastexecve = string(pt,unfin-pt);
				}else{
					bool success;
					string prog = getstrarg (pt+7,success,pt);
					if (success){
						debug_printf ("execve %s\n",prog.c_str());
						link_file (cwd,glocal.rootfs,prog,glocal.dirs);
					}
				}
			}else if (strncmp(pt,"stat(",5)==0){
				bool success;
				string file = getstrarg (pt+5,success,pt);
				if (success){
					link_file (cwd,glocal.rootfs,file,glocal.dirs);
				}
			}else if (strncmp(pt,"lstat(",6)==0){
				bool success;
				string file = getstrarg (pt+6,success,pt);
				if (success){
					link_file (cwd,glocal.rootfs,file,glocal.dirs);
				}
			}else if (strncmp(pt,"clone(",6)==0){
				const char *unfin = strstr(pt,"<unfinished ...>");
				if (unfin != NULL){
					glocal.pids[pid].lastclone = string(pt,unfin-pt);
				}else{
					const char *pteq = strstr(pt," = ");
					if (pteq == NULL){
						tlmp_error ("Can't locate clone() result: %s\n",line);
					}else{
						pteq = str_skip(pteq+3);;
						int newpid = atoi(pteq);
						glocal.pids[newpid].cwd = cwd;
					}
				}	
			}
			return 0;
		</f>
		</call>
		}
		for (auto x:glocal.pids){
			debug_printf ("pid=%d cwd=%s\n",x.first,x.second.cwd.c_str());
		}
		glocal set<string> savefiles;
		// Remove duplicates
		for (auto &f:glocal.preserve) glocal.savefiles.insert(f);
		if (glocal.savefile != NULL){
			<call savefile>(glocal.savefile,false);
			<f dowrite>
				savefile_init (fout,glocal.name);
				for (auto &file:glocal.savefiles){
					const char *fname = basename (file);
					fprintf (fout,"test ! -f %s%s || mv -f %s%s data/%s\n"
						,glocal.rootfs.c_str(),file.c_str()
						,glocal.rootfs.c_str(),file.c_str()
						,fname);
				}
				return 0;
			</f>
			</call>
		}
		if (glocal.restorefile != NULL){
			<call savefile>(glocal.restorefile,false);
			<f dowrite>
				savefile_init (fout,glocal.name);
				for (auto &file:glocal.savefiles){
					const char *fname = basename (file);
					fprintf (fout,"test ! -f data/%s || cp -af data/%s %s%s\n"
						,fname,fname,glocal.rootfs.c_str(),file.c_str());
				}
				return 0;
			</f>
			</call>
		}
		if (ffiles != NULL) fclose (ffiles);
		return ret;
	</f>
	</call>
	return glocal.ret;
}
</mod>