From bde21731933bbfb01bf85725ec004e4652a85480 Mon Sep 17 00:00:00 2001
From: anders_k <anders_k@212acab6-be3b-0410-9dea-997c60f758d6>
Date: Thu, 26 Oct 2017 16:37:42 +0000
Subject: [PATCH] Only print the PE flags string for actual PE executable
 files. Now also zeros the checksum just in case.

git-svn-id: https://svn.code.sf.net/p/nsis/code/NSIS/trunk@6943 212acab6-be3b-0410-9dea-997c60f758d6
---
 Docs/src/history.but |  2 +-
 SCons/utils.py       | 99 ++++++++++++++++++++++++++++++++------------
 SConstruct           | 19 +++++----
 3 files changed, 85 insertions(+), 35 deletions(-)

diff --git a/Docs/src/history.but b/Docs/src/history.but
index 7c2fb70c..40083e91 100644
--- a/Docs/src/history.but
+++ b/Docs/src/history.but
@@ -26,7 +26,7 @@ Released on ??? ??rd, 20??
 
 \b Added IntOp and System::Int64Op >>> operator
 
-\b Plug-ins now set the ASLR, DEP, LAA and NOSEH PE flags (\W{http://sf.net/p/nsis/bugs/1188}{bug #1188})
+\b Plug-ins now set the ASLR, DEP, LAA, NOSEH and TS PE flags (\W{http://sf.net/p/nsis/bugs/1188}{bug #1188})
 
 \b MakeNSIS exits with code 0 for various information commands (\W{http://sf.net/p/nsis/bugs/1193}{bug #1193})
 
diff --git a/SCons/utils.py b/SCons/utils.py
index b0f250d4..1d561e90 100644
--- a/SCons/utils.py
+++ b/SCons/utils.py
@@ -131,38 +131,85 @@ def GetOptionOrEnv(name, defval = None):
 		return os.environ[name]
 	return defval
 
-def EnablePESecurityFlagsHelper(filepath):
+import struct
+def FileUnpackRead(pack, size, f, fpos=None, defval=None):
+	r = defval
+	try:
+		if not fpos is None: f.seek(fpos)
+		r = struct.unpack(pack, f.read(size))[0]
+	finally:
+		return r
+def ReadU16LE(f, fpos=None, defval=None):
+	return FileUnpackRead("<H", 2, f, fpos, defval)
+def ReadU32LE(f, fpos=None, defval=None):
+	return FileUnpackRead("<I", 4, f, fpos, defval)
+def WriteU16LE(f, v, fpos):
+	if not fpos is None: f.seek(fpos)
+	f.write(struct.pack("<H", v))
+def WriteU32LE(f, v, fpos):
+	if not fpos is None: f.seek(fpos)
+	f.write(struct.pack("<I", v))
+
+class MSPE:
+	def __init__(self, path=None, open_for_write=False):
+		self._f = None
+		self._path = None
+		self.NTHOffset = 0
+		self.NTOHMagic = None
+		self.IsPEP = False # PE+ A.K.A IMAGE_NT_HEADERS64
+		if not path is None:
+			self.Open(path, open_for_write)
+	def __del__(self): self.Close()
+	def Close(self):
+		if getattr(self, '_f', None):
+			self._f.close()
+			self._f = None
+	def Open(self, path, open_for_write=False):
+		mode = "rb"
+		if open_for_write: mode = "r+b"
+		self._path = path
+		f = self._f = open(path, mode)
+		if not 0x5A4D == ReadU16LE(f, 0): return # IMAGE_DOS_SIGNATURE?
+		fanew = ReadU32LE(f, 60)
+		if not 0x00004550 == ReadU32LE(f, fanew): return # IMAGE_NT_SIGNATURE?
+		self.NTHOffset = fanew
+		self.NTOHMagic = ReadU16LE(f, fanew+4+20)
+		self.IsPEP = 0x20b == self.NTOHMagic # IMAGE_NT_OPTIONAL_HDR64_MAGIC?
+	def ReadCharacteristics(self):
+		return ReadU16LE(self._f, self.NTHOffset+4+18)
+	def WriteCharacteristics(self, value):
+		WriteU16LE(self._f, value, self.NTHOffset+4+18)
+	def ReadDllCharacteristics(self):
+		return ReadU16LE(self._f, self.NTHOffset+4+20+70)
+	def WriteDllCharacteristics(self, value):
+		WriteU16LE(self._f, value, self.NTHOffset+4+20+70)
+	def WriteChecksum(self, value):
+		WriteU32LE(self._f, value, self.NTHOffset+4+20+64)
+
+def IsPEExecutable(pe):
+	if not isinstance(pe, MSPE): pe = MSPE(pe)
+	if pe.ReadCharacteristics() & 0x0002: return True # IMAGE_FILE_EXECUTABLE_IMAGE?
+
+def SetPESecurityFlagsWorker(filepath):
 	"""
 	Sets the [HE]ASLR, DEP and LAA flags in the PE header
 	"""
-	import struct
-	def ReadU16LE(f, fpos):
-		if not fpos is None: f.seek(fpos)
-		return struct.unpack("<H", f.read(2))[0]
-	def ReadU32LE(f, fpos):
-		if not fpos is None: f.seek(fpos)
-		return struct.unpack("<I", f.read(4))[0]
-	def WriteU16LE(f, v, fpos):
-		if not fpos is None: f.seek(fpos)
-		f.write(struct.pack("<H", v))
-	f = open(filepath, "r+b")
+	pe = MSPE(filepath, open_for_write=True)
 	try:
-		if not 0x5A4D == ReadU16LE(f, 0): return
-		pepos = ReadU32LE(f, 60)
-		if not 0x00004550 == ReadU32LE(f, pepos): return
-		pe64 = 0x20b == ReadU16LE(f, pepos+4+20) # IMAGE_NT_OPTIONAL_HDR64_MAGIC?
-		ifh_c = ReadU16LE(f, pepos+4+18)
+		if not IsPEExecutable(pe): return
+		ifh_c = pe.ReadCharacteristics()
 		ifh_c |= 0x0020 # +IMAGE_FILE_LARGE_ADDRESS_AWARE
-		WriteU16LE(f, ifh_c, pepos+4+18)
-		ioh_dc = ReadU16LE(f, pepos+4+20+70)
-		ioh_dc |= 0x0100 # +IMAGE_DLLCHARACTERISTICS_NX_COMPAT
+		pe.WriteCharacteristics(ifh_c)
+		ioh_dc = pe.ReadDllCharacteristics()
+		ioh_dc |= 0x0100 # +IMAGE_DLLCHARACTERISTICS_NX_COMPAT (DEP)
 		ioh_dc |= 0x0400 # +IMAGE_DLLCHARACTERISTICS_NO_SEH
+		ioh_dc |= 0x8000 # +IMAGE_DLLCHARACTERISTICS_TERMINAL_SERVER_AWARE (TODO: Should we set this on .DLLs?)
 		if not (ifh_c & 0x0001): # IMAGE_FILE_RELOCS_STRIPPED?
-			ioh_dc |= 0x0040 # +IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE
-			if pe64: ioh_dc |= 0x0020 # +IMAGE_DLLCHARACTERISTICS_HIGH_ENTROPY_VA
-		# TODO: IMAGE_DLLCHARACTERISTICS_TERMINAL_SERVER_AWARE?
-		WriteU16LE(f, ioh_dc, pepos+4+20+70)
+			ioh_dc |= 0x0040 # +IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE (ASLR)
+			if pe.IsPEP: ioh_dc |= 0x0020 # +IMAGE_DLLCHARACTERISTICS_HIGH_ENTROPY_VA (HEASLR)
+		pe.WriteDllCharacteristics(ioh_dc)
+		pe.WriteChecksum(0)
 	finally:
-		f.close()
+		return
 
-Export('AddAvailableLibs AddZLib FlagsConfigure GetAvailableLibs GetOptionOrEnv EnablePESecurityFlagsHelper')
+Export('AddAvailableLibs AddZLib FlagsConfigure GetAvailableLibs GetOptionOrEnv IsPEExecutable SetPESecurityFlagsWorker')
diff --git a/SConstruct b/SConstruct
index f8edabd2..4a15a2a5 100644
--- a/SConstruct
+++ b/SConstruct
@@ -420,14 +420,17 @@ def Sign(targets):
 			a = defenv.Action('$CODESIGNER "%s"' % t.path)
 			defenv.AddPostAction(t, a)
 
-Import('EnablePESecurityFlagsHelper')
-def EnablePESecurityFlagsAction(target, source, env):
+Import('IsPEExecutable SetPESecurityFlagsWorker')
+def SetPESecurityFlagsAction(target, source, env):
 	for t in target:
-		EnablePESecurityFlagsHelper(t.path)
-
-def EnablePESecurityFlags(targets):
+		SetPESecurityFlagsWorker(t.path)
+def SetPESecurityFlagsActionEcho(target, source, env):
+	for t in target:
+		if IsPEExecutable(t.path):
+			print('Setting PE flags on %s' % t.name)
+def SetPESecurityFlags(targets):
 	for t in targets:
-		a = defenv.Action(EnablePESecurityFlagsAction, cmdstr=('Setting PE flags on %s' % (t)))
+		a = defenv.Action(SetPESecurityFlagsAction, strfunction=SetPESecurityFlagsActionEcho)
 		defenv.AddPostAction(t, a)
 
 def TestScript(scripts):
@@ -446,7 +449,7 @@ defenv.DistributeDoc = DistributeDoc
 defenv.DistributeDocs = DistributeDocs
 defenv.DistributeExamples = DistributeExamples
 defenv.Sign = Sign
-defenv.EnablePESecurityFlags = EnablePESecurityFlags
+defenv.SetPESecurityFlags = SetPESecurityFlags
 defenv.TestScript = TestScript
 
 def DistributeExtras(env, target, examples, docs):
@@ -686,7 +689,7 @@ def BuildPluginWorker(target, source, libs, examples = None, docs = None,
 	defenv.Alias(target, plugin)
 	defenv.Alias('plugins', plugin)
 
-	defenv.EnablePESecurityFlags(plugin)
+	defenv.SetPESecurityFlags(plugin)
 	defenv.Sign(plugin)
 
 	CleanMap(env, plugin, target)