做一个maven插件练练手

最近做的一个maven小插件

import com.google.common.base.Strings;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import com.google.common.reflect.ClassPath;
import org.apache.commons.io.FileUtils;
import org.apache.maven.plugin.AbstractMojo;
import org.apache.maven.plugin.MojoExecutionException;
import org.apache.maven.plugins.annotations.LifecyclePhase;
import org.apache.maven.plugins.annotations.Mojo;
import org.apache.maven.plugins.annotations.Parameter;
import org.apache.maven.plugins.annotations.ResolutionScope;
import org.apache.maven.project.MavenProject;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Type;
import org.objectweb.asm.tree.AnnotationNode;
import org.objectweb.asm.tree.ClassNode;
import org.objectweb.asm.tree.FieldNode;
import org.objectweb.asm.tree.MethodNode;

import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

/** * 将BeanType注解替换成Named注解 * * @author macrochen * @since 16/2/17 */
 * 将BeanType注解替换成Named注解
 *
 * @author macrochen
 * @since 16/2/17
 */
@Mojo(name = "beanType", defaultPhase = LifecyclePhase.PROCESS_CLASSES,
		requiresDependencyResolution = ResolutionScope.COMPILE)
public class BeanTypeMojo extends AbstractMojo {

	@Parameter(defaultValue = "${project.build.outputDirectory}", required = true)
	private File outputDirectory;

	/**	 * 扫描指定的包, 避免将整个classpath下的class都扫一遍, 如果有多个包, 用分号(;)分隔, 而且会扫描指定包的下级子包	 */
	 * 扫描指定的包, 避免将整个classpath下的class都扫一遍, 如果有多个包, 用分号(;)分隔, 而且会扫描指定包的下级子包
	 */
	@Parameter(defaultValue = "${beanType.scanPackage}", required = true)
	private String scanPackage;

	@Parameter(defaultValue = "${project}", readonly = true)
	private MavenProject project;

	public void execute()			throws MojoExecutionException {
			throws MojoExecutionException 

		try {
			ClassLoader classloader = getClassLoader(project);//Thread.currentThread().getContextClassLoader()
			List<ClassPath.ClassInfo> list = getClasses(classloader);
			for (ClassPath.ClassInfo classInfo : list) {
				String s = classInfo.getName();

				getLog().info("read class file: " + s);

				modify(classloader, s);
			}
		} catch (Exception e) {
			getLog().error("replace failure", e);
			throw new MojoExecutionException(e.getMessage(), e);
		}
	}

	private void modify(ClassLoader classloader, String s) throws Exception {
		// 更新内部类, 如果有的话
		modifyInnerClass(classloader, s);

		ClassNode cn = createClassNode(classloader, s);

		// 没修改不更新字节码
		if (!doModify(cn)) return;

		// 输出到class文件
		output(s, cn);
	}

	private void output(String s, ClassNode cn) throws IOException {
		s = outputDirectory.getAbsolutePath() + "/" + s.replace(".", "/") + ".class";

		FileUtils.writeByteArrayToFile(new File(s), getByteArray(cn));
		getLog().info("replace " + s + " is over");
	}

	private ClassNode createClassNode(ClassLoader classloader, String s) throws IOException {
		ClassReader cr = new ClassReader(classloader.getResourceAsStream(s.replace('.', '/') + ".class"));
		ClassNode cn = new ClassNode();
		cr.accept(cn, 0);
		return cn;
	}

	private void modifyInnerClass(ClassLoader classloader, String s) throws Exception {
		Class<?> clazz = Class.forName(s, true, classloader);
		Class<?>[] innserClasses = clazz.getDeclaredClasses();
		if (innserClasses != null && innserClasses.length > 0) {
			for (Class<?> innserClass : innserClasses) {
				modify(classloader, innserClass.getName());
			}
		}
	}

	private byte[] getByteArray(ClassNode cn) {
		ClassWriter cw = new ClassWriter(0);
		cn.accept(cw);
		return cw.toByteArray();
	}

	private List<ClassPath.ClassInfo> getClasses(ClassLoader classloader) throws IOException {
		List<ClassPath.ClassInfo> list = Lists.newArrayList();
		if (Strings.isNullOrEmpty(scanPackage)) throw new RuntimeException("scanPackage not be assigned.");
		String[] packages = scanPackage.split(";");
		for (String aPackage : packages) {
			ImmutableSet<ClassPath.ClassInfo> classes = ClassPath.from(classloader).getTopLevelClassesRecursive(aPackage);
			list.addAll(classes);
		}
		return list;
	}

