borked lmao

BlockGenerator.java

package org.nullbool.api.obfuscation.cfg;

import static org.objectweb.asm.tree.AbstractInsnNode.JUMP_INSN;
import static org.objectweb.asm.tree.AbstractInsnNode.LOOKUPSWITCH_INSN;
import static org.objectweb.asm.tree.AbstractInsnNode.TABLESWITCH_INSN;

import java.io.File;
import java.io.FileOutputStream;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.Deque;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import org.nullbool.api.obfuscation.TarjanSCC.TarjanBlock;
import org.nullbool.api.obfuscation.cfg.BlockGenerator.ControlFlowGraph.SSC;
import org.nullbool.api.util.ClassStructure;
import org.nullbool.api.util.InstructionUtil;
import org.nullbool.api.util.LabelHelper;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.ClassNode;
import org.objectweb.asm.tree.JumpInsnNode;
import org.objectweb.asm.tree.LabelNode;
import org.objectweb.asm.tree.LookupSwitchInsnNode;
import org.objectweb.asm.tree.MethodNode;
import org.objectweb.asm.tree.TableSwitchInsnNode;
import org.objectweb.asm.util.Printer;
import org.topdank.banalysis.asm.insn.InstructionPrinter;

/**
 * @author Bibl (don't ban me pls)
 * @created 25 Jul 2015 02:18:09
 */
public class BlockGenerator implements Opcodes {

	/* public static class TarjanBlock {
		public final FlowBlock block;
		public boolean marked;
		public int index;
		public int lowlink;
		public boolean onstack;

		public TarjanBlock(FlowBlock block) {
			this.block = block;
			marked = false;
			index = 0;
			lowlink = 0;
			onstack = false;
		}

		@Override
		public String toString() {
			StringBuilder sb = new StringBuilder();
			sb.append("Block: ").append(block.id()).append(", marked=").append(marked);
			sb.append(", index=").append(index).append(", lowlink=").append(lowlink);
			sb.append(", onstack=").append(onstack);
			return sb.toString();
		}
	} */

	public static class TarjanSCC {
		private final Map<FlowBlock, TarjanBlock> blocks;
		private final Deque<TarjanBlock> stack;
		private final Set<SSC> sscs;
		private int index;

		public TarjanSCC() {
			blocks = new HashMap<FlowBlock, TarjanBlock>();
			stack = new LinkedList<TarjanBlock>();
			sscs = new HashSet<SSC>();
			index = 0;
		}

		public Set<SSC> tarjan(ControlFlowGraph graph) {
			List<FlowBlock> blocks = graph.blocks;
			if (blocks == null || blocks.isEmpty())
				return null;

			this.blocks.clear();
			stack.clear();
			sscs.clear();
			index = 0;

			for (FlowBlock b : blocks) {
				this.blocks.put(b, new TarjanBlock(b));
			}

			for (Entry<FlowBlock, TarjanBlock> e : this.blocks.entrySet()) {
				TarjanBlock b = e.getValue();
				if (!b.marked) {
					strongconnect(b);
				}
			}

			return sscs;
		}

		private void strongconnect(TarjanBlock b) {
			b.marked = true;
			b.index = index;
			b.lowlink = index;
			index++;
			stack.push(b);

			b.onstack = true;

			for (FlowBlock _s : b.block.successors()) {
				if (!(_s instanceof DummyExitBlock)) {
					TarjanBlock s = blocks.get(_s);
					if (!s.marked) {
						strongconnect(s);
						b.lowlink = Math.min(b.lowlink, s.lowlink);
					} else if (s.onstack) {
						b.lowlink = Math.min(b.lowlink, s.index);
					}
				}
			}

			for (FlowBlock _s : b.block.exceptionSuccessors()) {
				if (!(_s instanceof DummyExitBlock)) {
					TarjanBlock s = blocks.get(_s);
					if (!s.marked) {
						strongconnect(s);
						b.lowlink = Math.min(b.lowlink, s.lowlink);
					} else if (s.onstack) {
						b.lowlink = Math.min(b.lowlink, s.index);
					}
				}
			}

			if (b.lowlink == b.index) {
				SSC ssc = new SSC();
				TarjanBlock w;
				do {
					w = stack.pop();
					ssc.blocks.add(w.block);
					w.onstack = false;
				} while(w != b);
				sscs.add(ssc);
			}
		}
	}

