From 2aedcc730c6c567c8eb32c7558ef44ba3192408e Mon Sep 17 00:00:00 2001
From: de4dot <de4dot@gmail.com>
Date: Sun, 4 Nov 2012 23:24:12 +0100
Subject: [PATCH] Preserve tokens if necessary

---
 de4dot.code/AssemblyModule.cs |  4 +++-
 de4dot.code/ObfuscatedFile.cs | 13 ++++++++++---
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/de4dot.code/AssemblyModule.cs b/de4dot.code/AssemblyModule.cs
index 4e1a273e..29b8dd99 100644
--- a/de4dot.code/AssemblyModule.cs
+++ b/de4dot.code/AssemblyModule.cs
@@ -51,10 +51,12 @@ namespace de4dot.code {
 			return module;
 		}
 
-		public void save(string newFilename, bool updateMaxStack, IModuleWriterListener writerListener) {
+		public void save(string newFilename, bool preserveTokens, bool updateMaxStack, IModuleWriterListener writerListener) {
 			var writerOptions = new ModuleWriterOptions(module, writerListener);
 			if (!updateMaxStack)
 				writerOptions.MetaDataOptions.Flags |= MetaDataFlags.KeepOldMaxStack;
+			if (preserveTokens)
+				writerOptions.MetaDataOptions.Flags |= MetaDataFlags.PreserveTokens | MetaDataFlags.PreserveUSOffsets | MetaDataFlags.PreserveExtraSignatureData;
 			module.Write(newFilename, writerOptions);
 		}
 
diff --git a/de4dot.code/ObfuscatedFile.cs b/de4dot.code/ObfuscatedFile.cs
index b2aa7b22..1b611995 100644
--- a/de4dot.code/ObfuscatedFile.cs
+++ b/de4dot.code/ObfuscatedFile.cs
@@ -304,9 +304,13 @@ namespace de4dot.code {
 			return detected;
 		}
 
+		bool ShouldPreserveTokens() {
+			return options.KeepObfuscatorTypes || deob.Type == "un";
+		}
+
 		public void save() {
 			Log.n("Saving {0}", options.NewFilename);
-			assemblyModule.save(options.NewFilename, options.ControlFlowDeobfuscation, deob as IModuleWriterListener);
+			assemblyModule.save(options.NewFilename, ShouldPreserveTokens(), options.ControlFlowDeobfuscation, deob as IModuleWriterListener);
 		}
 
 		IList<MethodDef> getAllMethods() {
@@ -524,7 +528,7 @@ namespace de4dot.code {
 			if (!options.ControlFlowDeobfuscation) {
 				// If it's the unknown type, we don't remove any types that could cause Mono.Cecil
 				// to throw an exception.
-				if (deob.Type == "un" || options.KeepObfuscatorTypes)
+				if (ShouldPreserveTokens())
 					return;
 			}
 
@@ -590,7 +594,10 @@ namespace de4dot.code {
 				cflowDeobfuscator.deobfuscate();
 
 			if (options.ControlFlowDeobfuscation) {
-				numRemovedLocals = blocks.optimizeLocals();
+				// Don't remove any locals if we should preserve tokens or we won't be able
+				// to always preserve StandAloneSig tokens.
+				if (!ShouldPreserveTokens())
+					numRemovedLocals = blocks.optimizeLocals();
 				blocks.repartitionBlocks();
 			}