	/**	 * 获取目标类库中的类加载器。目标类库就是指运行此插件的工程。	 */
	 * 获取目标类库中的类加载器。目标类库就是指运行此插件的工程。
	 */
	public static ClassLoader getClassLoader(MavenProject project) throws MojoExecutionException {
		try {
			List<String> classpathElements = project.getCompileClasspathElements();
			classpathElements.add(project.getBuild().getOutputDirectory());
			URL[] urls = new URL[classpathElements.size()];

			for (int i = 0; i < classpathElements.size(); ++i) {
				urls[i] = new File(classpathElements.get(i)).toURI().toURL();
			}

			ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader();
			return new URLClassLoader(urls, contextClassLoader);
		} catch (Exception e) {
			throw new MojoExecutionException("无法创建目标工程的类加载器", e);
		}
	}

	public boolean doModify(ClassNode cn) {
		boolean modified1 = forClass(cn);
		boolean modified2 = forMethods(cn);
		boolean modified3 = forFields(cn);
		return modified1 || modified2 || modified3;
	}

	private boolean forClass(ClassNode cn) {
		boolean modified = false;
		List<AnnotationNode> va = cn.visibleAnnotations;
		if (va != null && !va.isEmpty()) {
			Iterator<AnnotationNode> iterator = va.iterator();
			while (iterator.hasNext()) {
				AnnotationNode node = iterator.next();
				if (isBeanType(node)) {
					//getLog().info("find class: " + cn.name + " with beanType");
					modified = true;
					iterator.remove();
					break;
				}
			}
			if (modified) {
				String className = cn.name.replace("/", ".");
				addNamedNode(va, className);
				va.add(new AnnotationNode("Lorg/springframework/stereotype/Component;"));
			}
		}
		return modified;
	}

	private void addNamedNode(List<AnnotationNode> va, String className) {
		AnnotationNode node = new AnnotationNode("Ljavax/inject/Named;");
		node.values = new ArrayList<Object>();
		node.values.add("value");
		node.values.add(className);
		va.add(node);
	}

	private boolean isBeanType(AnnotationNode node) {
		return "Lcom/macrochen/annotations/BeanType;".equals(node.desc);
	}

	private boolean forMethods(ClassNode cn) {
		boolean exist = false;
		List<MethodNode> methods = cn.methods;
		for (MethodNode mn : methods) {
			List<AnnotationNode> list = mn.visibleAnnotations;
			if (list == null || list.isEmpty()) continue;
			Iterator<AnnotationNode> it = list.iterator();
			String clazzName = null;
			while (it.hasNext()) {
				AnnotationNode node = it.next();
				if (isBeanType(node)) {
					List<Object> values = node.values;
					if (values == null || values.isEmpty()) throw new RuntimeException("beanType value is empty!");

					//getLog().info("find method: " + mn.name + " with beanType");

					clazzName = ((Type) values.get(1)).getClassName();
					it.remove();

					exist = true;
					break;
				}
			}
			if (clazzName != null) {
				addNamedNode(list, clazzName);
				list.add(new AnnotationNode("Ljavax/inject/Inject;"));
			}
		}
		return exist;
	}

	private boolean forFields(ClassNode cn) {
		boolean exist = false;
		List<FieldNode> methods = cn.fields;
		for (FieldNode fn : methods) {
			List<AnnotationNode> list = fn.visibleAnnotations;
			if (list == null || list.isEmpty()) continue;
			Iterator<AnnotationNode> it = list.iterator();
			String clazzName = null;
			while (it.hasNext()) {
				AnnotationNode node = it.next();
				if (isBeanType(node)) {
					List<Object> values = node.values;
					if (values == null || values.isEmpty()) throw new RuntimeException("beanType value is empty!");

					//getLog().info("find method: " + fn.name + " with beanType");

					clazzName = ((Type) values.get(1)).getClassName();
					it.remove();

					exist = true;
					break;
				}
			}
			if (clazzName != null) {
				addNamedNode(list, clazzName);
				list.add(new AnnotationNode("Ljavax/inject/Inject;"));
			}
		}
		return exist;
	}

}