package gadgetinspector;

import gadgetinspector.config.ConfigRepository;
import gadgetinspector.config.GIConfig;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.Writer;
import java.nio.charset.StandardCharsets;
import java.nio.file.FileVisitResult;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.file.SimpleFileVisitor;
import java.nio.file.StandardOpenOption;
import java.nio.file.attribute.BasicFileAttributes;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Scanner;
import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
import org.apache.log4j.ConsoleAppender;
import org.apache.log4j.Level;
import org.apache.log4j.PatternLayout;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Main entry point for running an end-to-end analysis. Deletes all data files before starting and writes discovered
 * gadget chains to gadget-chains.txt.
 */
public class GadgetInspector {
    private static final Logger LOGGER = LoggerFactory.getLogger(GadgetInspector.class);

    private static void printUsage() {
        System.out.println("Usage:\n  Pass either a single argument which will be interpreted as a WAR, or pass " +
                "any number of arguments which will be intepretted as a list of JARs forming a classpath.");

    }

    static {
        try {
            if (!Files.exists(Paths.get("gadget-result"))) {
                Files.createDirectory(Paths.get("gadget-result"));
            }
            if (!Files.exists(Paths.get("conf"))) {
                Files.createDirectory(Paths.get("conf"));
            }
            if (!Files.exists(Paths.get("history"))) {
                Files.createDirectory(Paths.get("history"));
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static void main(String[] args) throws Exception {
        if (args.length == 0) {
            printUsage();
            System.exit(1);
        }

        //配置log4j用于输出日志
        configureLogging();

        //是否不删除所有的dat文件
        boolean resume = false;
        //是否Spring-Boot jar项目
        boolean boot = false;
        //fuzz类型，默认java原生序列化
        GIConfig config = ConfigRepository.getConfig("jserial");

        int argIndex = 0;
        while (argIndex < args.length) {
            String arg = args[argIndex];
            if (!arg.startsWith("--")) {
                break;
            }
            if (arg.equals("--resume")) {
                //不删除dat文件
                resume = true;
            } else if (arg.equals("--config")) {
                //--config参数指定fuzz类型
                config = ConfigRepository.getConfig(args[++argIndex]);
                if (config == null) {
                    throw new IllegalArgumentException("Invalid config name: " + args[argIndex]);
                }
                ConfigHelper.giConfig = config;
            } else if (arg.equals("--boot")) {
                //指定为Spring-Boot jar项目
                boot = true;
            } else if (arg.equals("--mybatis.xml")) {
                //mybatis mapper xml目录位置
                ConfigHelper.mybatisMapperXMLPath = args[++argIndex];
            } else if (arg.equals("--noTaintTrack")) {
                //是否污点分析，若不使用污点分析，将会把所有链都搜索出来，好处是不会遗漏，坏处是需要大量的人工审计
                ConfigHelper.taintTrack = false;
            } else if (arg.equals("--opLevel")) {
                //链聚合优化等级，--OpLevel 1表示一层优化，默认0不优化
                ConfigHelper.opLevel = Integer.parseInt(args[++argIndex]);
            } else if (arg.equals("--history")) {
                //启用历史扫描jar包记录，方便大规模扫描时不重复扫描旧jar包，好处时减少工作时间，坏处是遇到依赖组合的gadget可能扫不出来
                ConfigHelper.history = true;
                ConfigHelper.historyRecordFile += args[++argIndex];
            } else if (arg.equals("--max")) {
                //大规模扫描jar时，指定最大数量
                ConfigHelper.maxJarCount = Integer.parseInt(args[++argIndex]);
            } else if (arg.equals("--onlyJDK")) {
                //仅扫描jdk
                ConfigHelper.onlyJDK = true;
            } else if (arg.equals("--maxChainLength")) {
                //仅输出小于该链长度的链
                ConfigHelper.maxChainLength = Integer.parseInt(args[++argIndex]);
            } else if (arg.equals("--craw")) {
                //启用爬虫功能，每*分钟分析一遍，出一次报告。若配置--crawMaven，则使用内置爬虫功能，爬取maven仓库
                ConfigHelper.craw = true;
                ConfigHelper.crawMin = Integer.parseInt(args[++argIndex]);
            } else if (arg.equals("--crawMaven")) {
                //爬取maven仓库
                ConfigHelper.crawMaven = true;
                ConfigHelper.crawMavenJarPath = args[++argIndex];
                if (!ConfigHelper.crawMavenJarPath.endsWith("/"))
                    ConfigHelper.crawMavenJarPath += "/";
            } else if (arg.equals("--onlyCrawMaven")) {
                //只爬取maven仓库，不挖掘gadget
                ConfigHelper.onlyCrawMaven = true;
            } else if (arg.equals("--onlyCrawNexus")) {
                //只爬取nexus本地仓库，不挖掘gadget
                ConfigHelper.onlyCrawNexus = true;
            } else if (arg.equals("--onlyCrawMavenPopular")) {
                //只爬取maven仓库popular，不挖掘gadget
                ConfigHelper.onlyCrawMavenPopular = true;
            } else if (arg.equals("--skipSourcesFile")) {
                //跳过的sources
                ConfigHelper.skipSourcesFile = args[++argIndex];
            } else if (arg.equals("--slinksFile")) {
                //自定义slinks，每行空格分离三段（
                // 类 方法名 方法描述，例：java/rmi/registry/Registry lookup (Ljava/lang/String;)Ljava/rmi/Remote;），
                // 缺少的不判断（例：java/rmi/registry/Registry lookup，不判断方法描述）
                ConfigHelper.slinksFile = args[++argIndex];
            } else if (arg.equals("--slink")) {
              for (int i = argIndex + 1; i < args.length; i++) {
                if (!args[i].startsWith("--") && !args[i].startsWith("/")) {
                  ConfigHelper.slinks.add(args[++argIndex]);
                } else {
                  break;
                }
              }
            } else {
                throw new IllegalArgumentException("Unexpected argument: " + arg);
            }

            argIndex += 1;
        }

        if (ConfigHelper.crawMaven) {
            if (!ConfigHelper.onlyCrawMavenPopular && !ConfigHelper.onlyCrawNexus) {
                MavenCrawer.start();
            }
            if (!ConfigHelper.onlyCrawMaven && !ConfigHelper.onlyCrawNexus) {
                MavenCrawerPopular.start();
            }
            if (!ConfigHelper.onlyCrawMaven && !ConfigHelper.onlyCrawMavenPopular) {
//                MavenCrawerNexus.start();
            }
        }

        for (int i = 0; (i < 1 || ConfigHelper.craw) && (!ConfigHelper.onlyCrawMaven && !ConfigHelper.onlyCrawNexus && !ConfigHelper.onlyCrawMavenPopular); i++) {
            AtomicBoolean haveNewJar = new AtomicBoolean(false);
            List<Path> pathList = new ArrayList<>();
            ClassLoader classLoader = initJarData(args, boot, argIndex, haveNewJar, pathList);
            if (ConfigHelper.craw && !haveNewJar.getAndSet(false)) {
                Thread.sleep(60 * 1000L);
                continue;
            }

            //类枚举加载器，具有两个方法
            //getRuntimeClasses获取rt.jar的所有class
            //getAllClasses获取rt.jar以及classLoader加载的class
            final ClassResourceEnumerator classResourceEnumerator = new ClassResourceEnumerator(
                    classLoader);

            //删除所有的dat文件
            if (!resume) {
                // Delete all existing dat files
                LOGGER.info("Deleting stale data...");
                for (String datFile : Arrays
                    .asList("classes.dat", "methods.dat", "inheritanceMap.dat",
                        "passthrough.dat", "callgraph.dat", "sources.dat", "methodimpl.dat",
                        "slinks.dat")) {
                    final Path path = Paths.get(datFile);
                    if (Files.exists(path)) {
                        Files.delete(path);
                    }
                }
            }

            try {
              //扫描java runtime所有的class（rt.jar）和指定的jar或war中的所有class
              beginDiscovery(config, classResourceEnumerator, pathList);
            } catch (Throwable t) {
              t.printStackTrace();
              //捕获异常，避免异常导致程序终止
            }

            saveHistory(haveNewJar, pathList);
            Thread.sleep(ConfigHelper.crawMin * 60 * 1000L);
            LOGGER.info("Analysis complete!");
        }
    }

    /**
     * 保存已被扫描的jar包记录
     *
     * @param haveNewJar
     * @param pathList
     * @throws IOException
     */
    private static void saveHistory(AtomicBoolean haveNewJar, List<Path> pathList)
        throws IOException {
        if (ConfigHelper.history && pathList.size() > 0) {
            haveNewJar.set(true);
            Path filePath = Paths.get(ConfigHelper.historyRecordFile);
            if (!Files.exists(filePath))
                Files.createFile(filePath);
            try (OutputStream outputStream = Files
                .newOutputStream(filePath,
                    StandardOpenOption.APPEND);
                Writer writer = new OutputStreamWriter(outputStream,
                    StandardCharsets.UTF_8)) {
                for (String jar : pathList.stream().map(path -> path.getFileName().toString()).collect(
                    Collectors.toSet())) {
                    writer.write(jar);
                    writer.write("\n");
                }
                writer.flush();
            }
        }
    }

    private static ClassLoader initJarData(String[] args, boolean boot, int argIndex,
        AtomicBoolean haveNewJar, List<Path> pathList)
        throws IOException {
        ClassLoader classLoader = null;
        if (!ConfigHelper.onlyJDK) {
            //程序参数的最后一部分，即最后一个具有前缀--的参数（例：--resume）后
            if (args.length == argIndex + 1 && args[argIndex].toLowerCase().endsWith(".war")) {
                //加载war文件
                Path path = Paths.get(args[argIndex]);
                LOGGER.info("Using WAR classpath: " + path);
                //实现为URLClassLoader，加载war包下的WEB-INF/lib和WEB-INF/classes
                classLoader = Util.getWarClassLoader(path);
            } else if (args.length == argIndex + 1 && args[argIndex].toLowerCase().endsWith(".jar")
                && boot) {
                Path path = Paths.get(args[argIndex]);
                LOGGER.info("Using JAR classpath: " + path);
                //实现为URLClassLoader，加载jar包下的BOOT-INF/lib和BOOT-INF/classes
                classLoader = Util.getJarAndLibClassLoader(path);
            } else {
                //加载jar文件，java命令后部，可配置多个
                Set<String> scanJarHistory = new HashSet<>();
                Path filePath = Paths.get(ConfigHelper.historyRecordFile);
                if (ConfigHelper.history && Files.exists(filePath)) {
                    try (InputStream inputStream = Files
                        .newInputStream(filePath);
                        Scanner scanner = new Scanner(inputStream, StandardCharsets.UTF_8.name())) {
                        while (scanner.hasNext()) {
                            String jar = scanner.nextLine();
                            if (jar.length() > 0) {
                                scanJarHistory.add(jar.trim());
                            }
                        }
                    }
                }
                AtomicInteger jarCount = new AtomicInteger(0);
                for (int i = 0; i < args.length - argIndex; i++) {
                    String pathStr = args[argIndex + i];
                    if (!pathStr.endsWith(".jar")) {
                        //todo 主要用于大批量的挖掘链
                        //非.jar结尾，即目录，需要遍历目录找出所有jar文件
                        File file = Paths.get(pathStr).toFile();
                        if (file == null || !file.exists())
                            continue;
                        Files.walkFileTree(file.toPath(), new SimpleFileVisitor<Path>() {
                            @Override
                            public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) {
                                if (!file.getFileName().toString().endsWith(".jar"))
                                    return FileVisitResult.CONTINUE;
                                File readFile = file.toFile();
                                Path path = Paths.get(readFile.getAbsolutePath());
                                if (Files.exists(path)) {
                                    if (ConfigHelper.history) {
                                        if (!scanJarHistory.contains(path.getFileName().toString())) {
                                            if (jarCount.incrementAndGet() <= ConfigHelper.maxJarCount) {
                                                pathList.add(path);
                                            }
                                        }
                                    } else {
                                        if (jarCount.incrementAndGet() <= ConfigHelper.maxJarCount) {
                                            pathList.add(path);
                                        }
                                    }
                                }
                                return FileVisitResult.CONTINUE;
                            }
                        });

                        continue;
                    }
                    Path path = Paths.get(pathStr).toAbsolutePath();
                    if (!Files.exists(path)) {
                        throw new IllegalArgumentException("Invalid jar path: " + path);
                    }
                    pathList.add(path);
                }
                LOGGER.info("Using classpath: " + Arrays.toString(pathList.toArray()));
                //实现为URLClassLoader，加载所有指定的jar
                classLoader = Util.getJarClassLoader(pathList.toArray(new Path[0]));
            }
        }
        haveNewJar.set(pathList.size() != 0);
        return classLoader;
    }

    private static void beginDiscovery(GIConfig config,
        ClassResourceEnumerator classResourceEnumerator,
        List<Path> pathList) throws Exception {
        // Perform the various discovery steps
        if (!Files.exists(Paths.get("classes.dat")) || !Files.exists(Paths.get("methods.dat"))
                || !Files.exists(Paths.get("inheritanceMap.dat"))) {
            LOGGER.info("Running method discovery...");
            MethodDiscovery methodDiscovery = new MethodDiscovery();
            methodDiscovery.discover(classResourceEnumerator);
            //保存了类信息、方法信息、继承实现信息
            methodDiscovery.save();
        }

        if (!Files.exists(Paths.get("slinks.dat")) && config.getSlinkDiscovery() != null) {
            LOGGER.info("Running slink discovery...");
            SlinkDiscovery slinkDiscovery = config.getSlinkDiscovery();
            slinkDiscovery.discover();
            slinkDiscovery.save();
        }

        if (!Files.exists(Paths.get("passthrough.dat")) && ConfigHelper.taintTrack) {
            LOGGER.info("Analyzing methods for passthrough dataflow...");
            PassthroughDiscovery passthroughDiscovery = new PassthroughDiscovery();
            //记录参数在方法调用链中的流动关联（如：A、B、C、D四个方法，调用链为A->B B->C C->D，其中参数随着调用关系从A流向B，在B调用C过程中作为入参并随着方法结束返回，最后流向D）
            //该方法主要是追踪上面所说的"B调用C过程中作为入参并随着方法结束返回"，入参和返回值之间的关联
            passthroughDiscovery.discover(classResourceEnumerator, config);
            passthroughDiscovery.save();
        }

        if (!Files.exists(Paths.get("callgraph.dat"))) {
            LOGGER.info("Analyzing methods in order to build a call graph...");
            CallGraphDiscovery callGraphDiscovery = new CallGraphDiscovery();
            //记录参数在方法调用链中的流动关联（如：A、B、C三个方法，调用链为A->B B->C，其中参数随着调用关系从A流向B，最后流C）
            //该方法主要是追踪上面所说的参数流动，即A->B入参和B->C入参的关系，以确定参数可控
            callGraphDiscovery.discover(classResourceEnumerator, config);
            callGraphDiscovery.save();
        }

        if (!Files.exists(Paths.get("sources.dat"))) {
            LOGGER.info("Discovering gadget chain source methods...");
            SourceDiscovery sourceDiscovery = config.getSourceDiscovery();
            //查找利用链的入口（例：java原生反序列化的readObject）
            sourceDiscovery.discover();
            sourceDiscovery.save();
        }

        {
            LOGGER.info("Searching call graph for gadget chains...");
            GadgetChainDiscovery gadgetChainDiscovery = new GadgetChainDiscovery(config);
            //根据上面的数据收集，最终分析利用链
            gadgetChainDiscovery.discover(pathList);
        }
    }

    private static void configureLogging() {
        ConsoleAppender console = new ConsoleAppender();
        String PATTERN = "%d %c [%p] %m%n";
        console.setLayout(new PatternLayout(PATTERN));
        console.setThreshold(Level.INFO);
        console.activateOptions();
        org.apache.log4j.Logger.getRootLogger().addAppender(console);
    }
}
