package pf.bot.loader;

import java.io.IOException;
import java.io.InputStream;
import java.security.AllPermission;
import java.security.cert.Certificate;
import java.security.CodeSource;
import java.security.Permissions;
import java.security.ProtectionDomain;
import java.util.HashMap;
import java.util.Map;
import java.util.jar.JarFile;
import java.util.zip.ZipEntry;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;

import pf.bot.bot.Config;
import pf.bot.inject.ScriptModifier;

public class ClientClassLoader extends ClassLoader {

    private ProtectionDomain domain;
    private JarFile jfile;

    private Map<String, Class<?>> classes;

    private ScriptModifier mod;

    public ClientClassLoader() {
        jfile = Config.getJar(Config.client_jar);
        domain = getDomain();
        classes = new HashMap<String, Class<?>>();
        mod = new ScriptModifier();
        mod.load();
    }

    private final Permissions getPermissions() {
        Permissions permissions = new Permissions();
        permissions.add(new AllPermission());
        return permissions;
    }

    private final ProtectionDomain getDomain() {
        CodeSource code = new CodeSource(null, (Certificate[]) null);
        return new ProtectionDomain(code, getPermissions());
    }

    public final Class<?> loadClass(String name) throws ClassNotFoundException {
        if (classes.containsKey(name))
            return classes.get(name);
        Class<?> clas = null;
        try {
            clas = super.findSystemClass(name);
        } catch (Exception e) {
            ZipEntry entry = jfile.getEntry(name.replaceAll("\\.", "/") + ".class");
            if (entry != null) {
                try {
                    InputStream stream = jfile.getInputStream(entry);
                    ClassReader reader = new ClassReader(stream);
                    ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_MAXS);
                    reader.accept(writer, ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);

                    /* class injection method */
                    byte[] buffer = mod.process(name, writer.toByteArray());

                    clas = super.defineClass(name, buffer, 0, buffer.length, domain);

                } catch (IOException e1) {
                    clas = super.loadClass(name);
                }

            }
        }
        classes.put(name, clas);
        return clas;
    }
}