	public static class ControlFlowGraph {
		public final List<FlowBlock> blocks;
		public final Map<LabelNode, FlowBlock> labelBoundaries;
		public final Map<AbstractInsnNode, FlowBlock> boundaries;

		public ControlFlowGraph(MethodNode m) {
			blocks = new ArrayList<FlowBlock>();
			labelBoundaries = new HashMap<LabelNode, FlowBlock>();
			boundaries = new HashMap<AbstractInsnNode, FlowBlock>();

			create(m);
		}

		public void destroy() {
			blocks.clear();
			labelBoundaries.clear();
			boundaries.clear();
		}

		public void recreate(MethodNode m) {
			destroy();
			create(m);
		}

		void create(MethodNode m) {
			int count = 1;
			FlowBlock block = new FlowBlock(LabelHelper.createBlockName(count));
			for(AbstractInsnNode ain : m.instructions.toArray()) {
				if(!(ain instanceof LabelNode)) {
					block.insns().add(ain);
				}

				int type = ain.type();
				if(ain instanceof LabelNode || ain.opcode() == ATHROW || type == JUMP_INSN || type == TABLESWITCH_INSN || type == LOOKUPSWITCH_INSN || isExit(ain.opcode())) {
					blocks.add(block);
					block = new FlowBlock(LabelHelper.createBlockName(++count));
				}
			}

			if(block.size() > 0) {
				blocks.add(block);
			}

			Iterator<FlowBlock> it = blocks.iterator();
			while(it.hasNext()) {
				FlowBlock b = it.next();
				if(b.size() == 0) {
					it.remove();
					continue;
				}

				boundaries.put(b.first(), b);
			}

			for(int i=0; i < blocks.size(); i++) {
				FlowBlock b = blocks.get(i);
				AbstractInsnNode last = b.last();
				if(last instanceof JumpInsnNode) {
					JumpInsnNode jin = (JumpInsnNode) last;

					if(InstructionUtil.isConditional(jin.opcode())) {
						// add the fall-through block as a successor
						if((i + 1) <= blocks.size()) {
							FlowBlock next = blocks.get(i + 1);
							linkBlocks(b, next);
						}
					} else if(!InstructionUtil.isUnconditional(jin.opcode())) {
						throw new RuntimeException("jump: " + Printer.OPCODES[jin.opcode()] + " is not any?");
					}

					// add jump to block as a successor
					{
						LabelNode label = jin.label;
						map(b, label, "unaligned jump block boundary", "jump label blocks don't match");
					}

				} else if(last instanceof TableSwitchInsnNode) {
					TableSwitchInsnNode tsin = (TableSwitchInsnNode) last;
					// add switch jumps as successors
					for(LabelNode label : tsin.labels) {
						map(b, label, "unaligned tableswitch block boundary", "tableswitch label blocks don't match");
					}

					// add the default case jump as a successor
					LabelNode label = tsin.dflt;
					map(b, label, "unaligned tableswitch default block boundary", "default tableswitch label blocks don't match");
				} else if(last instanceof LookupSwitchInsnNode) {
					LookupSwitchInsnNode lsin = (LookupSwitchInsnNode) last;
					// add switch jumps as successors
					for(LabelNode label : lsin.labels) {
						map(b, label, "unaligned lookupwitch block boundary", "lookupwitch label blocks don't match");
					}

					// add the default case jump as a successor
					LabelNode label = lsin.dflt;
					map(b, label, "unaligned lookupwitch default block boundary", "default lookupwitch label blocks don't match");
				} else if(isExit(last.opcode())) {
					// don't add a fall through successor
				} else {
					// add fall through successors
					if((i + 1) <= blocks.size()) {
						FlowBlock next = blocks.get(i + 1);
						linkBlocks(b, next);
					}
				}
			}
		}

