/* ###
 * IP: GHIDRA
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *      http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
//Loads the contents of another program
//into the current program.
//The languages of the two programs must match!
//The information from the source program takes priority
//over the destination program. Meaning, all information in
//the destination program will be removed to make
//room for the information coming from the source
//program.
//@category Program

import java.util.Iterator;

import ghidra.app.script.GhidraScript;
import ghidra.app.util.MemoryBlockUtils;
import ghidra.app.util.NamespaceUtils;
import ghidra.app.util.importer.MessageLog;
import ghidra.program.model.listing.*;
import ghidra.program.model.mem.*;
import ghidra.program.model.symbol.*;

public class MergeTwoProgramsScript extends GhidraScript {

	@Override
	protected void run() throws Exception {

		if (currentProgram == null) {
			printerr("Please open a program first!");
			return;
		}

		Program otherProgram = askProgram("Select program from which to merge: ");

		if (otherProgram == null) {
			printerr("Please select the other program first!");
			return;
		}

		try {

			if (!currentProgram.getLanguage().equals(otherProgram.getLanguage())) {
				printerr("Incompatible program languages!");
				return;
			}

			if (currentProgram.getMemory().intersects(otherProgram.getMemory())) {
				printerr("Memory map of current program must be disjoint from other program!");
				return;
			}

			openProgram(currentProgram);

			mergeMemory(currentProgram, otherProgram);
			mergeSymbols(currentProgram, otherProgram);
			mergeBookmarks(currentProgram, otherProgram);
			mergeComments(currentProgram, otherProgram);
			mergeData(currentProgram, otherProgram);
			mergeInstructions(currentProgram, otherProgram);
			mergeEquates(currentProgram, otherProgram);
			mergeReferences(currentProgram, otherProgram);
		}
		finally {
			otherProgram.release(this);
		}
	}

	private void mergeReferences(Program currProgram, Program otherProgram) {
		monitor.setMessage("Merging references...");
		ReferenceManager currentReferenceManager = currProgram.getReferenceManager();
		ReferenceManager otherReferenceManager = otherProgram.getReferenceManager();
		ReferenceIterator otherReferenceIterator =
			otherReferenceManager.getReferenceIterator(otherProgram.getMinAddress());
		while (otherReferenceIterator.hasNext()) {
			if (monitor.isCancelled()) {
				break;
			}
			Reference otherReference = otherReferenceIterator.next();
			if (otherReference.isStackReference()) {
				continue;
			}
			currentReferenceManager.addReference(otherReference);
		}
	}

	private void mergeInstructions(Program currProgram, Program otherProgram) {
		monitor.setMessage("Merging instructions...");
		Listing currentListing = currProgram.getListing();
		Listing otherListing = otherProgram.getListing();
		InstructionIterator otherInstructions = otherListing.getInstructions(true);
		while (otherInstructions.hasNext()) {
			if (monitor.isCancelled()) {
				break;
			}
			Instruction otherInstruction = otherInstructions.next();
			if (currentListing.isUndefined(otherInstruction.getMinAddress(),
				otherInstruction.getMaxAddress())) {
				disassemble(otherInstruction.getMinAddress());
			}
		}
	}

	private void mergeEquates(Program currProgram, Program otherProgram) throws Exception {
		monitor.setMessage("Merging equates...");
		EquateTable currentEquateTable = currProgram.getEquateTable();
		EquateTable otherEquateTable = otherProgram.getEquateTable();
		Iterator<Equate> otherEquates = otherEquateTable.getEquates();
		while (otherEquates.hasNext()) {
			if (monitor.isCancelled()) {
				break;
			}
			Equate otherEquate = otherEquates.next();
			Equate currentEquate =
				currentEquateTable.createEquate(otherEquate.getName(), otherEquate.getValue());
			EquateReference[] otherEquateReferences = otherEquate.getReferences();
			for (EquateReference otherEquateReference : otherEquateReferences) {
				if (monitor.isCancelled()) {
					break;
				}
				currentEquate.addReference(otherEquateReference.getAddress(),
					otherEquateReference.getOpIndex());
			}
		}
	}

	private void mergeData(Program currProgram, Program otherProgram) throws Exception {
		monitor.setMessage("Merging data...");
		Listing currentListing = currProgram.getListing();
		Listing otherListing = otherProgram.getListing();
		DataIterator otherDataIterator = otherListing.getDefinedData(true);
		while (otherDataIterator.hasNext()) {
			if (monitor.isCancelled()) {
				break;
			}
			Data otherData = otherDataIterator.next();
			if (currentListing.isUndefined(otherData.getMinAddress(), otherData.getMaxAddress())) {
				currentListing.createData(otherData.getMinAddress(), otherData.getDataType());
			}
		}
	}

	private void mergeComments(Program currProgram, Program otherProgram) throws Exception {
		monitor.setMessage("Merging comments...");
		int[] commentTypes = { CodeUnit.EOL_COMMENT, CodeUnit.PRE_COMMENT, CodeUnit.POST_COMMENT,
			CodeUnit.PLATE_COMMENT, CodeUnit.REPEATABLE_COMMENT, };
		Listing currentListing = currProgram.getListing();
		Listing otherListing = otherProgram.getListing();
		CodeUnitIterator otherCodeUnits = otherListing.getCodeUnits(true);
		while (otherCodeUnits.hasNext()) {
			if (monitor.isCancelled()) {
				break;
			}
			CodeUnit otherCodeUnit = otherCodeUnits.next();
			for (int commentType : commentTypes) {
				if (monitor.isCancelled()) {
					break;
				}
				String otherComment = otherCodeUnit.getComment(commentType);
				if (otherComment != null) {
					currentListing.setComment(otherCodeUnit.getAddress(), commentType,
						otherComment);
				}
			}
		}
	}

	private void mergeBookmarks(Program currProgram, Program otherProgram) {
		monitor.setMessage("Merging bookmarks...");
		BookmarkManager currentBookmarkManager = currProgram.getBookmarkManager();
		BookmarkManager otherBookmarkManager = otherProgram.getBookmarkManager();
		Iterator<Bookmark> otherBookmarks = otherBookmarkManager.getBookmarksIterator();
		while (otherBookmarks.hasNext()) {
			if (monitor.isCancelled()) {
				break;
			}
			Bookmark otherBookmark = otherBookmarks.next();
			currentBookmarkManager.setBookmark(otherBookmark.getAddress(),
				otherBookmark.getTypeString(), otherBookmark.getCategory(),
				otherBookmark.getComment());
		}
	}

	private void mergeSymbols(Program currProgram, Program otherProgram) throws Exception {
		monitor.setMessage("Merging symbols...");
		SymbolTable currentSymbolTable = currProgram.getSymbolTable();
		SymbolTable otherSymbolTable = otherProgram.getSymbolTable();
		SymbolIterator otherSymbols = otherSymbolTable.getAllSymbols(false);
		while (otherSymbols.hasNext()) {
			if (monitor.isCancelled()) {
				break;
			}
			Symbol otherSymbol = otherSymbols.next();
			if (otherSymbol.isDynamic()) {
				continue;
			}
			try {
				Namespace otherNamespace = otherSymbol.getParentNamespace();
				Namespace currentNamespace =
					mirrorNamespace(currProgram, otherProgram, otherNamespace);
				if (otherSymbol.getSymbolType() == SymbolType.FUNCTION) {
					Function otherFunction =
						otherProgram.getListing().getFunctionAt(otherSymbol.getAddress());
					currProgram.getListing()
							.createFunction(otherSymbol.getName(), currentNamespace,
								otherFunction.getEntryPoint(), otherFunction.getBody(),
								SourceType.USER_DEFINED);
				}
				else {
					currentSymbolTable.createLabel(otherSymbol.getAddress(), otherSymbol.getName(),
						currentNamespace, SourceType.USER_DEFINED);
				}
			}
			catch (Exception e) {
				printerr("Unable to create symbol: " + otherSymbol.getName());
			}
		}
	}

	private Namespace mirrorNamespace(Program currProgram, Program otherProgram,
			Namespace otherNamespace) throws Exception {
		if (otherNamespace == null) {
			return currProgram.getGlobalNamespace();
		}
		SourceType source = SourceType.USER_DEFINED;//this will be default, since we are running a script!
		try {
			source = otherNamespace.getSymbol().getSource();
		}
		catch (Exception e) {
		}
		return NamespaceUtils.createNamespaceHierarchy(otherNamespace.getName(true), null,
			currProgram, source);
	}

	private void mergeMemory(Program currProgram, Program otherProgram) throws Exception {
		monitor.setMessage("Merging memory...");
		Memory otherMemory = otherProgram.getMemory();
		MemoryBlock[] otherBlocks = otherMemory.getBlocks();
		MessageLog log = new MessageLog();
		for (MemoryBlock otherBlock : otherBlocks) {
			if (monitor.isCancelled()) {
				break;
			}
			if (otherBlock.getType() != MemoryBlockType.DEFAULT) {
				printerr("Unhandled memory block type: " + otherBlock.getName());
				continue;
			}
			if (otherBlock.isInitialized()) {
				MemoryBlockUtils.createInitializedBlock(currProgram, false, otherBlock.getName(),
					otherBlock.getStart(), otherBlock.getData(), otherBlock.getSize(),
					otherBlock.getComment(), otherBlock.getSourceName(), otherBlock.isRead(),
					otherBlock.isWrite(), otherBlock.isExecute(), log, monitor);
			}
			else {
				MemoryBlockUtils.createUninitializedBlock(currProgram, false, otherBlock.getName(),
					otherBlock.getStart(), otherBlock.getSize(), otherBlock.getComment(),
					otherBlock.getSourceName(), otherBlock.isRead(), otherBlock.isWrite(),
					otherBlock.isExecute(), log);
			}
		}
	}

}