import org.objectweb.asm.*;
import java.io.*;
import java.util.*;
import java.util.jar.*;

public class MethodCallAnalyzer
{
   private static final Set<String> appcdsClasses = new TreeSet<>();
   private static final Map<String, byte[]> jarClasses = new HashMap<>();
   private static final Queue<String> analysisQueue = new LinkedList<>();
   private static final Set<String> alreadyAnalyzed = new HashSet<>();
   
   private static final List<String> includeFilters = new ArrayList<>();
   private static final List<String> excludeFilters = new ArrayList<>();
   private static final List<String> explicitStops = new ArrayList<>();
   private static boolean implicitBridgeMode = false;

   private static final Map<String, String> dependencyGraph = new HashMap<>();
   private static String traceTarget = null;

   public static void main(String[] args)
   {
      if (args.length < 1)
      {
         System.err.println("Usage: java MethodCallAnalyzer <jar> [includes] [-excludes] [!stops] [--implicit-bridge] [--trace=class]");
         System.exit(1);
      }

      for (int i = 1; i < args.length; i++)
      {
         String arg = args[i];
         if (arg.startsWith("--trace="))
         {
            traceTarget = arg.substring(8).replace('.', '/');
         }
         else if (arg.equals("--implicit-bridge"))
         {
            implicitBridgeMode = true;
         }
         else if (arg.startsWith("-"))
         {
            excludeFilters.add(arg.substring(1).replace('.', '/'));
         }
         else if (arg.startsWith("!"))
         {
            explicitStops.add(arg.substring(1).replace('.', '/'));
         }
         else
         {
            includeFilters.add(arg.replace('.', '/'));
         }
      }

      try (JarFile jarFile = new JarFile(args[0]))
      {
         System.out.println("Loading JAR and SPIs...");
         loadJarAndSPIs(jarFile);
         seedQueue();

         System.out.println("Scanning...");
         while (!analysisQueue.isEmpty())
         {
            String current = analysisQueue.poll();
            if (!alreadyAnalyzed.add(current))
            {
               continue;
            }

            byte[] bytes = jarClasses.get(current);
            if (bytes == null)
            {
               continue;
            }

            // Shallow scan if it's a bridge or outside include list (if mode enabled)
            boolean isBridge = isExplicitStop(current) || (implicitBridgeMode && !isIncluded(current));

            analyzeClassBytes(current, bytes, isBridge);
         }

         try (PrintWriter w = new PrintWriter(new FileWriter("appcds_classes.lst")))
         {
            appcdsClasses.forEach(w::println);
         }

         System.out.println("Done! Total classes: " + appcdsClasses.size());
         if (traceTarget != null)
         {
            printTrace();
         }

      }
      catch (Exception e)
      {
         e.printStackTrace();
      }
   }