		public void linkBlocks(FlowBlock pred, FlowBlock succ) {
			pred.addSuccessor(succ);
			succ.addPredecessor(pred);
		}

		public void map(FlowBlock block, LabelNode label, String errMsg1, String errMsg2) {
			AbstractInsnNode jump = label.getNext();
			FlowBlock jblock = boundaries.get(jump);
			if(jblock == null) {
				throw new RuntimeException(errMsg1);
			}

			linkBlocks(block, jblock);

			if(labelBoundaries.containsKey(label)) {
				FlowBlock prev = labelBoundaries.get(label);
				if(prev != jblock) {
					throw new RuntimeException(errMsg2);
				}
			} else {
				labelBoundaries.put(label, jblock);
			}
		}

		/* static class DeobfuscationStatus {
			final List<FlowBlock> blocks;
			final Set<FlowBlock> visited;
			final Stack<FlowBlock> stack;
			final int[] index;
			final int[] low;
			int currentDfsNum;

			DeobfuscationStatus(List<FlowBlock> blocks) {
				this.blocks = blocks;
				visited = new HashSet<FlowBlock>();
				stack = new Stack<FlowBlock>();
				int count = blocks.size();
				index = new int[count];
				low = new int[count];
				currentDfsNum = 0;
			}
		} */

		static class Block {
			final List<AbstractInsnNode> insns;
			final List<FlowEdge> in;
			final List<FlowEdge> out;
			final String id;
			
			Block(String id) {
				this.id = id;
				
				insns = new ArrayList<AbstractInsnNode>();
				in = new ArrayList<FlowEdge>();
				out = new ArrayList<FlowEdge>();
			}
		}
		
		static class FlowEdge {
			final FlowBlock src;
			final FlowBlock dst;
			final EdgeType type;
			
			FlowEdge(FlowBlock src, FlowBlock dst, EdgeType type) {
				this.src = src;
				this.dst = dst;
				this.type = type;
			}
			
			FlowEdge(FlowBlock src, FlowBlock dst) {
				this(src, dst, type(src, dst));
			}
			
			private static EdgeType type(FlowBlock src, FlowBlock dst) {
				int op = src.lastOpcode();
				if(InstructionUtil.isUnconditional(op)) {
					
				}
			}
		}
		
		static enum EdgeType {
			IMMEDIATE(), 
			CONDITIONAL(), 
			UNCONDITIONAL()
		}
		
		static class SSC {
			public final List<FlowBlock> blocks = new ArrayList<FlowBlock>();

			@Override
			public String toString() {
				StringBuilder sb = new StringBuilder("SSC: { ");
				Iterator<FlowBlock> it = blocks.iterator();
				while (it.hasNext()) {
					FlowBlock b = it.next();
					sb.append(b.id());
					if (it.hasNext()) {
						sb.append(',');
					}
					sb.append(' ');
				}
				sb.append('}');
				return sb.toString();
			}
		}

		static class Frame {
			static final int TODO = Integer.MIN_VALUE + 1;
			static final int DONE = Integer.MIN_VALUE + 2;
			final List<FlowBlock> newOrder;
			final List<FlowBlock> blocks;
			final Deque<FlowBlock> stack;
			final int[] index;
			final int[] low;
			int dfsnum;

			Frame(List<FlowBlock> newOrder, List<FlowBlock> blocks, Deque<FlowBlock> stack, int size, int dfsnum) {
				this.newOrder = newOrder;
				this.blocks = blocks;
				this.stack = stack;
				this.dfsnum = dfsnum;
				index = new int[size];
				low = new int[size];
			}

			int indexOf(FlowBlock b) {
				return blocks.indexOf(b);
			}
		}

