package ca.uwaterloo.cs858.analysis;

import com.ibm.wala.classLoader.*;
import com.ibm.wala.classLoader.Module;
import com.ibm.wala.dalvik.classLoader.DexFileModule;
import com.ibm.wala.dalvik.classLoader.DexIRFactory;
import com.ibm.wala.dalvik.classLoader.WDexClassLoaderImpl;
import com.ibm.wala.dalvik.ipa.callgraph.impl.AndroidEntryPoint;
import com.ibm.wala.dalvik.ipa.callgraph.impl.DexEntryPoint;
import com.ibm.wala.dalvik.util.AndroidAnalysisScope;
import com.ibm.wala.dalvik.util.AndroidEntryPointLocator;
import com.ibm.wala.ipa.callgraph.*;
import com.ibm.wala.ipa.callgraph.impl.*;
import com.ibm.wala.ipa.callgraph.propagation.SSAPropagationCallGraphBuilder;
import com.ibm.wala.ipa.cha.ClassHierarchy;
import com.ibm.wala.ipa.slicer.*;
import com.ibm.wala.ssa.*;
import com.ibm.wala.types.*;
import com.ibm.wala.util.Predicate;
import com.ibm.wala.util.WalaException;
import com.ibm.wala.util.config.AnalysisScopeReader;
import com.ibm.wala.util.graph.Graph;
import com.ibm.wala.util.graph.GraphSlicer;
import com.ibm.wala.util.graph.traverse.SlowDFSDiscoverTimeIterator;
import com.ibm.wala.viz.DotUtil;
import com.ibm.wala.viz.NodeDecorator;
import edu.purdue.cs.toydroid.utils.*;
import org.apache.commons.cli.*;
import org.apache.commons.io.FileUtils;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.io.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.stream.Collectors;


public class Main implements Callable<Main> {

    static class AnalysisScopeUtil {
        private static final Logger logger = LogManager.getLogger(AnalysisScopeUtil.class);
        private static final ClassLoader WALA_CLASSLOADER = AnalysisScopeReader.class.getClassLoader();
        private static final String BASIC_FILE = "primordial.txt";
        private static final String EXCLUSIONS = "AndroidRegressionExclusions.txt";
        private static final String PrimordialTag = "Primordial";
        private static final String ApplicationTag = "Application";

        public static AnalysisScope makeAnalysisScope() throws IOException, FileNotFoundException {
            AnalysisScope scope = null;
            String android_jar = SimpleConfig.getAndroidJar();
            String exclusion_file = SimpleConfig.getExclusionFile();
            logger.info("AndroidJar: {}", new Object[]{android_jar});
            if (exclusion_file == null) {
                exclusion_file = "AndroidRegressionExclusions.txt";
            }

            logger.info("ExclusionFile: {}", new Object[]{exclusion_file});
            scope = AnalysisScopeReader.readJavaScope("primordial.txt", new File(exclusion_file), WALA_CLASSLOADER);
            AndroidAnalysisScope.addClassPathToScope(android_jar, scope, ClassLoaderReference.Primordial);
            scope.setLoaderImpl(ClassLoaderReference.Application, "com.ibm.wala.dalvik.classLoader.WDexClassLoaderImpl");
            Iterator iter = SimpleConfig.iteratorAdditionalJars();

            while (iter.hasNext()) {
                String additional_jar = (String) iter.next();
                int idx = additional_jar.indexOf(44);
                if (idx > 0) {
                    String ref = additional_jar.substring(0, idx);
                    String path = additional_jar.substring(idx + 1);
                    ClassLoaderReference clRef = null;
                    logger.info("AdditionalJar[{}]: {}", new Object[]{ref, path});
                    if (ref.startsWith("Primordial")) {
                        clRef = scope.getPrimordialLoader();
                    } else if (ref.startsWith("Application")) {
                        clRef = scope.getApplicationLoader();
                    } else {
                        logger.warn("Unrecognized ADDITIONAL_JARS in Config.properties.");
                    }

                    if (clRef != null) {
                        AndroidAnalysisScope.addClassPathToScope(path, scope, clRef);
                    }
                }
            }
            return scope;
        }

        public static AnalysisScope makeAnalysisScope(String apkFile) throws IOException, FileNotFoundException {
            AnalysisScope scope = null;
            String android_jar = SimpleConfig.getAndroidJar();
            String exclusion_file = SimpleConfig.getExclusionFile();
            logger.info("AndroidJar: {}", new Object[]{android_jar});
            if (exclusion_file == null) {
                exclusion_file = "AndroidRegressionExclusions.txt";
            }

            logger.info("ExclusionFile: {}", new Object[]{exclusion_file});
            scope = AnalysisScopeReader.readJavaScope("primordial.txt", new File(exclusion_file), WALA_CLASSLOADER);
            AndroidAnalysisScope.addClassPathToScope(android_jar, scope, ClassLoaderReference.Primordial);

            scope.setLoaderImpl(ClassLoaderReference.Application, "com.ibm.wala.dalvik.classLoader.WDexClassLoaderImpl");
            scope.addToScope(ClassLoaderReference.Application, DexFileModule.make(new File(apkFile)));

            return scope;
        }
    }

