From c18bd53ed5b423c994a261736453160284af860e Mon Sep 17 00:00:00 2001
From: hengyunabc <hengyunabc@gmail.com>
Date: Thu, 7 Feb 2019 12:43:46 +0800
Subject: [PATCH] dump command support output directory. #419

---
 .../klass100/ClassDumpTransformer.java        | 14 +++++++++-
 .../command/klass100/DumpClassCommand.java    | 28 +++++++++++++++++--
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/core/src/main/java/com/taobao/arthas/core/command/klass100/ClassDumpTransformer.java b/core/src/main/java/com/taobao/arthas/core/command/klass100/ClassDumpTransformer.java
index 5bf50a512..62919bf47 100644
--- a/core/src/main/java/com/taobao/arthas/core/command/klass100/ClassDumpTransformer.java
+++ b/core/src/main/java/com/taobao/arthas/core/command/klass100/ClassDumpTransformer.java
@@ -24,10 +24,17 @@ class ClassDumpTransformer implements ClassFileTransformer {
     private Map<Class<?>, File> dumpResult;
     private File arthasLogHome;
 
+    private File directory;
+
     public ClassDumpTransformer(Set<Class<?>> classesToEnhance) {
+        this(classesToEnhance, null);
+    }
+
+    public ClassDumpTransformer(Set<Class<?>> classesToEnhance, File directory) {
         this.classesToEnhance = classesToEnhance;
         this.dumpResult = new HashMap<Class<?>, File>();
         this.arthasLogHome = new File(LogUtil.LOGGER_FILE).getParentFile();
+        this.directory = directory;
     }
 
     @Override
@@ -50,7 +57,12 @@ class ClassDumpTransformer implements ClassFileTransformer {
         String classDumpDir = "classdump";
 
         // 创建类所在的包路径
-        File dumpDir = new File(arthasLogHome, classDumpDir);
+        File dumpDir = null;
+        if (directory != null) {
+            dumpDir = directory;
+        } else {
+            dumpDir = new File(arthasLogHome, classDumpDir);
+        }
         if (!dumpDir.mkdirs() && !dumpDir.exists()) {
             logger.warn("create dump directory:{} failed.", dumpDir.getAbsolutePath());
             return;
diff --git a/core/src/main/java/com/taobao/arthas/core/command/klass100/DumpClassCommand.java b/core/src/main/java/com/taobao/arthas/core/command/klass100/DumpClassCommand.java
index 1d25408d0..b56954c73 100644
--- a/core/src/main/java/com/taobao/arthas/core/command/klass100/DumpClassCommand.java
+++ b/core/src/main/java/com/taobao/arthas/core/command/klass100/DumpClassCommand.java
@@ -37,6 +37,7 @@ import static com.taobao.text.ui.Element.label;
 @Summary("Dump class byte array from JVM")
 @Description(Constants.EXAMPLE +
         "  dump java.lang.String\n" +
+        "  dump -d /tmp/output java.lang.String\n" +
         "  dump org/apache/commons/lang/StringUtils\n" +
         "  dump *StringUtils\n" +
         "  dump -E org\\\\.apache\\\\.commons\\\\.lang\\\\.StringUtils\n" +
@@ -48,6 +49,8 @@ public class DumpClassCommand extends AnnotatedCommand {
     private String code = null;
     private boolean isRegEx = false;
 
+    private String directory;
+
     @Argument(index = 0, argName = "class-pattern")
     @Description("Class name pattern, use either '.' or '/' as separator")
     public void setClassPattern(String classPattern) {
@@ -66,13 +69,27 @@ public class DumpClassCommand extends AnnotatedCommand {
         isRegEx = regEx;
     }
 
+    @Option(shortName = "d", longName = "directory")
+    @Description("Sets the destination directory for class files")
+    public void setDirectory(String directory) {
+        this.directory = directory;
+    }
+
     @Override
     public void process(CommandProcess process) {
         RowAffect effect = new RowAffect();
-        Instrumentation inst = process.session().getInstrumentation();
 
-        Set<Class<?>> matchedClasses = SearchUtils.searchClass(inst, classPattern, isRegEx, code);
         try {
+            if (directory != null) {
+                File dir = new File(directory);
+                if (dir.isFile()) {
+                    process.write(directory + " :is not a directory, please check it\n");
+                    return;
+                }
+            }
+            Instrumentation inst = process.session().getInstrumentation();
+            Set<Class<?>> matchedClasses = SearchUtils.searchClass(inst, classPattern, isRegEx, code);
+
             if (matchedClasses == null || matchedClasses.isEmpty()) {
                 processNoMatch(process);
             } else if (matchedClasses.size() > 5) {
@@ -134,7 +151,12 @@ public class DumpClassCommand extends AnnotatedCommand {
     }
 
     private Map<Class<?>, File> dump(Instrumentation inst, Set<Class<?>> classes) throws UnmodifiableClassException {
-        ClassDumpTransformer transformer = new ClassDumpTransformer(classes);
+        ClassDumpTransformer transformer = null;
+        if (directory != null) {
+            transformer = new ClassDumpTransformer(classes, new File(directory));
+        } else {
+            transformer = new ClassDumpTransformer(classes);
+        }
         Enhancer.enhance(inst, transformer, classes);
         return transformer.getDumpResult();
     }