		public void visit(FlowBlock cur, Frame frame, Collection<FlowBlock> consider) {
			int index = frame.indexOf(cur);
			int dfs = frame.dfsnum;
			frame.low[index] = dfs;
			frame.index[index] = dfs;
			frame.dfsnum++;

			frame.stack.push(cur);

			List<FlowBlock> succs = new ArrayList<FlowBlock>();
			succs.addAll(cur.successors());
			succs.retainAll(consider);
			Collections.sort(succs, new Comparator<FlowBlock>(){
				@Override
				public int compare(FlowBlock o1, FlowBlock o2) {
					return Integer.compare(frame.blocks.indexOf(o2), frame.blocks.indexOf(o1));
				}
			});

			for(FlowBlock s : succs) {
				int sIndex = frame.indexOf(s);
				if(frame.index[sIndex] == Frame.TODO) {
					visit(s, frame, consider);
					frame.low[index] = Math.min(frame.low[index], frame.low[sIndex]);
				} else if(frame.index[sIndex] == Frame.DONE) {
					/* Do nothing */
				} else {
					frame.low[index] = Math.min(frame.low[index], frame.index[sIndex]);
				}
			}

			if(frame.low[index] == frame.index[index]) {
				SSC ssc = new SSC();
				FlowBlock w;
				do {
					w = frame.stack.pop();
					ssc.blocks.add(w);
					int pIndex = frame.indexOf(w);
					frame.index[pIndex] = Frame.DONE;
				} while(w != cur);

				System.out.println(ssc);

				if(ssc.blocks.size() == 1) {
					frame.newOrder.add(0, cur);
				} else {
					frame.newOrder.addAll(0, deobfuscate(ssc.blocks.get(0), ssc.blocks));
				}
			}
		}

		public List<FlowBlock> deobfuscate(FlowBlock v, Collection<FlowBlock> consider) {			
			Frame frame = new Frame(new ArrayList<FlowBlock>(), blocks, new LinkedList<FlowBlock>(), blocks.size(), 0);

			for(FlowBlock b : consider) {
				int index = frame.indexOf(b);
				frame.index[index] = Frame.TODO;
			}

			int vIndex = frame.indexOf(v);
			frame.index[vIndex] = Frame.DONE;

			List<FlowBlock> succs = new ArrayList<FlowBlock>();
			succs.addAll(v.successors());
			succs.retainAll(consider);
			Collections.sort(succs, new Comparator<FlowBlock>(){
				@Override
				public int compare(FlowBlock o1, FlowBlock o2) {
					return Integer.compare(frame.blocks.indexOf(o2), frame.blocks.indexOf(o1));
				}
			});

			for(FlowBlock s : succs) {
				int sIndex = frame.indexOf(s);
				if(frame.index[sIndex] == Frame.TODO) {
					visit(s, frame, consider);
				}
			}

			frame.newOrder.add(0, v);

			return frame.newOrder;
		}
		
		public List<FlowBlock> reorder(List<FlowBlock> blocks) {
			List<FlowBlock> order = new ArrayList<FlowBlock>();
			Deque<FlowBlock> stack = new LinkedList<FlowBlock>();
			stack.push(blocks.get(0));
			
			while(!stack.isEmpty()) {
				FlowBlock cur = stack.pop();
				if(!order.contains(cur)) {
					order.add(cur);
					List<FlowBlock> toAdd = new ArrayList<FlowBlock>();
					
//					if()
				}
			}
			
			return order;
		}

		public static boolean isExit(int op) {
			return (op >= IRETURN && op <= RETURN);
		}

		public String toString(List<FlowBlock> blocks) {
			StringBuilder sb = new StringBuilder();
			for(FlowBlock b : blocks) {
				sb.append(b.toVerboseString(labelBoundaries)).append('\n');
			}
			return sb.toString();
		}

