From e798ba1104e085f4a2e4f99b986bffd928f58a98 Mon Sep 17 00:00:00 2001
From: Matt Nadareski <mnadareski@outlook.com>
Date: Wed, 8 Sep 2021 00:51:25 -0700
Subject: [PATCH] Convert WTM to section based

---
 BurnOutSharp/ProtectionType/WTMCDProtect.cs | 102 ++++++++++++++------
 1 file changed, 73 insertions(+), 29 deletions(-)

diff --git a/BurnOutSharp/ProtectionType/WTMCDProtect.cs b/BurnOutSharp/ProtectionType/WTMCDProtect.cs
index 8e741ef2..7024ee94 100644
--- a/BurnOutSharp/ProtectionType/WTMCDProtect.cs
+++ b/BurnOutSharp/ProtectionType/WTMCDProtect.cs
@@ -1,5 +1,8 @@
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using BurnOutSharp.ExecutableType.Microsoft;
 using BurnOutSharp.Matching;
 
 namespace BurnOutSharp.ProtectionType
@@ -7,37 +10,78 @@ namespace BurnOutSharp.ProtectionType
     public class WTMCDProtect : IContentCheck, IPathCheck
     {
         /// <inheritdoc/>
-        public List<ContentMatchSet> GetContentMatchSets()
-        {
-            return new List<ContentMatchSet>
-            {
-                // This string is found in the .imp files associated with this protection.
-                // WTM76545
-                new ContentMatchSet(new byte?[] { 0x57, 0x54, 0x4D, 0x37, 0x36, 0x35, 0x34, 0x35 }, "WTM CD Protect"),
-
-                // Found in the copy protected setup used by this protection.
-                // wtmdum.imp
-                new ContentMatchSet(new byte?[] { 0x77, 0x74, 0x6D, 0x64, 0x75, 0x6D, 0x2E, 0x69, 0x6D, 0x70 }, "WTM CD Protect"),
-
-                // WTM DIGITAL Photo Protect
-                new ContentMatchSet(new byte?[]
-                {
-                    0x57, 0x54, 0x4D, 0x20, 0x44, 0x49, 0x47, 0x49,
-                    0x54, 0x41, 0x4C, 0x20, 0x50, 0x68, 0x6F, 0x74,
-                    0x6F, 0x20, 0x50, 0x72, 0x6F, 0x74, 0x65, 0x63,
-                    0x74
-                }, "WTM Protection Viewer"),
-
-                // WTM Copy Protection Viewer
-                new ContentMatchSet(new byte?[]
-                {
-                    0x48, 0x61, 0x6E, 0x73, 0x70, 0x65, 0x74, 0x65, 0x72
-                }, "WTM Protection Viewer"),
-            };
-        }
+        public List<ContentMatchSet> GetContentMatchSets() => null;
 
         /// <inheritdoc/>
-        public string CheckContents(string file, byte[] fileContent, bool includeDebug = false) => null;
+        public string CheckContents(string file, byte[] fileContent, bool includeDebug = false)
+        {
+            // Get the sections from the executable, if possible
+            PortableExecutable pex = PortableExecutable.Deserialize(fileContent, 0);
+            var sections = pex?.SectionTable;
+            if (sections == null)
+                return null;
+
+            // Get the CODE section, if it exists
+            var codeSection = sections.FirstOrDefault(s => Encoding.ASCII.GetString(s.Name).StartsWith("CODE"));
+            if (codeSection != null)
+            {
+                int sectionAddr = (int)codeSection.PointerToRawData;
+                int sectionEnd = sectionAddr + (int)codeSection.VirtualSize;
+                var matchers = new List<ContentMatchSet>
+                {
+                    // wtmdum.imp
+                    new ContentMatchSet(
+                        new ContentMatch(new byte?[]
+                        {
+                            0x77, 0x74, 0x6D, 0x64, 0x75, 0x6D, 0x2E, 0x69,
+                            0x6D, 0x70
+                        }, start: sectionAddr, end: sectionEnd),
+                    "WTM CD Protect"),
+                };
+
+                string match = MatchUtil.GetFirstMatch(file, fileContent, matchers, includeDebug);
+                if (!string.IsNullOrWhiteSpace(match))
+                    return match;
+            }
+
+            // Get the .text section, if it exists
+            var textSection = sections.FirstOrDefault(s => Encoding.ASCII.GetString(s.Name).StartsWith(".text"));
+            if (textSection != null)
+            {
+                int sectionAddr = (int)textSection.PointerToRawData;
+                int sectionEnd = sectionAddr + (int)textSection.VirtualSize;
+                var matchers = new List<ContentMatchSet>
+                {
+                    // WTM DIGITAL Photo Protect
+                    new ContentMatchSet(
+                        new ContentMatch(new byte?[]
+                        {
+                            0x57, 0x54, 0x4D, 0x20, 0x44, 0x49, 0x47, 0x49,
+                            0x54, 0x41, 0x4C, 0x20, 0x50, 0x68, 0x6F, 0x74,
+                            0x6F, 0x20, 0x50, 0x72, 0x6F, 0x74, 0x65, 0x63,
+                            0x74
+                        }, start: sectionAddr, end: sectionEnd),
+                    "WTM Protection Viewer"),
+
+                    // WTM Copy Protection Viewer
+                    new ContentMatchSet(
+                        new ContentMatch(new byte?[]
+                        {
+                            0x57, 0x54, 0x4D, 0x20, 0x43, 0x6F, 0x70, 0x79,
+                            0x20, 0x50, 0x72, 0x6F, 0x74, 0x65, 0x63, 0x74,
+                            0x69, 0x6F, 0x6E, 0x20, 0x56, 0x69, 0x65, 0x77,
+                            0x65, 0x72
+                        }, start: sectionAddr, end: sectionEnd),
+                    "WTM Protection Viewer"),
+                };
+
+                string match = MatchUtil.GetFirstMatch(file, fileContent, matchers, includeDebug);
+                if (!string.IsNullOrWhiteSpace(match))
+                    return match;
+            }
+
+            return null;
+        }
 
         /// <inheritdoc/>
         public ConcurrentQueue<string> CheckDirectoryPath(string path, IEnumerable<string> files)