   private static void analyzeClassBytes(final String sourceClass, byte[] bytes, boolean shallowScan)
   {
      new ClassReader(bytes).accept(new ClassVisitor(Opcodes.ASM9)
      {
         private void reg(String name)
         {
            if (name == null || name.startsWith("[") || isExcluded(name))
            {
               return;
            }
            if (appcdsClasses.add(name))
            {
               analysisQueue.add(name);
               if (!dependencyGraph.containsKey(name))
               {
                  dependencyGraph.put(name, sourceClass);
               }
            }
         }

         private void regType(Type t)
         {
            if (t.getSort() == Type.OBJECT)
            {
               reg(t.getInternalName());
            }
            else if (t.getSort() == Type.ARRAY)
            {
               regType(t.getElementType());
            }
            else if (t.getSort() == Type.METHOD)
            {
               regType(t.getReturnType());
               for (Type a : t.getArgumentTypes())
               {
                  regType(a);
               }
            }
         }

         private void regDesc(String d)
         {
            if (d != null)
            {
               try
               {
                  regType(Type.getType(d));
               }
               catch (Exception e)
               {
               }
            }
         }

         @Override
         public void visit(int v, int a, String n, String s, String sn, String[] i)
         {
            reg(n);
            reg(sn);
            if (i != null)
            {
               for (String iface : i)
               {
                  reg(iface);
               }
            }
         }

         @Override
         public FieldVisitor visitField(int a, String n, String d, String s, Object v)
         {
            if (!shallowScan)
            {
               regDesc(d);
            }
            return null;
         }

         @Override
         public MethodVisitor visitMethod(int a, String n, String d, String s, String[] e)
         {
            if (!shallowScan)
            {
               regDesc(d);
               if (e != null)
               {
                  for (String ex : e)
                  {
                     reg(ex);
                  }
               }
            }
            if (shallowScan)
            {
               return null; // STOP HERE for bridge classes
            }

            return new MethodVisitor(Opcodes.ASM9)
            {
               @Override
               public void visitMethodInsn(int o, String ow, String n, String d, boolean i)
               {
                  reg(ow);
                  regDesc(d);
                  checkSmartTriggers(ow.replace('/', '.'));
               }

               @Override
               public void visitFieldInsn(int o, String ow, String n, String d)
               {
                  reg(ow);
                  regDesc(d);
               }

               @Override
               public void visitTypeInsn(int o, String t)
               {
                  if (t.startsWith("["))
                  {
                     regDesc(t);
                  }
                  else
                  {
                     reg(t);
                  }
               }

               @Override
               public void visitLdcInsn(Object v)
               {
                  if (v instanceof Type)
                  {
                     regType((Type) v);
                  }
               }

               @Override
               public void visitTryCatchBlock(Label s, Label e, Label h, String t)
               {
                  if (t != null)
                  {
                     reg(t);
                  }
               }

               @Override
               public void visitInvokeDynamicInsn(String name, String descriptor, Handle bsm, Object... bsmArgs)
               {
                  regDesc(descriptor);
                  reg(bsm.getOwner());
                  for (Object arg : bsmArgs)
                  {
                     if (arg instanceof Type)
                     {
                        regType((Type) arg);
                     }
                     else if (arg instanceof Handle)
                     {
                        Handle h = (Handle) arg;
                        reg(h.getOwner());
                        regDesc(h.getDesc());
                     }
                  }
               }
            };
         }

         private void checkSmartTriggers(String cleanOwner)
         {
            // XML / Xerces
            if (cleanOwner.startsWith("javax.xml") || cleanOwner.startsWith("org.w3c"))
            {
               reg("com/sun/org/apache/xerces/internal/jaxp/DocumentBuilderFactoryImpl");
               reg("com/sun/org/apache/xerces/internal/parsers/DOMParser");
            }
            // Security / SSL
            else if (cleanOwner.startsWith("java.security") || cleanOwner.startsWith("javax.crypto") || cleanOwner.startsWith("javax.net.ssl"))
            {
               reg("sun/security/provider/Sun");
               reg("com/sun/crypto/provider/SunJCE");
               reg("sun/security/ssl/SSLContextImpl");
            }
            // NIO
            else if (cleanOwner.startsWith("java.nio.file"))
            {
               reg("sun/nio/fs/UnixFileSystemProvider");
               reg("sun/nio/ch/FileChannelImpl");
            }
         }
      }, 0);
   }

   private static void loadJarAndSPIs(JarFile jarFile) throws IOException
   {
      Enumeration<JarEntry> entries = jarFile.entries();
      while (entries.hasMoreElements())
      {
         JarEntry entry = entries.nextElement();
         String name = entry.getName();
         if (name.startsWith("META-INF/services/"))
         {
            try (BufferedReader r = new BufferedReader(new InputStreamReader(jarFile.getInputStream(entry))))
            {
               String line;
               while ((line = r.readLine()) != null)
               {
                  line = line.trim();
                  if (!line.isEmpty() && !line.startsWith("#"))
                  {
                     String spi = line.replace('.', '/');
                     if (!isExcluded(spi) && appcdsClasses.add(spi))
                     {
                        analysisQueue.add(spi);
                        dependencyGraph.put(spi, "[SPI]");
                     }
                  }
               }
            }
         }
         else if (name.endsWith(".class"))
         {
            String cls = name.substring(0, name.length() - 6);
            InputStream is = jarFile.getInputStream(entry);
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            byte[] buf = new byte[8192];
            int n;
            while ((n = is.read(buf)) != -1)
            {
               bos.write(buf, 0, n);
            }
            jarClasses.put(cls, bos.toByteArray());
         }
      }
   }

   private static void seedQueue()
   {
      for (String cls : jarClasses.keySet())
      {
         if (isIncluded(cls) && !isExcluded(cls))
         {
            if (appcdsClasses.add(cls))
            {
               analysisQueue.add(cls);
               dependencyGraph.put(cls, "[Root]");
            }
         }
      }
   }

   private static boolean isIncluded(String cls)
   {
      for (String f : includeFilters)
      {
         if (cls.startsWith(f + "/") || cls.equals(f) || cls.startsWith(f + "$"))
         {
            return true;
         }
      }
      return false;
   }

   private static boolean isExcluded(String cls)
   {
      for (String ex : excludeFilters)
      {
         if (cls.startsWith(ex + "/") || cls.equals(ex) || cls.startsWith(ex + "$"))
         {
            return true;
         }
      }
      return false;
   }

   private static boolean isExplicitStop(String cls)
   {
      for (String st : explicitStops)
      {
         if (cls.startsWith(st + "/") || cls.equals(st) || cls.startsWith(st + "$"))
         {
            return true;
         }
      }
      return false;
   }

   private static void printTrace()
   {
      System.out.println("\nTRACE FOR: " + traceTarget);
      String c = traceTarget;
      while (c != null)
      {
         System.out.println(" ↳ " + c);
         c = dependencyGraph.get(c);
         if (c != null && c.startsWith("["))
         {
            System.out.println(" ↳ " + c);
            break;
         }
      }
   }
}