		public String toDotGraphString() {
			StringBuilder sb = new StringBuilder();
			sb.append("digraph g {").append('\n');
			for(int i=0; i < blocks.size(); i++) {
				FlowBlock b = blocks.get(i);
				sb.append(i).append('[');
				sb.append("shape = box, ").append("label = ").append('"');

				StringBuilder labelSb = new StringBuilder();
				for(AbstractInsnNode ain : b.insns()) {
					labelSb.append(ain.method.instructions.indexOf(ain));
					labelSb.append(' ');
					labelSb.append(Printer.OPCODES[ain.opcode()]);
					labelSb.append("\\l");
				}
				sb.append(labelSb.toString());
				sb.append('"');
				sb.append(']');
				sb.append('\n');
			}

			for(int i=0; i < blocks.size(); i++) {
				FlowBlock b = blocks.get(i);
				for(FlowBlock s : b.successors()) {
					sb.append(i).append(" -> ").append(blocks.indexOf(s)).append('\n');
				}
			}

			sb.append("}");
			return sb.toString();
		}

		@Override
		public String toString() {
			return toString(blocks); 
		}
	}

	public static void main(String[] args) throws Exception {
		ClassNode cn = ClassStructure.create(BlockGenerator.class.getResourceAsStream("/TestClass.class"));

		for(MethodNode m : cn.methods) {
			if(!m.name.equals("method796"))
				continue;

			System.out.println(m);
			InstructionPrinter.consolePrint(m);
			//			InstructionPrinter.consolePrint(m);

			ControlFlowGraph cfg = new ControlFlowGraph(m);
			System.out.println(cfg);
			System.out.println("---------------------------------------------------------------------");
			Map<FlowBlock, LabelNode> labels = new HashMap<FlowBlock, LabelNode>();
			for(Entry<LabelNode, FlowBlock> e : cfg.labelBoundaries.entrySet()) {
				labels.put(e.getValue(), e.getKey());
			}

			List<FlowBlock> order = cfg.deobfuscate(cfg.blocks.get(0), cfg.blocks);
			cfg.blocks.clear();
			cfg.blocks.addAll(order);

			m.instructions.removeAll(true);
			for(FlowBlock b : order) {
				LabelNode label = labels.get(b);
				if(label != null) {
					b.insns().add(0, label);
				}

				b.transfer(m.instructions);
			}

			System.out.println(cfg.toString(order));



			//			List<FlowBlock> order = new ArrayList<FlowBlock>();
			//			Iterator<FlowBlock> it = new DFSIterator(cfg.blocks.get(0));
			//			while(it.hasNext()) {
			//				FlowBlock b = it.next();
			//				if(!order.contains(b)) {
			//					SSC ssc = null;
			//					for(SSC s : set) {
			//						if(s.blocks.contains(b)) {
			//							ssc = s;
			//							break;
			//						}
			//					}
			//					if(ssc != null) {
			//						order.addAll(ssc.blocks);
			//					} else {
			//						order.add(b);
			//					}
			//				}
			//			}
			//			
			//			System.out.println();

			//			Path path = Paths.get(new File("C:/Users/Bibl/Desktop/tests/test.dot").toURI());
			//			Files.write(path, cfg.toDotGraphString().getBytes(), StandardOpenOption.CREATE_NEW);

			//			for(FlowBlock b : order) {
			//				System.out.println(b.toVerboseString(cfg.labelBoundaries));
			//			}
			//			System.out.println(cfg.toString(cfg.blocks));
			//			m.instructions.removeAll(true);
			//			for(FlowBlock b : cfg.deobfuscate(cfg.blocks.get(0), cfg.blocks)) {
			//				System.out.println(b.id());
			//			}
			//			
			//			for(FlowBlock b : new Sorter().topologicalSort(cfg)) {
			//				System.out.println(b.id());
			//			}
		}

		ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
		cn.accept(cw);
		FileOutputStream fos = new FileOutputStream(new File("C:/Users/Bibl/Desktop/TestClass2.class"));
		fos.write(cw.toByteArray());
		fos.close();

		//		CheckClassAdapter cca = new CheckClassAdapter(cn);
		//		cn.accept(cca);

	}
}