    private static Logger logger = LogManager.getLogger(Main.class);
    private static boolean taskTimeout = false;
    private static long timeout = 20;
    private static String outputDir = "C:\\Users\\mrafu\\IdeaProjects\\Analysis\\results";
    private static long startTask;
    private static List<String> ApkFiles = Arrays.asList(
            "C:\\Users\\mrafu\\IdeaProjects\\Analysis\\apks\\app-host-shamir-debug\\classes.dex",
            "C:\\Users\\mrafu\\IdeaProjects\\Analysis\\apks\\app-host-shamir-debug\\classes2.dex"
    );

    public static void main(String[] args) {
        try {
            logger.info("APK Files: {}", ApkFiles);
            List<File> files = ApkFiles.stream().map((file) -> new File(file)).collect(Collectors.toList());
            for (File file : files) {
                if (!(file.isFile() && file.canRead())) {
                    throw new Exception("Invalid APK file: " + file.getAbsolutePath());
                }
            }
            logger.info("Set TIMEOUT as {} minutes.", timeout);

            String fName = files.get(0).getParentFile().getName();
            int idxDot = fName.lastIndexOf('.');
            if (idxDot >= 0) {
                fName = fName.substring(0, idxDot);
            }

            outputDir = outputDir + File.separator + fName;
            logger.info("Writing to: " + outputDir);

            File file = new File(outputDir);
            if (!file.exists()) {
                file.mkdirs();// cascade mkdir
            } else if (!file.isDirectory()) {
                throw new Exception("Non-directory path exists as OUTPUT dir: " + outputDir);
            } else {
                FileUtils.cleanDirectory(file);
            }
        } catch (Exception pex) {
            logger.error("Error parsing command line arguments: {}", pex.getMessage());
        }

        startTask = System.currentTimeMillis();
        try {
            doAnalysis();
        } catch (Throwable e) {
            logger.error("Crashed: {}", e.getMessage());
        }

        long endTask = System.currentTimeMillis();
        String time = String.format("%d.%03d", (endTask - startTask) / 1000, (endTask - startTask) % 1000);
        logger.info("Total Time: {} seconds.", time);
        System.exit(0);
    }

    public static void doAnalysis() throws Throwable {
        logger.info("Start Analysis...");
        Main analysis = new Main();
        Future<Main> future = Executors.newSingleThreadExecutor().submit(analysis);
        try {
            future.get(timeout, TimeUnit.MINUTES);
        } catch (InterruptedException ire) {
            logger.error("InterruptedException: {}", ire.getMessage());
        } catch (ExecutionException exe) {
            logger.error("ExecutionException: {}", exe.getMessage());
            exe.printStackTrace();
        } catch (TimeoutException toe) {
            taskTimeout = true;
            logger.warn("Analysis Timeout After {} {}!", timeout, TimeUnit.MINUTES);
        }
        long analysisEnd = System.currentTimeMillis();
        String time = String.format("%d.%03d",
                (analysisEnd - startTask) / 1000,
                (analysisEnd - startTask) % 1000);
        logger.info("Total Analysis Time: {} seconds.", time);
    }

    ////////////////////////////////////
    private AnalysisScope scope;
    private ClassHierarchy cha;
    private IClass asyncTask;

    private enum InvokeType {
        NONE,
        CANCAL,
        ISCANCELLED
    }

    public Main() {

    }

    @Override
    public Main call() throws Exception {
        initialize();
        analyze();
        return this;
    }

    private void initialize() throws Exception {
        scope = AnalysisScopeUtil.makeAnalysisScope();
        cha = ClassHierarchy.make(scope);
        // make class hierarchy for the target module only
        AnalysisScope appScope = AnalysisScopeUtil.makeAnalysisScope(ApkFiles.get(0));
        ClassHierarchy appCha = ClassHierarchy.make(appScope);
        IClassLoader appLoader = appCha.getFactory().getLoader(ClassLoaderReference.Application, appCha, appScope);
        // remove classes (in target module) from existing class hierarchy's class loader
        Set<IClass> toRemove = new HashSet<>();
        Iterator<IClass> appIter = appLoader.iterateAllClasses();
        while (appIter.hasNext()) {
            IClass k = appIter.next();
            toRemove.add(k);
        }

        appScope = null;
        appCha = null;
        IClassLoader oriLoader = cha.getFactory().getLoader(ClassLoaderReference.Application, cha, scope);
        oriLoader.removeAll(toRemove);
        toRemove = null;

        // add target module to scope
        scope.setLoaderImpl(ClassLoaderReference.Application, "com.ibm.wala.dalvik.classLoader.WDexClassLoaderImpl");
        ArrayList<Module> list = new ArrayList();
        for (String apkFile : ApkFiles) {
            Module dexModule = DexFileModule.make(new File(apkFile));
            scope.addToScope(ClassLoaderReference.Application, dexModule);
            list.add(dexModule);

        }

        // init the module s.t. the classes are included in the scope's class loader
        oriLoader.init(list);

        // make new class hierarchy
        cha = ClassHierarchy.make(scope, cha.getFactory());
    }

    private void analyze() throws Exception {
        Entrypoint entrypoint;
        SSAPropagationCallGraphBuilder cgBuilder;
        CallGraph cg;
        AnalysisCache cache = new AnalysisCache(new DexIRFactory());
        ArrayList<Entrypoint> entryList = new ArrayList<Entrypoint>();


//        for (IClass klass : cha) {
//            if (klass.getName().toString().startsWith("Lca")) {
//                System.out.println(klass.getName());
//            }
//        }

//        try {
//            IClass klass = cha.lookupClass(TypeReference.find(ClassLoaderReference.Application, "Lca/uwaterloo/cs858shamirhost/MainActivity"));
//            if (klass != null) {
//                System.out.println(klass.getName());
//            }
//        } catch (Exception e) {
//            e.printStackTrace();
//        }

        // build call graph for SmsManager.sendTextMessage(String, String, String, PendingIntent, PendingIntent)
        try {
            IClass klass = cha.lookupClass(TypeReference.find(ClassLoaderReference.Application, "Lca/uwaterloo/cs858shamirhost/MainActivity"));
            if (klass != null) {
                System.err.println(klass);
                Collection<IMethod> coll = klass.getDeclaredMethods();
                for (IMethod m : coll) {
                    if (m.getName().toString().equals("sendData")) {
                        DexEntryPoint dep = new DexEntryPoint(m, cha);
                        entryList.add(dep);
                        AnalysisOptions options = new AnalysisOptions(scope, entryList);
                        // context insensitive call graph
                        CallGraphBuilder builder = Util.makeVanillaZeroOneCFABuilder(options, cache, cha, scope);
                        cg = builder.makeCallGraph(options, null);
                        DotUtil.dotify(cg, new NodeDecorator<CGNode>() {
                            @Override
                            public String getLabel(CGNode cgNode) throws WalaException {
                                StringBuilder builder = new StringBuilder();
                                builder.append(cgNode.getMethod().getDeclaringClass().getName().toString());
                                builder.append("\n");
                                builder.append(cgNode.getMethod().getName().toString());
                                builder.append("\n");
                                builder.append(cgNode.getMethod().getSelector().toString());
                                return builder.toString();
                            }
                        }, outputDir + File.separator + "sendData.dot", null, null);
                        break;
                    }
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }

        // build call graph for SmsStorageMonitor$2.onReceive()
        // the call graph might be in-complete. For example, some methods are invoked but the corresponding method body is
        // not in the call graph. Then we can iterate each cg node in the CG and analyze the control flow graph to get more
        // information.
//        try {
//            IClass receiver = cha.lookupClass(TypeReference.find(ClassLoaderReference.Application, "Lcom/android/internal/telephony/SmsStorageMonitor$2"));
//            if (receiver != null) {
//                IMethod onReceive = receiver.getMethod(Selector.make("onReceive(Landroid/content/Context;Landroid/content/Intent;)V"));
//                if (onReceive != null) {
//                    DexEntryPoint dep = new DexEntryPoint(onReceive, cha);
//                    entryList.clear();
//                    entryList.add(dep);
//                    AnalysisOptions options = new AnalysisOptions(scope, entryList);
//                    // context insensitive call graph
//                    CallGraphBuilder builder = Util.makeVanillaZeroOneCFABuilder(options, cache, cha, scope);
//                    cg = builder.makeCallGraph(options, null);
//                    // dump the call graph to a graphviz dot file; use graphviz dot to convert it to other formats, e.g, pdf, png, svg.
//                    DotUtil.dotify(cg, new NodeDecorator<CGNode>() {
//                        @Override
//                        public String getLabel(CGNode cgNode) throws WalaException {
//                            StringBuilder builder = new StringBuilder();
//                            builder.append(cgNode.getMethod().getDeclaringClass().getName().toString());
//                            builder.append("\n");
//                            builder.append(cgNode.getMethod().getName().toString());
//                            builder.append("\n");
//                            builder.append(cgNode.getMethod().getSelector().toString());
//                            return builder.toString();
//                        }
//                    }, outputDir + File.separator + "onReceive.dot", null, null);
//                }
//            }
//        } catch (Exception e) {
//            e.printStackTrace();
//        }
    }

}
